From 4f9da625478653ba9ee08d23d816b05db3c672f0 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 16 Aug 2026 10:03:48 +0800 Subject: [PATCH] [diffusion] chore: use native hunyuan3d paint and delight models (#34980) --- .../configs/models/vaes/stable_diffusion.py | 37 + .../configs/pipeline_configs/hunyuan3d.py | 7 + .../layerwise_offload_components.py | 6 + .../runtime/models/dits/hunyuan3d.py | 843 +-------- .../runtime/models/dits/hunyuan3d_paint.py | 386 +++++ .../runtime/models/dits/stable_diffusion.py | 918 ++++++++++ .../runtime/models/vaes/autoencoder.py | 4 +- .../runtime/pipelines/hunyuan3d_pipeline.py | 263 ++- .../model_specific_stages/hunyuan3d/paint.py | 1530 +++++++---------- .../model_specific_stages/hunyuan3d/shape.py | 5 + .../test_hunyuan3d_native_texture_models.py | 210 +++ 11 files changed, 2442 insertions(+), 1767 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/vaes/stable_diffusion.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d_paint.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py diff --git a/python/sglang/multimodal_gen/configs/models/vaes/stable_diffusion.py b/python/sglang/multimodal_gen/configs/models/vaes/stable_diffusion.py new file mode 100644 index 000000000..520ca80ba --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/stable_diffusion.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Diffusion AutoencoderKL configuration.""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class StableDiffusionVAEArchConfig(VAEArchConfig): + scaling_factor: float = 0.18215 + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 4 + sample_size: int = 32 + block_out_channels: tuple[int, ...] = (64,) + layers_per_block: int = 1 + act_fn: str = "silu" + norm_num_groups: int = 32 + down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",) + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",) + mid_block_add_attention: bool = True + use_quant_conv: bool = True + use_post_quant_conv: bool = True + force_upcast: bool = True + + +@dataclass +class StableDiffusionVAEConfig(VAEConfig): + arch_config: StableDiffusionVAEArchConfig = field( + default_factory=StableDiffusionVAEArchConfig + ) + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + use_parallel_decode: bool = False + use_temporal_scaling_frames: bool = False diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py index 38ac3954b..952cb96a7 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py @@ -4,6 +4,7 @@ from typing import Optional from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig +from sglang.multimodal_gen.configs.models.encoders.clip import CLIPTextConfig from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, @@ -30,6 +31,12 @@ class Hunyuan3D2PipelineConfig(PipelineConfig): vae_config: VAEConfig = field(default_factory=Hunyuan3DVAEConfig) vae_precision: str = "fp32" + text_encoder_configs: tuple[CLIPTextConfig, ...] = field( + default_factory=lambda: (CLIPTextConfig(),) + ) + text_encoder_precisions: tuple[str, ...] = ("fp16",) + native_only_components = ("delight_text_encoder",) + # Shape model configuration shape_model_path: Optional[str] = None shape_use_safetensors: bool = True diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py index 80bbac252..098ce2cca 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py @@ -29,6 +29,8 @@ DIT_COMPONENT_NAMES = frozenset( "video_dit_2", "audio_dit", "dual_tower_bridge", + "delight_transformer", + "paint_transformer", } ) VAE_COMPONENT_NAMES = frozenset( @@ -40,6 +42,8 @@ VAE_COMPONENT_NAMES = frozenset( "spatial_upsampler", "condition_image_encoder", "diffusion_decoder", + "delight_vae", + "paint_vae", } ) DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset( @@ -47,6 +51,8 @@ DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset( "vae", "video_vae", "condition_image_encoder", + "delight_vae", + "paint_vae", } ) CPU_OFFLOAD_FLAG_NAMES = ( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py index 293ad21a8..1feac8f8e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py @@ -3,12 +3,11 @@ from __future__ import annotations import math from dataclasses import dataclass -from typing import List, Optional, Tuple +from typing import Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F -from einops import rearrange from sglang.multimodal_gen.configs.models.dits.hunyuan3d import ( Hunyuan3DDiTArchConfig, @@ -612,842 +611,4 @@ class Hunyuan3D2DiT(CachableDiT, LayerwiseOffloadableModuleMixin): return latent -import copy -import json -import os as _os - -from diffusers.models import UNet2DConditionModel -from diffusers.models.attention_processor import Attention as DiffusersAttention -from diffusers.models.transformers.transformer_2d import BasicTransformerBlock - - -def _chunked_feed_forward( - ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int -): - """Feed forward with chunking to save memory.""" - if hidden_states.shape[chunk_dim] % chunk_size != 0: - raise ValueError( - f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]}" - f"has to be divisible by chunk size: {chunk_size}." - f" Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." - ) - - num_chunks = hidden_states.shape[chunk_dim] // chunk_size - ff_output = torch.cat( - [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)], - dim=chunk_dim, - ) - return ff_output - - -class SGLangAttentionWrapper(torch.nn.Module): - """Drop-in replacement for DiffusersAttention that uses sglang's attention backend.""" - - _SUPPORTED_BACKENDS = {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA} - - def __init__( - self, - query_dim: int, - heads: int = 8, - dim_head: int = 64, - dropout: float = 0.0, - bias: bool = False, - cross_attention_dim: int | None = None, - out_bias: bool = True, - ) -> None: - super().__init__() - self.inner_dim = dim_head * heads - self.heads = heads - self.dim_head = dim_head - self.query_dim = query_dim - cross_attention_dim = cross_attention_dim or query_dim - - self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_k = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias) - self.to_v = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias) - self.to_out = nn.ModuleList( - [nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)] - ) - - from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( - wrap_attention_impl_forward, - ) - from sglang.multimodal_gen.runtime.layers.attention.selector import ( - get_attn_backend, - ) - - attn_backend = get_attn_backend( - dim_head, torch.float16, self._SUPPORTED_BACKENDS - ) - impl_cls = attn_backend.get_impl_cls() - self.attn_impl = impl_cls( - num_heads=heads, - head_size=dim_head, - softmax_scale=dim_head**-0.5, - num_kv_heads=heads, - causal=False, - ) - wrap_attention_impl_forward(self.attn_impl) - self._attn_backend_name = attn_backend.get_enum().name - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - - B, N_q, _ = hidden_states.shape - _, N_kv, _ = encoder_hidden_states.shape - - q = self.to_q(hidden_states).view(B, N_q, self.heads, self.dim_head) - k = self.to_k(encoder_hidden_states).view(B, N_kv, self.heads, self.dim_head) - v = self.to_v(encoder_hidden_states).view(B, N_kv, self.heads, self.dim_head) - - from sglang.multimodal_gen.runtime.managers.forward_context import ( - get_forward_context, - ) - - ctx = get_forward_context() - out = self.attn_impl.forward(q, k, v, attn_metadata=ctx.attn_metadata) - out = out.reshape(B, N_q, self.inner_dim) - - out = self.to_out[0](out) - out = self.to_out[1](out) - return out - - -class Basic2p5DTransformerBlock(torch.nn.Module): - """2.5D Transformer block with Multiview Attention (MVA) and Reference View Attention (RVA).""" - - def __init__( - self, - transformer: BasicTransformerBlock, - layer_name: str, - use_ma: bool = True, - use_ra: bool = True, - is_turbo: bool = False, - use_sglang_attn: bool = True, - ) -> None: - super().__init__() - self.transformer = transformer - self.layer_name = layer_name - self.use_ma = use_ma - self.use_ra = use_ra - self.is_turbo = is_turbo - self.use_sglang_attn = use_sglang_attn and not is_turbo - - attn_cls = ( - SGLangAttentionWrapper if self.use_sglang_attn else DiffusersAttention - ) - attn_kwargs = dict( - query_dim=self.dim, - heads=self.num_attention_heads, - dim_head=self.attention_head_dim, - dropout=self.dropout, - bias=self.attention_bias, - cross_attention_dim=None, - upcast_attention=self.attn1.upcast_attention, - out_bias=True, - ) - if self.use_sglang_attn: - attn_kwargs.pop("upcast_attention") - - if self.use_ma: - self.attn_multiview = attn_cls(**attn_kwargs) - - if self.use_ra: - self.attn_refview = attn_cls(**attn_kwargs) - - if self.is_turbo: - self._initialize_attn_weights() - - def _initialize_attn_weights(self): - """Initialize attention weights for turbo mode.""" - if self.use_ma: - self.attn_multiview.load_state_dict(self.attn1.state_dict()) - with torch.no_grad(): - for layer in self.attn_multiview.to_out: - for param in layer.parameters(): - param.zero_() - if self.use_ra: - self.attn_refview.load_state_dict(self.attn1.state_dict()) - with torch.no_grad(): - for layer in self.attn_refview.to_out: - for param in layer.parameters(): - param.zero_() - - def __getattr__(self, name: str): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self.transformer, name) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, - cross_attention_kwargs: dict = None, - class_labels: Optional[torch.LongTensor] = None, - added_cond_kwargs: Optional[dict] = None, - ) -> torch.Tensor: - """Forward pass with MVA and RVA support.""" - batch_size = hidden_states.shape[0] - - cross_attention_kwargs = ( - cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} - ) - num_in_batch = cross_attention_kwargs.pop("num_in_batch", 1) - mode = cross_attention_kwargs.pop("mode", None) - - if not self.is_turbo: - mva_scale = cross_attention_kwargs.pop("mva_scale", 1.0) - ref_scale = cross_attention_kwargs.pop("ref_scale", 1.0) - else: - position_attn_mask = cross_attention_kwargs.pop("position_attn_mask", None) - position_voxel_indices = cross_attention_kwargs.pop( - "position_voxel_indices", None - ) - mva_scale = 1.0 - ref_scale = 1.0 - - condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None) - - # Normalization - if self.norm_type == "ada_norm": - norm_hidden_states = self.norm1(hidden_states, timestep) - elif self.norm_type == "ada_norm_zero": - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( - hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype - ) - elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]: - norm_hidden_states = self.norm1(hidden_states) - elif self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm1( - hidden_states, added_cond_kwargs["pooled_text_emb"] - ) - elif self.norm_type == "ada_norm_single": - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) - ).chunk(6, dim=1) - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - else: - raise ValueError("Incorrect norm used") - - if self.pos_embed is not None: - norm_hidden_states = self.pos_embed(norm_hidden_states) - - # Prepare GLIGEN inputs - cross_attention_kwargs = ( - cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} - ) - gligen_kwargs = cross_attention_kwargs.pop("gligen", None) - - # Self-attention - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=( - encoder_hidden_states if self.only_cross_attention else None - ), - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - - if self.norm_type == "ada_norm_zero": - attn_output = gate_msa.unsqueeze(1) * attn_output - elif self.norm_type == "ada_norm_single": - attn_output = gate_msa * attn_output - - hidden_states = attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # Reference Attention - Write mode - if mode is not None and "w" in mode: - condition_embed_dict[self.layer_name] = rearrange( - norm_hidden_states, "(b n) l c -> b (n l) c", n=num_in_batch - ) - - # Reference Attention - Read mode - if mode is not None and "r" in mode and self.use_ra: - condition_embed = ( - condition_embed_dict[self.layer_name] - .unsqueeze(1) - .repeat(1, num_in_batch, 1, 1) - ) - condition_embed = rearrange(condition_embed, "b n l c -> (b n) l c") - - attn_output = self.attn_refview( - norm_hidden_states, - encoder_hidden_states=condition_embed, - attention_mask=None, - **cross_attention_kwargs, - ) - - if not self.is_turbo: - ref_scale_timing = ref_scale - if isinstance(ref_scale, torch.Tensor): - ref_scale_timing = ( - ref_scale.unsqueeze(1).repeat(1, num_in_batch).view(-1) - ) - for _ in range(attn_output.ndim - 1): - ref_scale_timing = ref_scale_timing.unsqueeze(-1) - - hidden_states = ref_scale_timing * attn_output + hidden_states - - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # Multiview Attention - if num_in_batch > 1 and self.use_ma: - multivew_hidden_states = rearrange( - norm_hidden_states, "(b n) l c -> b (n l) c", n=num_in_batch - ) - - if self.is_turbo: - position_mask = None - if position_attn_mask is not None: - if multivew_hidden_states.shape[1] in position_attn_mask: - position_mask = position_attn_mask[ - multivew_hidden_states.shape[1] - ] - position_indices = None - if position_voxel_indices is not None: - if multivew_hidden_states.shape[1] in position_voxel_indices: - position_indices = position_voxel_indices[ - multivew_hidden_states.shape[1] - ] - attn_output = self.attn_multiview( - multivew_hidden_states, - encoder_hidden_states=multivew_hidden_states, - attention_mask=position_mask, - position_indices=position_indices, - **cross_attention_kwargs, - ) - else: - attn_output = self.attn_multiview( - multivew_hidden_states, - encoder_hidden_states=multivew_hidden_states, - **cross_attention_kwargs, - ) - - attn_output = rearrange( - attn_output, "b (n l) c -> (b n) l c", n=num_in_batch - ) - - hidden_states = mva_scale * attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # GLIGEN Control - if gligen_kwargs is not None: - hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) - - # Cross-Attention - if self.attn2 is not None: - if self.norm_type == "ada_norm": - norm_hidden_states = self.norm2(hidden_states, timestep) - elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]: - norm_hidden_states = self.norm2(hidden_states) - elif self.norm_type == "ada_norm_single": - norm_hidden_states = hidden_states - elif self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm2( - hidden_states, added_cond_kwargs["pooled_text_emb"] - ) - else: - raise ValueError("Incorrect norm") - - if self.pos_embed is not None and self.norm_type != "ada_norm_single": - norm_hidden_states = self.pos_embed(norm_hidden_states) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - - hidden_states = attn_output + hidden_states - - # Feed-forward - if self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm3( - hidden_states, added_cond_kwargs["pooled_text_emb"] - ) - elif not self.norm_type == "ada_norm_single": - norm_hidden_states = self.norm3(hidden_states) - - if self.norm_type == "ada_norm_zero": - norm_hidden_states = ( - norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - ) - - if self.norm_type == "ada_norm_single": - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp - - if self._chunk_size is not None: - ff_output = _chunked_feed_forward( - self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size - ) - else: - ff_output = self.ff(norm_hidden_states) - - if self.norm_type == "ada_norm_zero": - ff_output = gate_mlp.unsqueeze(1) * ff_output - elif self.norm_type == "ada_norm_single": - ff_output = gate_mlp * ff_output - - hidden_states = ff_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states - - -@torch.no_grad() -def compute_voxel_grid_mask(position: torch.Tensor, grid_resolution: int = 8): - """Compute voxel grid mask for position-aware attention.""" - position = position.half() - B, N, _, H, W = position.shape - assert H % grid_resolution == 0 and W % grid_resolution == 0 - - valid_mask = (position != 1).all(dim=2, keepdim=True) - valid_mask = valid_mask.expand_as(position) - position[valid_mask == False] = 0 - - position = rearrange( - position, - "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", - num_h=grid_resolution, - num_w=grid_resolution, - ) - valid_mask = rearrange( - valid_mask, - "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", - num_h=grid_resolution, - num_w=grid_resolution, - ) - - grid_position = position.sum(dim=(-2, -1)) - count_masked = valid_mask.sum(dim=(-2, -1)) - - grid_position = grid_position / count_masked.clamp(min=1) - grid_position[count_masked < 5] = 0 - - grid_position = grid_position.permute(0, 1, 4, 2, 3) - grid_position = rearrange(grid_position, "b n c h w -> b n (h w) c") - - grid_position_expanded_1 = grid_position.unsqueeze(2).unsqueeze(4) - grid_position_expanded_2 = grid_position.unsqueeze(1).unsqueeze(3) - - distances = torch.norm(grid_position_expanded_1 - grid_position_expanded_2, dim=-1) - - weights = distances - grid_distance = 1.73 / grid_resolution - - weights = weights < grid_distance - - return weights - - -def compute_multi_resolution_mask( - position_maps: torch.Tensor, grid_resolutions: List[int] = [32, 16, 8] -) -> dict: - """Compute multi-resolution position attention masks.""" - position_attn_mask = {} - with torch.no_grad(): - for grid_resolution in grid_resolutions: - position_mask = compute_voxel_grid_mask(position_maps, grid_resolution) - position_mask = rearrange( - position_mask, "b ni nj li lj -> b (ni li) (nj lj)" - ) - position_attn_mask[position_mask.shape[1]] = position_mask - return position_attn_mask - - -@torch.no_grad() -def compute_discrete_voxel_indice( - position: torch.Tensor, grid_resolution: int = 8, voxel_resolution: int = 128 -): - """Compute discrete voxel indices for position encoding.""" - position = position.half() - B, N, _, H, W = position.shape - assert H % grid_resolution == 0 and W % grid_resolution == 0 - - valid_mask = (position != 1).all(dim=2, keepdim=True) - valid_mask = valid_mask.expand_as(position) - position[valid_mask == False] = 0 - - position = rearrange( - position, - "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", - num_h=grid_resolution, - num_w=grid_resolution, - ) - valid_mask = rearrange( - valid_mask, - "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", - num_h=grid_resolution, - num_w=grid_resolution, - ) - - grid_position = position.sum(dim=(-2, -1)) - count_masked = valid_mask.sum(dim=(-2, -1)) - - grid_position = grid_position / count_masked.clamp(min=1) - grid_position[count_masked < 5] = 0 - - grid_position = grid_position.permute(0, 1, 4, 2, 3).clamp(0, 1) - voxel_indices = grid_position * (voxel_resolution - 1) - voxel_indices = torch.round(voxel_indices).long() - return voxel_indices - - -def compute_multi_resolution_discrete_voxel_indice( - position_maps: torch.Tensor, - grid_resolutions: List[int] = [64, 32, 16, 8], - voxel_resolutions: List[int] = [512, 256, 128, 64], -) -> dict: - """Compute multi-resolution discrete voxel indices.""" - voxel_indices = {} - with torch.no_grad(): - for grid_resolution, voxel_resolution in zip( - grid_resolutions, voxel_resolutions - ): - voxel_indice = compute_discrete_voxel_indice( - position_maps, grid_resolution, voxel_resolution - ) - voxel_indice = rearrange(voxel_indice, "b n c h w -> b (n h w) c") - voxel_indices[voxel_indice.shape[1]] = { - "voxel_indices": voxel_indice, - "voxel_resolution": voxel_resolution, - } - return voxel_indices - - -class UNet2p5DConditionModel(torch.nn.Module): - """2.5D UNet for multi-view texture generation.""" - - def __init__(self, unet: UNet2DConditionModel) -> None: - super().__init__() - self.unet = unet - - self.use_ma = True - self.use_ra = True - self.use_camera_embedding = True - self.use_dual_stream = True - self.is_turbo = False - - if self.use_dual_stream: - self.unet_dual = copy.deepcopy(unet) - self.init_attention(self.unet_dual) - self.init_attention( - self.unet, use_ma=self.use_ma, use_ra=self.use_ra, is_turbo=self.is_turbo - ) - self.init_condition() - self.init_camera_embedding() - - @staticmethod - def from_pretrained(pretrained_model_name_or_path: str, **kwargs): - """Load a pretrained UNet2p5DConditionModel.""" - torch_dtype = kwargs.pop("dtype", kwargs.pop("torch_dtype", torch.float32)) - config_path = _os.path.join(pretrained_model_name_or_path, "config.json") - unet_ckpt_path = _os.path.join( - pretrained_model_name_or_path, "diffusion_pytorch_model.bin" - ) - - with open(config_path, "r", encoding="utf-8") as file: - config = json.load(file) - - unet = UNet2DConditionModel(**config) - unet = UNet2p5DConditionModel(unet) - unet_ckpt = torch.load(unet_ckpt_path, map_location="cpu", weights_only=True) - unet.load_state_dict(unet_ckpt, strict=True) - unet = unet.to(torch_dtype) - return unet - - def init_condition(self): - """Initialize condition-related modules.""" - self.unet.conv_in = torch.nn.Conv2d( - 12, # 4 (latent) + 4 (normal) + 4 (position) - self.unet.conv_in.out_channels, - kernel_size=self.unet.conv_in.kernel_size, - stride=self.unet.conv_in.stride, - padding=self.unet.conv_in.padding, - dilation=self.unet.conv_in.dilation, - groups=self.unet.conv_in.groups, - bias=self.unet.conv_in.bias is not None, - ) - - self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1, 77, 1024)) - self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1, 77, 1024)) - - def init_camera_embedding(self): - """Initialize camera embedding module.""" - if self.use_camera_embedding: - time_embed_dim = 1280 - self.max_num_ref_image = 5 - self.max_num_gen_image = 12 * 3 + 4 * 2 - self.unet.class_embedding = nn.Embedding( - self.max_num_ref_image + self.max_num_gen_image, time_embed_dim - ) - - def init_attention( - self, - unet: UNet2DConditionModel, - use_ma: bool = False, - use_ra: bool = False, - is_turbo: bool = False, - use_sglang_attn: bool = True, - ): - """Initialize attention blocks with MVA and RVA support.""" - block_kwargs = dict( - use_ma=use_ma, - use_ra=use_ra, - is_turbo=is_turbo, - use_sglang_attn=use_sglang_attn, - ) - - # Down blocks - for down_block_i, down_block in enumerate(unet.down_blocks): - if ( - hasattr(down_block, "has_cross_attention") - and down_block.has_cross_attention - ): - for attn_i, attn in enumerate(down_block.attentions): - for transformer_i, transformer in enumerate( - attn.transformer_blocks - ): - if isinstance(transformer, BasicTransformerBlock): - attn.transformer_blocks[transformer_i] = ( - Basic2p5DTransformerBlock( - transformer, - f"down_{down_block_i}_{attn_i}_{transformer_i}", - **block_kwargs, - ) - ) - - # Mid block - if ( - hasattr(unet.mid_block, "has_cross_attention") - and unet.mid_block.has_cross_attention - ): - for attn_i, attn in enumerate(unet.mid_block.attentions): - for transformer_i, transformer in enumerate(attn.transformer_blocks): - if isinstance(transformer, BasicTransformerBlock): - attn.transformer_blocks[transformer_i] = ( - Basic2p5DTransformerBlock( - transformer, - f"mid_{attn_i}_{transformer_i}", - **block_kwargs, - ) - ) - - # Up blocks - for up_block_i, up_block in enumerate(unet.up_blocks): - if ( - hasattr(up_block, "has_cross_attention") - and up_block.has_cross_attention - ): - for attn_i, attn in enumerate(up_block.attentions): - for transformer_i, transformer in enumerate( - attn.transformer_blocks - ): - if isinstance(transformer, BasicTransformerBlock): - attn.transformer_blocks[transformer_i] = ( - Basic2p5DTransformerBlock( - transformer, - f"up_{up_block_i}_{attn_i}_{transformer_i}", - **block_kwargs, - ) - ) - - if use_sglang_attn and (use_ma or use_ra): - backend = "unknown" - for block in self._iter_2p5d_blocks(unet): - for attr in ("attn_multiview", "attn_refview"): - wrapper = getattr(block, attr, None) - if isinstance(wrapper, SGLangAttentionWrapper): - backend = wrapper._attn_backend_name - break - if backend != "unknown": - break - count = sum(1 for _ in self._iter_2p5d_blocks(unet)) - logger.info( - "Initialized %d Basic2p5DTransformerBlocks with sglang %s attention", - count, - backend, - ) - - @staticmethod - def _iter_2p5d_blocks(unet): - """Yield all Basic2p5DTransformerBlock instances in a UNet.""" - for block_group in (unet.down_blocks, [unet.mid_block], unet.up_blocks): - for block in block_group: - if not hasattr(block, "attentions"): - continue - for attn in block.attentions: - for tb in attn.transformer_blocks: - if isinstance(tb, Basic2p5DTransformerBlock): - yield tb - - def __getattr__(self, name: str): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self.unet, name) - - def forward( - self, - sample: torch.Tensor, - timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, - *args, - down_intrablock_additional_residuals=None, - down_block_res_samples=None, - mid_block_res_sample=None, - **cached_condition, - ): - """Forward pass for multi-view texture generation.""" - B, N_gen, _, H, W = sample.shape - assert H == W - - if self.use_camera_embedding: - camera_info_gen = ( - cached_condition["camera_info_gen"] + self.max_num_ref_image - ) - camera_info_gen = rearrange(camera_info_gen, "b n -> (b n)") - else: - camera_info_gen = None - - # Concatenate latents with normal and position maps - sample = [sample] - if "normal_imgs" in cached_condition: - sample.append(cached_condition["normal_imgs"]) - if "position_imgs" in cached_condition: - sample.append(cached_condition["position_imgs"]) - sample = torch.cat(sample, dim=2) - - sample = rearrange(sample, "b n c h w -> (b n) c h w") - - encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(1).repeat( - 1, N_gen, 1, 1 - ) - encoder_hidden_states_gen = rearrange( - encoder_hidden_states_gen, "b n l c -> (b n) l c" - ) - - # Process reference images for RVA - if self.use_ra: - if "condition_embed_dict" in cached_condition: - condition_embed_dict = cached_condition["condition_embed_dict"] - else: - condition_embed_dict = {} - ref_latents = cached_condition["ref_latents"] - N_ref = ref_latents.shape[1] - - if self.use_camera_embedding: - camera_info_ref = cached_condition["camera_info_ref"] - camera_info_ref = rearrange(camera_info_ref, "b n -> (b n)") - else: - camera_info_ref = None - - ref_latents = rearrange(ref_latents, "b n c h w -> (b n) c h w") - - encoder_hidden_states_ref = self.unet.learned_text_clip_ref.unsqueeze( - 1 - ).repeat(B, N_ref, 1, 1) - encoder_hidden_states_ref = rearrange( - encoder_hidden_states_ref, "b n l c -> (b n) l c" - ) - - noisy_ref_latents = ref_latents - timestep_ref = 0 - - if self.use_dual_stream: - unet_ref = self.unet_dual - else: - unet_ref = self.unet - - unet_ref( - noisy_ref_latents, - timestep_ref, - encoder_hidden_states=encoder_hidden_states_ref, - class_labels=camera_info_ref, - return_dict=False, - cross_attention_kwargs={ - "mode": "w", - "num_in_batch": N_ref, - "condition_embed_dict": condition_embed_dict, - }, - ) - cached_condition["condition_embed_dict"] = condition_embed_dict - else: - condition_embed_dict = None - - mva_scale = cached_condition.get("mva_scale", 1.0) - ref_scale = cached_condition.get("ref_scale", 1.0) - - if self.is_turbo: - position_attn_mask = cached_condition.get("position_attn_mask", None) - position_voxel_indices = cached_condition.get( - "position_voxel_indices", None - ) - cross_attention_kwargs_ = { - "mode": "r", - "num_in_batch": N_gen, - "condition_embed_dict": condition_embed_dict, - "position_attn_mask": position_attn_mask, - "position_voxel_indices": position_voxel_indices, - "mva_scale": mva_scale, - "ref_scale": ref_scale, - } - else: - cross_attention_kwargs_ = { - "mode": "r", - "num_in_batch": N_gen, - "condition_embed_dict": condition_embed_dict, - "mva_scale": mva_scale, - "ref_scale": ref_scale, - } - - return self.unet( - sample, - timestep, - encoder_hidden_states_gen, - *args, - class_labels=camera_info_gen, - down_intrablock_additional_residuals=( - [ - s.to(dtype=self.unet.dtype) - for s in down_intrablock_additional_residuals - ] - if down_intrablock_additional_residuals is not None - else None - ), - down_block_additional_residuals=( - [s.to(dtype=self.unet.dtype) for s in down_block_res_samples] - if down_block_res_samples is not None - else None - ), - mid_block_additional_residual=( - mid_block_res_sample.to(dtype=self.unet.dtype) - if mid_block_res_sample is not None - else None - ), - return_dict=False, - cross_attention_kwargs=cross_attention_kwargs_, - ) - - -# Entry class for model registry -EntryClass = [Hunyuan3D2DiT, UNet2p5DConditionModel] +EntryClass = Hunyuan3D2DiT diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d_paint.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d_paint.py new file mode 100644 index 000000000..53e968da8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d_paint.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native multi-view UNet used by Hunyuan3D Paint.""" + +from __future__ import annotations + +import copy +from typing import Any + +import torch +from einops import rearrange +from torch import nn + +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.models.dits.stable_diffusion import ( + BasicTransformerBlock, + CrossAttnDownBlock2D, + CrossAttnUpBlock2D, + StableDiffusionAttention, + StableDiffusionUNet2DConditionModel, + StableDiffusionUNetConfig, + StableDiffusionUNetOutput, + Transformer2DModel, + UNetMidBlock2DCrossAttn, +) + + +class Hunyuan3DPaintTransformerBlock(nn.Module): + def __init__( + self, + transformer: BasicTransformerBlock, + layer_name: str, + *, + use_multiview_attention: bool, + use_reference_attention: bool, + is_turbo: bool, + ) -> None: + super().__init__() + self.transformer = transformer + self.layer_name = layer_name + self.use_multiview_attention = use_multiview_attention + self.use_reference_attention = use_reference_attention + self.is_turbo = is_turbo + self.attn_multiview = ( + StableDiffusionAttention( + transformer.dim, + transformer.num_attention_heads, + transformer.attention_head_dim, + ) + if use_multiview_attention + else None + ) + self.attn_refview = ( + StableDiffusionAttention( + transformer.dim, + transformer.num_attention_heads, + transformer.attention_head_dim, + ) + if use_reference_attention + else None + ) + if is_turbo: + self._initialize_added_attention() + + def _initialize_added_attention(self) -> None: + for attention in (self.attn_multiview, self.attn_refview): + if attention is None: + continue + attention.load_state_dict(self.transformer.attn1.state_dict()) + with torch.no_grad(): + for parameter in attention.to_out[0].parameters(): + parameter.zero_() + + @staticmethod + def _broadcast_scale( + scale: float | torch.Tensor, + output: torch.Tensor, + num_views: int, + ) -> float | torch.Tensor: + if not isinstance(scale, torch.Tensor): + return scale + scale = scale.unsqueeze(1).repeat(1, num_views).reshape(-1) + for _ in range(output.ndim - 1): + scale = scale.unsqueeze(-1) + return scale + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor: + options = {} if cross_attention_kwargs is None else cross_attention_kwargs + num_views = int(options.get("num_in_batch", 1)) + mode = options.get("mode") + condition_embeddings = options.get("condition_embed_dict") + if mode is not None and not isinstance(condition_embeddings, dict): + raise ValueError("Hunyuan3D reference attention requires a shared cache.") + + normalized = self.transformer.norm1(hidden_states) + hidden_states = hidden_states + self.transformer.attn1( + normalized, attention_mask=attention_mask + ) + + if mode is not None and "w" in mode: + condition_embeddings[self.layer_name] = rearrange( + normalized, "(b n) l c -> b (n l) c", n=num_views + ) + + if mode is not None and "r" in mode and self.use_reference_attention: + if self.attn_refview is None: + raise RuntimeError("Reference attention was not initialized.") + reference = condition_embeddings[self.layer_name] + reference = reference.unsqueeze(1).repeat(1, num_views, 1, 1) + reference = rearrange(reference, "b n l c -> (b n) l c") + reference_output = self.attn_refview( + normalized, encoder_hidden_states=reference + ) + reference_scale = self._broadcast_scale( + 1.0 if self.is_turbo else options.get("ref_scale", 1.0), + reference_output, + num_views, + ) + hidden_states = hidden_states + reference_scale * reference_output + + if num_views > 1 and self.use_multiview_attention: + if self.attn_multiview is None: + raise RuntimeError("Multiview attention was not initialized.") + multiview = rearrange(normalized, "(b n) l c -> b (n l) c", n=num_views) + position_masks = options.get("position_attn_mask") + position_mask = None + if isinstance(position_masks, dict): + position_mask = position_masks.get(multiview.shape[1]) + multiview_output = self.attn_multiview( + multiview, + encoder_hidden_states=multiview, + attention_mask=position_mask, + ) + multiview_output = rearrange( + multiview_output, "b (n l) c -> (b n) l c", n=num_views + ) + multiview_scale = 1.0 if self.is_turbo else options.get("mva_scale", 1.0) + hidden_states = hidden_states + multiview_scale * multiview_output + + hidden_states = hidden_states + self.transformer.attn2( + self.transformer.norm2(hidden_states), + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + return hidden_states + self.transformer.ff( + self.transformer.norm3(hidden_states) + ) + + +def _replace_transformer_blocks( + unet: StableDiffusionUNet2DConditionModel, + *, + use_multiview_attention: bool, + use_reference_attention: bool, + is_turbo: bool, +) -> None: + def replace(model: Transformer2DModel, layer_name: str) -> None: + transformer = model.transformer_blocks[0] + if not isinstance(transformer, BasicTransformerBlock): + raise TypeError( + f"Expected BasicTransformerBlock, got {type(transformer).__name__}." + ) + model.transformer_blocks[0] = Hunyuan3DPaintTransformerBlock( + transformer, + layer_name, + use_multiview_attention=use_multiview_attention, + use_reference_attention=use_reference_attention, + is_turbo=is_turbo, + ) + + for block_index, block in enumerate(unet.down_blocks): + if not isinstance(block, CrossAttnDownBlock2D): + continue + for attention_index, attention in enumerate(block.attentions): + replace(attention, f"down_{block_index}_{attention_index}_0") + + mid_block = unet.mid_block + if not isinstance(mid_block, UNetMidBlock2DCrossAttn): + raise TypeError(f"Unexpected SD2 mid block: {type(mid_block).__name__}.") + replace(mid_block.attentions[0], "mid_0_0") + + for block_index, block in enumerate(unet.up_blocks): + if not isinstance(block, CrossAttnUpBlock2D): + continue + for attention_index, attention in enumerate(block.attentions): + replace(attention, f"up_{block_index}_{attention_index}_0") + + +@torch.no_grad() +def compute_voxel_grid_mask( + position: torch.Tensor, grid_resolution: int = 8 +) -> torch.Tensor: + position = position.half() + _, _, _, height, width = position.shape + if height % grid_resolution != 0 or width % grid_resolution != 0: + raise ValueError( + f"Position map {height}x{width} is not divisible by {grid_resolution}." + ) + valid_mask = (position != 1).all(dim=2, keepdim=True).expand_as(position) + position = position.masked_fill(~valid_mask, 0) + position = rearrange( + position, + "b n c (nh gh) (nw gw) -> b n nh nw c gh gw", + nh=grid_resolution, + nw=grid_resolution, + ) + valid_mask = rearrange( + valid_mask, + "b n c (nh gh) (nw gw) -> b n nh nw c gh gw", + nh=grid_resolution, + nw=grid_resolution, + ) + counts = valid_mask.sum(dim=(-2, -1)) + grid_position = position.sum(dim=(-2, -1)) / counts.clamp(min=1) + grid_position = grid_position.masked_fill(counts < 5, 0) + grid_position = rearrange(grid_position, "b n h w c -> b n (h w) c") + lhs = grid_position.unsqueeze(2).unsqueeze(4) + rhs = grid_position.unsqueeze(1).unsqueeze(3) + return torch.linalg.vector_norm(lhs - rhs, dim=-1) < 1.73 / grid_resolution + + +def compute_multi_resolution_mask( + position_maps: torch.Tensor, + grid_resolutions: tuple[int, ...] = (32, 16, 8), +) -> dict[int, torch.Tensor]: + masks: dict[int, torch.Tensor] = {} + with torch.no_grad(): + for grid_resolution in grid_resolutions: + mask = compute_voxel_grid_mask(position_maps, grid_resolution) + mask = rearrange(mask, "b ni nj li lj -> b (ni li) (nj lj)") + masks[mask.shape[1]] = mask + return masks + + +class Hunyuan3DPaintUNet(nn.Module, LayerwiseOffloadableModuleMixin): + layer_names = [ + "unet.down_blocks", + "unet.up_blocks", + "unet_dual.down_blocks", + "unet_dual.up_blocks", + ] + layerwise_offload_dit_group_enabled = True + + def __init__( + self, + config: StableDiffusionUNetConfig, + *, + is_turbo: bool = False, + ) -> None: + super().__init__() + base_unet = StableDiffusionUNet2DConditionModel(config) + self.unet = base_unet + self.unet_dual = copy.deepcopy(base_unet) + + _replace_transformer_blocks( + self.unet_dual, + use_multiview_attention=False, + use_reference_attention=False, + is_turbo=is_turbo, + ) + _replace_transformer_blocks( + self.unet, + use_multiview_attention=True, + use_reference_attention=True, + is_turbo=is_turbo, + ) + + self.unet.conv_in = nn.Conv2d(12, self.unet.conv_in.out_channels, 3, padding=1) + self.unet.learned_text_clip_gen = nn.Parameter( + torch.randn(1, 77, config.cross_attention_dim) + ) + self.unet.learned_text_clip_ref = nn.Parameter( + torch.randn(1, 77, config.cross_attention_dim) + ) + self.max_num_ref_images = 5 + self.max_num_generated_images = 44 + time_embedding_dim = config.block_out_channels[0] * 4 + self.unet.class_embedding = nn.Embedding( + self.max_num_ref_images + self.max_num_generated_images, + time_embedding_dim, + ) + + @property + def config(self) -> StableDiffusionUNetConfig: + return self.unet.config + + @property + def dtype(self) -> torch.dtype: + return self.unet.dtype + + @property + def learned_text_clip_gen(self) -> torch.Tensor: + return self.unet.learned_text_clip_gen + + def forward( + self, + sample: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + *, + ref_latents: torch.Tensor, + num_in_batch: int, + condition_embed_dict: dict[str, torch.Tensor], + normal_imgs: torch.Tensor | None = None, + position_imgs: torch.Tensor | None = None, + camera_info_gen: torch.Tensor, + camera_info_ref: torch.Tensor, + ref_scale: float | torch.Tensor = 1.0, + mva_scale: float | torch.Tensor = 1.0, + position_attn_mask: dict[int, torch.Tensor] | None = None, + timestep_cond: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + return_dict: bool = True, + ) -> StableDiffusionUNetOutput | tuple[torch.Tensor]: + if timestep_cond is not None or cross_attention_kwargs is not None: + raise ValueError("Hunyuan3D Paint does not use extra UNet conditioning.") + if added_cond_kwargs is not None: + raise ValueError("Hunyuan3D Paint does not use added conditioning.") + batch_size, num_generated, _, height, width = sample.shape + if height != width or num_generated != num_in_batch: + raise ValueError( + "Hunyuan3D Paint expects square latents and a matching view count." + ) + + camera_gen = rearrange( + camera_info_gen + self.max_num_ref_images, "b n -> (b n)" + ) + inputs = [sample] + if normal_imgs is not None: + inputs.append(normal_imgs) + if position_imgs is not None: + inputs.append(position_imgs) + sample = rearrange(torch.cat(inputs, dim=2), "b n c h w -> (b n) c h w") + encoder_gen = encoder_hidden_states.unsqueeze(1).repeat(1, num_generated, 1, 1) + encoder_gen = rearrange(encoder_gen, "b n l c -> (b n) l c") + + if not condition_embed_dict: + num_reference = ref_latents.shape[1] + camera_ref = rearrange(camera_info_ref, "b n -> (b n)") + reference = rearrange(ref_latents, "b n c h w -> (b n) c h w") + encoder_ref = self.unet.learned_text_clip_ref.unsqueeze(1).repeat( + batch_size, num_reference, 1, 1 + ) + encoder_ref = rearrange(encoder_ref, "b n l c -> (b n) l c") + self.unet_dual( + reference, + 0, + encoder_ref, + class_labels=camera_ref, + return_dict=False, + cross_attention_kwargs={ + "mode": "w", + "num_in_batch": num_reference, + "condition_embed_dict": condition_embed_dict, + }, + ) + + options: dict[str, Any] = { + "mode": "r", + "num_in_batch": num_generated, + "condition_embed_dict": condition_embed_dict, + "mva_scale": mva_scale, + "ref_scale": ref_scale, + } + if position_attn_mask is not None: + options["position_attn_mask"] = position_attn_mask + return self.unet( + sample, + timestep, + encoder_gen, + class_labels=camera_gen, + return_dict=return_dict, + cross_attention_kwargs=options, + ) + + +EntryClass = Hunyuan3DPaintUNet diff --git a/python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py b/python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py new file mode 100644 index 000000000..2b23425a8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native Stable Diffusion 2.1 UNet blocks used by Hunyuan3D texture models.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) + + +@dataclass(frozen=True) +class StableDiffusionUNetConfig: + sample_size: int + in_channels: int + out_channels: int + center_input_sample: bool + flip_sin_to_cos: bool + freq_shift: float + down_block_types: tuple[str, ...] + up_block_types: tuple[str, ...] + block_out_channels: tuple[int, ...] + layers_per_block: int + downsample_padding: int + dropout: float + norm_num_groups: int + norm_eps: float + cross_attention_dim: int + attention_head_dim: tuple[int, ...] + transformer_layers_per_block: int + use_linear_projection: bool + + @classmethod + def from_dict(cls, config: dict[str, Any]) -> StableDiffusionUNetConfig: + block_out_channels = tuple(config["block_out_channels"]) + attention_head_dim_value = config["attention_head_dim"] + attention_head_dim = ( + (attention_head_dim_value,) * len(block_out_channels) + if isinstance(attention_head_dim_value, int) + else tuple(attention_head_dim_value) + ) + parsed = cls( + sample_size=int(config["sample_size"]), + in_channels=int(config["in_channels"]), + out_channels=int(config["out_channels"]), + center_input_sample=bool(config.get("center_input_sample", False)), + flip_sin_to_cos=bool(config.get("flip_sin_to_cos", True)), + freq_shift=float(config.get("freq_shift", 0.0)), + down_block_types=tuple(config["down_block_types"]), + up_block_types=tuple(config["up_block_types"]), + block_out_channels=block_out_channels, + layers_per_block=int(config["layers_per_block"]), + downsample_padding=int(config.get("downsample_padding", 1)), + dropout=float(config.get("dropout", 0.0)), + norm_num_groups=int(config["norm_num_groups"]), + norm_eps=float(config["norm_eps"]), + cross_attention_dim=int(config["cross_attention_dim"]), + attention_head_dim=attention_head_dim, + transformer_layers_per_block=int( + config.get("transformer_layers_per_block", 1) + ), + use_linear_projection=bool(config.get("use_linear_projection", False)), + ) + parsed.validate() + return parsed + + def validate(self) -> None: + expected_down = ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ) + expected_up = ( + "UpBlock2D", + "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", + ) + if self.down_block_types != expected_down or self.up_block_types != expected_up: + raise ValueError( + "The native SD2 UNet currently supports only the Hunyuan3D " + "four-level SD2.1 block layout." + ) + if len(self.block_out_channels) != 4 or len(self.attention_head_dim) != 4: + raise ValueError("Hunyuan3D SD2.1 UNet requires four channel stages.") + if self.layers_per_block != 2 or self.transformer_layers_per_block != 1: + raise ValueError( + "Hunyuan3D SD2.1 UNet requires two ResNet layers and one " + "transformer layer per block." + ) + if not self.use_linear_projection: + raise ValueError("Hunyuan3D SD2.1 checkpoints require linear projection.") + + +@dataclass +class StableDiffusionUNetOutput: + sample: torch.Tensor + + +def timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + *, + flip_sin_to_cos: bool, + downscale_freq_shift: float, +) -> torch.Tensor: + half_dim = embedding_dim // 2 + exponent = -math.log(10000) * torch.arange( + half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + embedding = timesteps[:, None].float() * torch.exp(exponent)[None, :] + embedding = torch.cat([torch.sin(embedding), torch.cos(embedding)], dim=-1) + if flip_sin_to_cos: + embedding = torch.cat( + [embedding[:, half_dim:], embedding[:, :half_dim]], dim=-1 + ) + if embedding_dim % 2 == 1: + embedding = F.pad(embedding, (0, 1)) + return embedding + + +class TimestepEmbedding(nn.Module): + def __init__(self, input_dim: int, embedding_dim: int) -> None: + super().__init__() + self.linear_1 = nn.Linear(input_dim, embedding_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(embedding_dim, embedding_dim) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.linear_2(self.act(self.linear_1(sample))) + + +class StableDiffusionAttention(nn.Module): + def __init__( + self, + query_dim: int, + num_heads: int, + head_dim: int, + cross_attention_dim: int | None = None, + ) -> None: + super().__init__() + inner_dim = num_heads * head_dim + context_dim = cross_attention_dim or query_dim + self.heads = num_heads + self.head_dim = head_dim + self.to_q = nn.Linear(query_dim, inner_dim, bias=False) + self.to_k = nn.Linear(context_dim, inner_dim, bias=False) + self.to_v = nn.Linear(context_dim, inner_dim, bias=False) + self.to_out = nn.ModuleList([nn.Linear(inner_dim, query_dim), nn.Dropout(0.0)]) + + def _prepare_mask(self, attention_mask: torch.Tensor | None) -> torch.Tensor | None: + if attention_mask is None: + return None + if attention_mask.ndim == 2: + return attention_mask[:, None, None, :] + if attention_mask.ndim == 3: + return attention_mask[:, None, :, :] + if attention_mask.ndim == 4: + return attention_mask + raise ValueError( + f"Expected a 2D, 3D, or 4D attention mask, got {attention_mask.ndim}D." + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + context = ( + hidden_states if encoder_hidden_states is None else encoder_hidden_states + ) + batch_size = hidden_states.shape[0] + query = self.to_q(hidden_states).view(batch_size, -1, self.heads, self.head_dim) + key = self.to_k(context).view(batch_size, -1, self.heads, self.head_dim) + value = self.to_v(context).view(batch_size, -1, self.heads, self.head_dim) + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=self._prepare_mask(attention_mask), + dropout_p=0.0, + is_causal=False, + ) + output = output.transpose(1, 2).reshape( + batch_size, -1, self.heads * self.head_dim + ) + return self.to_out[1](self.to_out[0](output)) + + +class GEGLU(nn.Module): + def __init__(self, input_dim: int, output_dim: int) -> None: + super().__init__() + self.proj = nn.Linear(input_dim, output_dim * 2) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) + return hidden_states * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + inner_dim = dim * 4 + self.net = nn.ModuleList( + [GEGLU(dim, inner_dim), nn.Dropout(0.0), nn.Linear(inner_dim, dim)] + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for layer in self.net: + hidden_states = layer(hidden_states) + return hidden_states + + +class BasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + head_dim: int, + cross_attention_dim: int, + ) -> None: + super().__init__() + self.dim = dim + self.num_attention_heads = num_heads + self.attention_head_dim = head_dim + self.norm1 = nn.LayerNorm(dim, eps=1e-5) + self.attn1 = StableDiffusionAttention(dim, num_heads, head_dim) + self.norm2 = nn.LayerNorm(dim, eps=1e-5) + self.attn2 = StableDiffusionAttention( + dim, num_heads, head_dim, cross_attention_dim + ) + self.norm3 = nn.LayerNorm(dim, eps=1e-5) + self.ff = FeedForward(dim) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor: + if cross_attention_kwargs is not None and cross_attention_kwargs: + unsupported = set(cross_attention_kwargs) - {"scale"} + if unsupported: + raise ValueError( + "Unsupported native SD2 cross-attention arguments: " + f"{sorted(unsupported)}" + ) + hidden_states = hidden_states + self.attn1( + self.norm1(hidden_states), attention_mask=attention_mask + ) + hidden_states = hidden_states + self.attn2( + self.norm2(hidden_states), + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + return hidden_states + self.ff(self.norm3(hidden_states)) + + +class Transformer2DModel(nn.Module): + def __init__( + self, + channels: int, + num_heads: int, + cross_attention_dim: int, + norm_num_groups: int, + ) -> None: + super().__init__() + head_dim = channels // num_heads + self.norm = nn.GroupNorm(norm_num_groups, channels, eps=1e-6, affine=True) + self.proj_in = nn.Linear(channels, channels) + self.transformer_blocks = nn.ModuleList( + [BasicTransformerBlock(channels, num_heads, head_dim, cross_attention_dim)] + ) + self.proj_out = nn.Linear(channels, channels) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor: + batch_size, channels, height, width = hidden_states.shape + residual = hidden_states + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch_size, height * width, channels + ) + hidden_states = self.proj_in(hidden_states) + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape( + batch_size, height, width, channels + ).permute(0, 3, 1, 2) + return hidden_states.contiguous() + residual + + +class ResnetBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embedding_dim: int, + norm_num_groups: int, + norm_eps: float, + dropout: float, + ) -> None: + super().__init__() + self.norm1 = nn.GroupNorm( + norm_num_groups, in_channels, eps=norm_eps, affine=True + ) + self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1) + self.time_emb_proj = nn.Linear(time_embedding_dim, out_channels) + self.norm2 = nn.GroupNorm( + norm_num_groups, out_channels, eps=norm_eps, affine=True + ) + self.dropout = nn.Dropout(dropout) + self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1) + self.nonlinearity = nn.SiLU() + self.conv_shortcut = ( + nn.Conv2d(in_channels, out_channels, 1) + if in_channels != out_channels + else None + ) + + def forward( + self, hidden_states: torch.Tensor, time_embedding: torch.Tensor + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.conv1(self.nonlinearity(self.norm1(hidden_states))) + time_states = self.time_emb_proj(self.nonlinearity(time_embedding)) + hidden_states = hidden_states + time_states[:, :, None, None] + hidden_states = self.conv2( + self.dropout(self.nonlinearity(self.norm2(hidden_states))) + ) + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + return residual + hidden_states + + +class Downsample2D(nn.Module): + def __init__(self, channels: int, padding: int) -> None: + super().__init__() + self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=padding) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.conv.padding == (0, 0): + hidden_states = F.pad(hidden_states, (0, 1, 0, 1)) + return self.conv(hidden_states) + + +class Upsample2D(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.conv = nn.Conv2d(channels, channels, 3, padding=1) + + def forward( + self, hidden_states: torch.Tensor, output_size: tuple[int, int] | None = None + ) -> torch.Tensor: + if output_size is None: + hidden_states = F.interpolate( + hidden_states, scale_factor=2.0, mode="nearest" + ) + else: + hidden_states = F.interpolate( + hidden_states, size=output_size, mode="nearest" + ) + return self.conv(hidden_states) + + +class DownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embedding_dim: int, + norm_num_groups: int, + norm_eps: float, + dropout: float, + add_downsample: bool, + downsample_padding: int, + ) -> None: + super().__init__() + self.resnets = nn.ModuleList( + [ + ResnetBlock2D( + in_channels if index == 0 else out_channels, + out_channels, + time_embedding_dim, + norm_num_groups, + norm_eps, + dropout, + ) + for index in range(2) + ] + ) + self.downsamplers = ( + nn.ModuleList([Downsample2D(out_channels, downsample_padding)]) + if add_downsample + else None + ) + + def forward( + self, hidden_states: torch.Tensor, time_embedding: torch.Tensor + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: + output_states: tuple[torch.Tensor, ...] = () + for resnet in self.resnets: + hidden_states = resnet(hidden_states, time_embedding) + output_states += (hidden_states,) + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + output_states += (hidden_states,) + return hidden_states, output_states + + +class CrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embedding_dim: int, + num_heads: int, + cross_attention_dim: int, + norm_num_groups: int, + norm_eps: float, + dropout: float, + add_downsample: bool, + downsample_padding: int, + ) -> None: + super().__init__() + self.resnets = nn.ModuleList( + [ + ResnetBlock2D( + in_channels if index == 0 else out_channels, + out_channels, + time_embedding_dim, + norm_num_groups, + norm_eps, + dropout, + ) + for index in range(2) + ] + ) + self.attentions = nn.ModuleList( + [ + Transformer2DModel( + out_channels, + num_heads, + cross_attention_dim, + norm_num_groups, + ) + for _ in range(2) + ] + ) + self.downsamplers = ( + nn.ModuleList([Downsample2D(out_channels, downsample_padding)]) + if add_downsample + else None + ) + + def forward( + self, + hidden_states: torch.Tensor, + time_embedding: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + encoder_attention_mask: torch.Tensor | None, + cross_attention_kwargs: dict[str, Any] | None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: + output_states: tuple[torch.Tensor, ...] = () + for resnet, attention in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, time_embedding) + hidden_states = attention( + hidden_states, + encoder_hidden_states, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + output_states += (hidden_states,) + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + output_states += (hidden_states,) + return hidden_states, output_states + + +class UNetMidBlock2DCrossAttn(nn.Module): + def __init__( + self, + channels: int, + time_embedding_dim: int, + num_heads: int, + cross_attention_dim: int, + norm_num_groups: int, + norm_eps: float, + dropout: float, + ) -> None: + super().__init__() + self.resnets = nn.ModuleList( + [ + ResnetBlock2D( + channels, + channels, + time_embedding_dim, + norm_num_groups, + norm_eps, + dropout, + ) + for _ in range(2) + ] + ) + self.attentions = nn.ModuleList( + [ + Transformer2DModel( + channels, + num_heads, + cross_attention_dim, + norm_num_groups, + ) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + time_embedding: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + encoder_attention_mask: torch.Tensor | None, + cross_attention_kwargs: dict[str, Any] | None, + ) -> torch.Tensor: + hidden_states = self.resnets[0](hidden_states, time_embedding) + hidden_states = self.attentions[0]( + hidden_states, + encoder_hidden_states, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + return self.resnets[1](hidden_states, time_embedding) + + +class UpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + previous_output_channels: int, + time_embedding_dim: int, + norm_num_groups: int, + norm_eps: float, + dropout: float, + add_upsample: bool, + ) -> None: + super().__init__() + resnets = [] + for index in range(3): + skip_channels = in_channels if index == 2 else out_channels + hidden_channels = previous_output_channels if index == 0 else out_channels + resnets.append( + ResnetBlock2D( + hidden_channels + skip_channels, + out_channels, + time_embedding_dim, + norm_num_groups, + norm_eps, + dropout, + ) + ) + self.resnets = nn.ModuleList(resnets) + self.upsamplers = ( + nn.ModuleList([Upsample2D(out_channels)]) if add_upsample else None + ) + + def forward( + self, + hidden_states: torch.Tensor, + residual_states: tuple[torch.Tensor, ...], + time_embedding: torch.Tensor, + upsample_size: tuple[int, int] | None, + ) -> torch.Tensor: + for resnet in self.resnets: + residual = residual_states[-1] + residual_states = residual_states[:-1] + hidden_states = resnet( + torch.cat([hidden_states, residual], dim=1), time_embedding + ) + if self.upsamplers is not None: + hidden_states = self.upsamplers[0](hidden_states, upsample_size) + return hidden_states + + +class CrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + previous_output_channels: int, + time_embedding_dim: int, + num_heads: int, + cross_attention_dim: int, + norm_num_groups: int, + norm_eps: float, + dropout: float, + add_upsample: bool, + ) -> None: + super().__init__() + resnets = [] + for index in range(3): + skip_channels = in_channels if index == 2 else out_channels + hidden_channels = previous_output_channels if index == 0 else out_channels + resnets.append( + ResnetBlock2D( + hidden_channels + skip_channels, + out_channels, + time_embedding_dim, + norm_num_groups, + norm_eps, + dropout, + ) + ) + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList( + [ + Transformer2DModel( + out_channels, + num_heads, + cross_attention_dim, + norm_num_groups, + ) + for _ in range(3) + ] + ) + self.upsamplers = ( + nn.ModuleList([Upsample2D(out_channels)]) if add_upsample else None + ) + + def forward( + self, + hidden_states: torch.Tensor, + residual_states: tuple[torch.Tensor, ...], + time_embedding: torch.Tensor, + encoder_hidden_states: torch.Tensor, + upsample_size: tuple[int, int] | None, + attention_mask: torch.Tensor | None, + encoder_attention_mask: torch.Tensor | None, + cross_attention_kwargs: dict[str, Any] | None, + ) -> torch.Tensor: + for resnet, attention in zip(self.resnets, self.attentions): + residual = residual_states[-1] + residual_states = residual_states[:-1] + hidden_states = resnet( + torch.cat([hidden_states, residual], dim=1), time_embedding + ) + hidden_states = attention( + hidden_states, + encoder_hidden_states, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + if self.upsamplers is not None: + hidden_states = self.upsamplers[0](hidden_states, upsample_size) + return hidden_states + + +DownBlock = DownBlock2D | CrossAttnDownBlock2D +UpBlock = UpBlock2D | CrossAttnUpBlock2D + + +class StableDiffusionUNet2DConditionModel(nn.Module, LayerwiseOffloadableModuleMixin): + """Native SD2.1 UNet with Diffusers-compatible parameter names.""" + + layer_names = ["down_blocks", "up_blocks"] + layerwise_offload_dit_group_enabled = True + + def __init__(self, config: StableDiffusionUNetConfig) -> None: + super().__init__() + self.config = config + channels = config.block_out_channels + time_embedding_dim = channels[0] * 4 + self.conv_in = nn.Conv2d(config.in_channels, channels[0], 3, padding=1) + self.time_embedding = TimestepEmbedding(channels[0], time_embedding_dim) + self.class_embedding: nn.Embedding | None = None + + down_blocks: list[DownBlock] = [] + output_channels = channels[0] + for index, block_type in enumerate(config.down_block_types): + input_channels = output_channels + output_channels = channels[index] + common = dict( + in_channels=input_channels, + out_channels=output_channels, + time_embedding_dim=time_embedding_dim, + norm_num_groups=config.norm_num_groups, + norm_eps=config.norm_eps, + dropout=config.dropout, + add_downsample=index != len(channels) - 1, + downsample_padding=config.downsample_padding, + ) + if block_type == "CrossAttnDownBlock2D": + down_blocks.append( + CrossAttnDownBlock2D( + **common, + num_heads=config.attention_head_dim[index], + cross_attention_dim=config.cross_attention_dim, + ) + ) + else: + down_blocks.append(DownBlock2D(**common)) + self.down_blocks = nn.ModuleList(down_blocks) + + self.mid_block = UNetMidBlock2DCrossAttn( + channels[-1], + time_embedding_dim, + config.attention_head_dim[-1], + config.cross_attention_dim, + config.norm_num_groups, + config.norm_eps, + config.dropout, + ) + + reversed_channels = tuple(reversed(channels)) + reversed_heads = tuple(reversed(config.attention_head_dim)) + up_blocks: list[UpBlock] = [] + output_channels = reversed_channels[0] + for index, block_type in enumerate(config.up_block_types): + previous_output_channels = output_channels + output_channels = reversed_channels[index] + input_channels = reversed_channels[min(index + 1, len(channels) - 1)] + common = dict( + in_channels=input_channels, + out_channels=output_channels, + previous_output_channels=previous_output_channels, + time_embedding_dim=time_embedding_dim, + norm_num_groups=config.norm_num_groups, + norm_eps=config.norm_eps, + dropout=config.dropout, + add_upsample=index != len(channels) - 1, + ) + if block_type == "CrossAttnUpBlock2D": + up_blocks.append( + CrossAttnUpBlock2D( + **common, + num_heads=reversed_heads[index], + cross_attention_dim=config.cross_attention_dim, + ) + ) + else: + up_blocks.append(UpBlock2D(**common)) + self.up_blocks = nn.ModuleList(up_blocks) + self.conv_norm_out = nn.GroupNorm( + config.norm_num_groups, channels[0], eps=config.norm_eps + ) + self.conv_act = nn.SiLU() + self.conv_out = nn.Conv2d(channels[0], config.out_channels, 3, padding=1) + + @property + def dtype(self) -> torch.dtype: + return self.conv_in.weight.dtype + + def _time_embedding( + self, sample: torch.Tensor, timestep: torch.Tensor | float | int + ) -> torch.Tensor: + if not torch.is_tensor(timestep): + timestep = torch.tensor([timestep], device=sample.device) + elif timestep.ndim == 0: + timestep = timestep[None].to(sample.device) + else: + timestep = timestep.to(sample.device) + timestep = timestep.expand(sample.shape[0]) + projected = timestep_embedding( + timestep, + self.config.block_out_channels[0], + flip_sin_to_cos=self.config.flip_sin_to_cos, + downscale_freq_shift=self.config.freq_shift, + ).to(dtype=sample.dtype) + return self.time_embedding(projected) + + @staticmethod + def _attention_bias( + mask: torch.Tensor | None, dtype: torch.dtype + ) -> torch.Tensor | None: + if mask is None: + return None + if mask.ndim == 2: + return ((1 - mask.to(dtype)) * -10000.0).unsqueeze(1) + return mask + + def forward( + self, + sample: torch.Tensor, + timestep: torch.Tensor | float | int, + encoder_hidden_states: torch.Tensor, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + down_block_additional_residuals: tuple[torch.Tensor, ...] | None = None, + mid_block_additional_residual: torch.Tensor | None = None, + down_intrablock_additional_residuals: tuple[torch.Tensor, ...] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + return_dict: bool = True, + ) -> StableDiffusionUNetOutput | tuple[torch.Tensor]: + if timestep_cond is not None or added_cond_kwargs is not None: + raise ValueError("The Hunyuan3D SD2.1 UNet has no added conditioning.") + if down_intrablock_additional_residuals is not None: + raise ValueError("T2I adapter residuals are not supported by Hunyuan3D.") + if (down_block_additional_residuals is None) != ( + mid_block_additional_residual is None + ): + raise ValueError( + "ControlNet down and mid residuals must be provided together." + ) + + attention_mask = self._attention_bias(attention_mask, sample.dtype) + encoder_attention_mask = self._attention_bias( + encoder_attention_mask, sample.dtype + ) + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + time_embedding = self._time_embedding(sample, timestep) + if self.class_embedding is not None: + if class_labels is None: + raise ValueError("class_labels are required by this UNet.") + time_embedding = time_embedding + self.class_embedding(class_labels).to( + sample.dtype + ) + + forward_upsample_size = any( + dimension % 8 != 0 for dimension in sample.shape[-2:] + ) + sample = self.conv_in(sample) + down_residuals = (sample,) + for block in self.down_blocks: + if isinstance(block, CrossAttnDownBlock2D): + sample, residuals = block( + sample, + time_embedding, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + cross_attention_kwargs, + ) + else: + sample, residuals = block(sample, time_embedding) + down_residuals += residuals + + if down_block_additional_residuals is not None: + down_residuals = tuple( + residual + additional + for residual, additional in zip( + down_residuals, down_block_additional_residuals + ) + ) + + sample = self.mid_block( + sample, + time_embedding, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + cross_attention_kwargs, + ) + if mid_block_additional_residual is not None: + sample = sample + mid_block_additional_residual + + for index, block in enumerate(self.up_blocks): + residuals = down_residuals[-len(block.resnets) :] + down_residuals = down_residuals[: -len(block.resnets)] + upsample_size = ( + down_residuals[-1].shape[-2:] + if index != len(self.up_blocks) - 1 and forward_upsample_size + else None + ) + if isinstance(block, CrossAttnUpBlock2D): + sample = block( + sample, + residuals, + time_embedding, + encoder_hidden_states, + upsample_size, + attention_mask, + encoder_attention_mask, + cross_attention_kwargs, + ) + else: + sample = block(sample, residuals, time_embedding, upsample_size) + + sample = self.conv_out(self.conv_act(self.conv_norm_out(sample))) + if not return_dict: + return (sample,) + return StableDiffusionUNetOutput(sample=sample) + + +EntryClass = StableDiffusionUNet2DConditionModel diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py index 437518a77..5e1b834e0 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py @@ -21,7 +21,7 @@ from diffusers.models.autoencoders.vae import ( from diffusers.models.modeling_outputs import AutoencoderKLOutput from torch import nn -from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig +from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) @@ -76,7 +76,7 @@ class AutoencoderKL(nn.Module, LayerwiseOffloadableModuleMixin): def __init__( self, - config: FluxVAEConfig, + config: VAEConfig, ): super().__init__() self.config = config diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py index c36ef347f..235916361 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py @@ -9,22 +9,45 @@ from __future__ import annotations import glob import importlib +import json import os from itertools import chain from typing import Any import torch import torch.nn as nn +import yaml +from diffusers import EulerAncestralDiscreteScheduler, LCMScheduler +from huggingface_hub import snapshot_download +from safetensors.torch import load_file as load_safetensors +from transformers import AutoTokenizer +from sglang.multimodal_gen.configs.models.vaes.stable_diffusion import ( + StableDiffusionVAEConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import ( + TextEncoderLoader, +) from sglang.multimodal_gen.runtime.loader.fsdp_load import ( load_model_from_full_model_state_dict, set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.models.dits.hunyuan3d_paint import ( + Hunyuan3DPaintUNet, +) +from sglang.multimodal_gen.runtime.models.dits.stable_diffusion import ( + StableDiffusionUNet2DConditionModel, + StableDiffusionUNetConfig, +) +from sglang.multimodal_gen.runtime.models.vaes.autoencoder import ( + AutoencoderKL as StableDiffusionAutoencoderKL, +) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) @@ -39,6 +62,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.h ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) @@ -132,11 +156,9 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): "Local path %s not found, downloading from HuggingFace Hub", local_path, ) - from huggingface_hub import snapshot_download - downloaded = snapshot_download( repo_id=model_path, - allow_patterns=[f"{subfolder}/*"], + allow_patterns=[f"{subfolder}/**"], ) local_path = os.path.join(downloaded, subfolder) @@ -167,8 +189,11 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): return config_path, ckpt_path @staticmethod - def _resolve_paint_dir(model_path: str, subfolder: str) -> str: - """Locate (or download) the paint subfolder and return its local path.""" + def _resolve_model_subfolder( + model_path: str, + subfolder: str, + required_files: tuple[str, ...], + ) -> str: local_path = os.path.join(model_path, subfolder) if not os.path.exists(local_path): local_path = os.path.expanduser(local_path) @@ -178,23 +203,21 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): "Local path %s not found, downloading from HuggingFace Hub", local_path, ) - from huggingface_hub import snapshot_download - downloaded = snapshot_download( repo_id=model_path, - allow_patterns=[f"{subfolder}/*"], + allow_patterns=[f"{subfolder}/**"], ) local_path = os.path.join(downloaded, subfolder) - for subdir in ("vae", "unet"): - config_file = os.path.join(local_path, subdir, "config.json") - if not os.path.exists(config_file): + for relative_path in required_files: + required_file = os.path.join(local_path, relative_path) + if not os.path.exists(required_file): raise FileNotFoundError( - f"Paint model incomplete: {config_file} not found. " + f"Hunyuan3D model incomplete: {required_file} not found. " "Download the model or check network connectivity." ) - logger.info("Resolved paint model directory: %s", local_path) + logger.debug("Resolved Hunyuan3D model directory: %s", local_path) return local_path @staticmethod @@ -203,9 +226,7 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): ) -> dict[str, dict[str, torch.Tensor]]: """Load a bundled checkpoint and split by the first '.' in each key.""" if use_safetensors: - import safetensors.torch - - flat = safetensors.torch.load_file(ckpt_path, device="cpu") + flat = load_safetensors(ckpt_path, device="cpu") ckpt: dict[str, dict[str, torch.Tensor]] = {} for key, value in flat.items(): component = key.split(".")[0] @@ -291,16 +312,193 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): params = cfg.get("params", {}) return target_cls(**params) + @staticmethod + def _read_json(path: str) -> dict[str, Any]: + with open(path, encoding="utf-8") as file: + return json.load(file) + + @staticmethod + def _load_component_weights(component_dir: str) -> dict[str, torch.Tensor]: + safetensors_path = os.path.join( + component_dir, "diffusion_pytorch_model.safetensors" + ) + pytorch_path = os.path.join(component_dir, "diffusion_pytorch_model.bin") + if os.path.isfile(safetensors_path): + return load_safetensors(safetensors_path, device="cpu") + if os.path.isfile(pytorch_path): + return torch.load(pytorch_path, map_location="cpu", weights_only=True) + raise FileNotFoundError( + f"No diffusion_pytorch_model weights found in {component_dir}." + ) + + @staticmethod + def _component_device(server_args: ServerArgs, component_name: str) -> torch.device: + if server_args.should_cpu_offload_component(component_name): + return torch.device("cpu") + return get_local_torch_device() + + @staticmethod + def _freeze(module: nn.Module) -> nn.Module: + for parameter in module.parameters(): + parameter.requires_grad = False + return module.eval() + + @staticmethod + def _maybe_compile_texture_transformer( + module: nn.Module, + server_args: ServerArgs, + config: Hunyuan3D2PipelineConfig, + ) -> None: + if not server_args.enable_torch_compile: + return + compile_mode = ( + os.environ.get("SGLANG_TORCH_COMPILE_MODE") + or config.dit_config.torch_compile_mode + ) + logger.info( + "Compiling %s with mode: %s", + module.__class__.__name__, + compile_mode, + ) + module.compile(mode=compile_mode, fullgraph=False, dynamic=None) + + @classmethod + def _load_stable_diffusion_vae( + cls, + component_dir: str, + dtype: torch.dtype, + device: torch.device, + ) -> StableDiffusionAutoencoderKL: + config_data = cls._read_json(os.path.join(component_dir, "config.json")) + vae_config = StableDiffusionVAEConfig() + vae_config.update_model_arch(config_data) + with set_default_torch_dtype(dtype): + vae = StableDiffusionAutoencoderKL(vae_config) + weights = cls._load_component_weights(component_dir) + vae.load_state_dict(weights, strict=True) + vae.to(device=device, dtype=dtype) + return cls._freeze(vae) + + @classmethod + def _load_stable_diffusion_unet( + cls, + component_dir: str, + dtype: torch.dtype, + device: torch.device, + *, + paint: bool, + is_turbo: bool = False, + ) -> nn.Module: + config_data = cls._read_json(os.path.join(component_dir, "config.json")) + unet_config = StableDiffusionUNetConfig.from_dict(config_data) + with set_default_torch_dtype(dtype), torch.device("meta"): + unet: nn.Module + if paint: + unet = Hunyuan3DPaintUNet(unet_config, is_turbo=is_turbo) + else: + unet = StableDiffusionUNet2DConditionModel(unet_config) + weights = cls._load_component_weights(component_dir) + unet.load_state_dict(weights, strict=True, assign=True) + unet.to(device=device, dtype=dtype) + return cls._freeze(unet) + + @classmethod + def _load_texture_components( + cls, + server_args: ServerArgs, + config: Hunyuan3D2PipelineConfig, + ) -> dict[str, Any]: + dtype = PRECISION_TO_TYPE[config.dit_precision] + components: dict[str, Any] = {} + + paint_dir = cls._resolve_model_subfolder( + server_args.model_path, + config.paint_subfolder, + ( + "vae/config.json", + "unet/config.json", + "scheduler/scheduler_config.json", + ), + ) + components["paint_vae"] = cls._load_stable_diffusion_vae( + os.path.join(paint_dir, "vae"), + dtype, + cls._component_device(server_args, "paint_vae"), + ) + components["paint_transformer"] = cls._load_stable_diffusion_unet( + os.path.join(paint_dir, "unet"), + dtype, + cls._component_device(server_args, "paint_transformer"), + paint=True, + is_turbo=config.paint_turbo_mode, + ) + cls._maybe_compile_texture_transformer( + components["paint_transformer"], server_args, config + ) + paint_scheduler_config = cls._read_json( + os.path.join(paint_dir, "scheduler", "scheduler_config.json") + ) + scheduler_class = ( + LCMScheduler if config.paint_turbo_mode else EulerAncestralDiscreteScheduler + ) + components["paint_scheduler"] = scheduler_class.from_config( + paint_scheduler_config, + **({} if config.paint_turbo_mode else {"timestep_spacing": "trailing"}), + ) + + if config.delight_enable: + delight_dir = cls._resolve_model_subfolder( + server_args.model_path, + config.delight_subfolder, + ( + "vae/config.json", + "unet/config.json", + "scheduler/scheduler_config.json", + "text_encoder/config.json", + "tokenizer/tokenizer_config.json", + ), + ) + components["delight_vae"] = cls._load_stable_diffusion_vae( + os.path.join(delight_dir, "vae"), + dtype, + cls._component_device(server_args, "delight_vae"), + ) + components["delight_transformer"] = cls._load_stable_diffusion_unet( + os.path.join(delight_dir, "unet"), + dtype, + cls._component_device(server_args, "delight_transformer"), + paint=False, + ) + cls._maybe_compile_texture_transformer( + components["delight_transformer"], server_args, config + ) + delight_text_encoder, _ = TextEncoderLoader().load( + os.path.join(delight_dir, "text_encoder"), + server_args, + "delight_text_encoder", + "transformers", + ) + components["delight_text_encoder"] = delight_text_encoder + components["delight_tokenizer"] = AutoTokenizer.from_pretrained( + os.path.join(delight_dir, "tokenizer") + ) + delight_scheduler_config = cls._read_json( + os.path.join(delight_dir, "scheduler", "scheduler_config.json") + ) + components["delight_scheduler"] = ( + EulerAncestralDiscreteScheduler.from_config(delight_scheduler_config) + ) + + return components + # Module loading override def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: - """Load all Hunyuan3D shape components from a bundled checkpoint.""" - import yaml - - from sglang.multimodal_gen.runtime.distributed import get_local_torch_device + """Load Hunyuan3D shape and optional texture components.""" + del loaded_modules config = server_args.pipeline_config if not isinstance(config, Hunyuan3D2PipelineConfig): @@ -348,16 +546,10 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): model_config["image_processor"] ) - logger.info("All Hunyuan3D shape components loaded successfully") - if config.paint_enable: - try: - paint_dir = self._resolve_paint_dir( - server_args.model_path, config.paint_subfolder - ) - components["hy3dpaint_dir"] = paint_dir - except Exception as e: - logger.warning("Failed to resolve paint model path: %s", e) + components.update(self._load_texture_components(server_args, config)) + + logger.info("Loaded Hunyuan3D pipeline components: %s", sorted(components)) return components @@ -411,13 +603,22 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase): if config.paint_enable: self.add_stage( stage_name="paint_preprocess", - stage=Hunyuan3DPaintPreprocessStage(config=config), + stage=Hunyuan3DPaintPreprocessStage( + config=config, + delight_transformer=self.get_module("delight_transformer"), + delight_vae=self.get_module("delight_vae"), + delight_text_encoder=self.get_module("delight_text_encoder"), + delight_tokenizer=self.get_module("delight_tokenizer"), + delight_scheduler=self.get_module("delight_scheduler"), + ), ) self.add_stage( stage_name="paint_texgen", stage=Hunyuan3DPaintTexGenStage( config=config, - paint_dir=self.get_module("hy3dpaint_dir"), + transformer=self.get_module("paint_transformer"), + scheduler=self.get_module("paint_scheduler"), + vae=self.get_module("paint_vae"), ), ) self.add_stage( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py index 15101a371..ca680fc9a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py @@ -1,25 +1,41 @@ -""" -Hunyuan3D paint/texture generation stages. - -Three-stage pipeline: Preprocess -> TexGen -> Postprocess. -""" +# SPDX-License-Identifier: Apache-2.0 +"""Hunyuan3D texture-generation stages.""" from __future__ import annotations +import concurrent.futures +import inspect import os +from dataclasses import dataclass from typing import Any +import cv2 import numpy as np import torch -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import ( - retrieve_timesteps, -) +from diffusers.image_processor import VaeImageProcessor +from diffusers.utils.torch_utils import randn_tensor from einops import rearrange +from PIL import Image +from torch import nn +from transformers import PreTrainedTokenizerBase +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.models.dits.hunyuan3d_paint import ( + Hunyuan3DPaintUNet, + compute_multi_resolution_mask, +) +from sglang.multimodal_gen.runtime.models.dits.stable_diffusion import ( + StableDiffusionUNet2DConditionModel, +) +from sglang.multimodal_gen.runtime.models.encoders.clip import CLIPTextModel +from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, @@ -32,458 +48,392 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.utils import PRECISION_TO_TYPE - -logger = init_logger(__name__) -# Utility functions -def guidance_scale_embedding( - w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 -) -> torch.Tensor: - """Generate guidance scale embeddings.""" - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb +def _module_dtype(module: nn.Module) -> torch.dtype: + return next(module.parameters()).dtype -def extract_into_tensor( - a: torch.Tensor, t: torch.Tensor, x_shape: tuple, n_gen: int -) -> torch.Tensor: - """Extract values from tensor and reshape for multi-view generation.""" - out = a.gather(-1, t) - out = out.repeat(n_gen) - out = rearrange(out, "(b n) -> b n", n=n_gen) - b, c, *_ = out.shape - return out.reshape(b, c, *((1,) * (len(x_shape) - 2))) - - -def get_predicted_original_sample( - model_output: torch.Tensor, - timesteps: torch.Tensor, - sample: torch.Tensor, - prediction_type: str, - alphas: torch.Tensor, - sigmas: torch.Tensor, - n_gen: int, -) -> torch.Tensor: - """Get predicted original sample from model output.""" - alphas = extract_into_tensor(alphas, timesteps, sample.shape, n_gen) - sigmas = extract_into_tensor(sigmas, timesteps, sample.shape, n_gen) - model_output = rearrange(model_output, "(b n) c h w -> b n c h w", n=n_gen) - - if prediction_type == "epsilon": - pred_x_0 = (sample - sigmas * model_output) / alphas - elif prediction_type == "sample": - pred_x_0 = model_output - elif prediction_type == "v_prediction": - pred_x_0 = alphas * sample - sigmas * model_output - else: - raise ValueError( - f"Prediction type {prediction_type} is not supported; " - "currently, `epsilon`, `sample`, and `v_prediction` are supported." - ) - - return pred_x_0 - - -def get_predicted_noise( - model_output: torch.Tensor, - timesteps: torch.Tensor, - sample: torch.Tensor, - prediction_type: str, - alphas: torch.Tensor, - sigmas: torch.Tensor, - n_gen: int, -) -> torch.Tensor: - """Get predicted noise from model output.""" - alphas = extract_into_tensor(alphas, timesteps, sample.shape, n_gen) - sigmas = extract_into_tensor(sigmas, timesteps, sample.shape, n_gen) - model_output = rearrange(model_output, "(b n) c h w -> b n c h w", n=n_gen) - - if prediction_type == "epsilon": - pred_epsilon = model_output - elif prediction_type == "sample": - pred_epsilon = (sample - alphas * model_output) / sigmas - elif prediction_type == "v_prediction": - pred_epsilon = alphas * model_output + sigmas * sample - else: - raise ValueError( - f"Prediction type {prediction_type} is not supported; " - "currently, `epsilon`, `sample`, and `v_prediction` are supported." - ) - - return pred_epsilon - - -def to_rgb_image(maybe_rgba): - """Convert RGBA image to RGB.""" - from PIL import Image - - if maybe_rgba.mode == "RGB": - return maybe_rgba - if maybe_rgba.mode == "RGBA": - rgba = maybe_rgba - img = np.random.randint( - 127, 128, size=[rgba.size[1], rgba.size[0], 3], dtype=np.uint8 - ) - img = Image.fromarray(img, "RGB") - img.paste(rgba, mask=rgba.getchannel("A")) - return img - raise ValueError(f"Unsupported image type: {maybe_rgba.mode}") - - -class DDIMSolver: - """DDIM solver for fast sampling.""" - - def __init__( - self, - alpha_cumprods: np.ndarray, - timesteps: int = 1000, - ddim_timesteps: int = 50, - ): - step_ratio = timesteps // ddim_timesteps - self.ddim_timesteps = ( - np.arange(1, ddim_timesteps + 1) * step_ratio - ).round().astype(np.int64) - 1 - self.ddim_alpha_cumprods = alpha_cumprods[self.ddim_timesteps] - self.ddim_alpha_cumprods_prev = np.asarray( - [alpha_cumprods[0]] + alpha_cumprods[self.ddim_timesteps[:-1]].tolist() - ) - self.ddim_timesteps = torch.from_numpy(self.ddim_timesteps).long() - self.ddim_alpha_cumprods = torch.from_numpy(self.ddim_alpha_cumprods) - self.ddim_alpha_cumprods_prev = torch.from_numpy(self.ddim_alpha_cumprods_prev) - - def to(self, device: torch.device) -> DDIMSolver: - self.ddim_timesteps = self.ddim_timesteps.to(device) - self.ddim_alpha_cumprods = self.ddim_alpha_cumprods.to(device) - self.ddim_alpha_cumprods_prev = self.ddim_alpha_cumprods_prev.to(device) - return self - - def ddim_step( - self, - pred_x0: torch.Tensor, - pred_noise: torch.Tensor, - timestep_index: torch.Tensor, - n_gen: int, - ) -> torch.Tensor: - alpha_cumprod_prev = extract_into_tensor( - self.ddim_alpha_cumprods_prev, timestep_index, pred_x0.shape, n_gen - ) - dir_xt = (1.0 - alpha_cumprod_prev).sqrt() * pred_noise - x_prev = alpha_cumprod_prev.sqrt() * pred_x0 + dir_xt - return x_prev +def _to_rgb_image(image: Image.Image, background: int = 127) -> Image.Image: + if image.mode == "RGB": + return image + if image.mode != "RGBA": + raise ValueError(f"Unsupported image mode: {image.mode}") + background_image = Image.new("RGB", image.size, (background,) * 3) + background_image.paste(image, mask=image.getchannel("A")) + return background_image def _recorrect_rgb( - src_image: torch.Tensor, - target_image: torch.Tensor, - alpha_channel: torch.Tensor, + source: torch.Tensor, + target: torch.Tensor, + alpha: torch.Tensor, scale: float = 0.95, ) -> torch.Tensor: - """Correct RGB values to match target color distribution.""" + mask = alpha[..., 0] > 0.5 + if not torch.any(mask): + return torch.cat([target, alpha], dim=-1) - def flat_and_mask(bgr, a): - mask = torch.where(a > 0.5, True, False) - bgr_flat = bgr.reshape(-1, bgr.shape[-1]) - mask_flat = mask.reshape(-1) - bgr_flat_masked = bgr_flat[mask_flat, :] - return bgr_flat_masked + source_masked = source[mask] + target_masked = target[mask] + corrected = torch.empty_like(source) + epsilon = torch.finfo(source.dtype).eps + for channel in range(3): + source_mean = source_masked[:, channel].mean() + source_std = source_masked[:, channel].std().clamp_min(epsilon) + target_mean = target_masked[:, channel].mean() + target_std = target_masked[:, channel].std() + corrected[..., channel] = ( + (source[..., channel] - scale * source_mean) * (target_std / source_std) + + scale * target_mean + ).clamp(0, 1) - src_flat = flat_and_mask(src_image, alpha_channel) - target_flat = flat_and_mask(target_image, alpha_channel) - corrected_bgr = torch.zeros_like(src_image) + if torch.mean((source - target) ** 2) < torch.mean((corrected - target) ** 2): + corrected = source + return torch.cat([corrected, alpha], dim=-1) - for i in range(3): - src_mean, src_stddev = torch.mean(src_flat[:, i]), torch.std(src_flat[:, i]) - target_mean, target_stddev = torch.mean(target_flat[:, i]), torch.std( - target_flat[:, i] - ) - corrected_bgr[:, :, i] = torch.clamp( - (src_image[:, :, i] - scale * src_mean) * (target_stddev / src_stddev) - + scale * target_mean, - 0, - 1, - ) - src_mse = torch.mean((src_image - target_image) ** 2) - modify_mse = torch.mean((corrected_bgr - target_image) ** 2) - if src_mse < modify_mse: - corrected_bgr = torch.cat([src_image, alpha_channel], dim=-1) - else: - corrected_bgr = torch.cat([corrected_bgr, alpha_channel], dim=-1) +def _scheduler_step_kwargs( + scheduler: Any, generator: torch.Generator +) -> dict[str, Any]: + parameters = inspect.signature(scheduler.step).parameters + kwargs: dict[str, Any] = {} + if "eta" in parameters: + kwargs["eta"] = 0.0 + if "generator" in parameters: + kwargs["generator"] = generator + return kwargs - return corrected_bgr + +@dataclass(slots=True) +class PaintDenoisingInputs: + timesteps: torch.Tensor + latents: torch.Tensor + model_kwargs: dict[str, Any] + num_views: int + guidance_scale: float + use_cfg: bool + generator: torch.Generator + latent_channels: int -# Stage 1: Preprocess (UV unwrap + delight + multi-view rendering) class Hunyuan3DPaintPreprocessStage(PipelineStage): - """Preprocessing: UV unwrap + delight in parallel, then multi-view rendering.""" + """Unwrap the mesh, remove image lighting, and render geometry controls.""" CAMERA_AZIMS = [0, 90, 180, 270, 0, 180] CAMERA_ELEVS = [0, 0, 0, 0, 90, -90] VIEW_WEIGHTS = [1, 0.1, 0.5, 0.1, 0.05, 0.05] + def __init__( + self, + config: Hunyuan3D2PipelineConfig, + delight_transformer: StableDiffusionUNet2DConditionModel | None, + delight_vae: AutoencoderKL | None, + delight_text_encoder: CLIPTextModel | None, + delight_tokenizer: PreTrainedTokenizerBase | None, + delight_scheduler: Any, + ) -> None: + super().__init__() + self.config = config + self.delight_transformer = delight_transformer + self.delight_vae = delight_vae + self.delight_text_encoder = delight_text_encoder + self.delight_tokenizer = delight_tokenizer + self.delight_scheduler = delight_scheduler + self._renderer: Any = None + self._delight_image_processor = VaeImageProcessor(vae_scale_factor=8) + + if config.delight_enable and any( + component is None + for component in ( + delight_transformer, + delight_vae, + delight_text_encoder, + delight_tokenizer, + delight_scheduler, + ) + ): + raise ValueError( + "Delight is enabled, but its model components are missing." + ) + @property def parallelism_type(self) -> StageParallelismType: return StageParallelismType.MAIN_RANK_ONLY - def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: - super().__init__() - self.config = config - self._delight_pipeline = None - self._delight_loaded = False - self._renderer = None - self._renderer_loaded = False - - # --- UV unwrap --- - - def _do_uv_unwrap(self, batch: Req, server_args: ServerArgs) -> Req: - import time + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + del server_args + if not self.config.delight_enable: + return [] + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse(stage_name, "delight_text_encoder", phase="prompt"), + ComponentUse(stage_name, "delight_vae", phase="encode"), + ComponentUse( + stage_name, + "delight_transformer", + phase="denoise", + memory_intensive=True, + ), + ComponentUse(stage_name, "delight_vae", phase="decode"), + ] + @staticmethod + def _unwrap_mesh(mesh: Any) -> Any: from sglang.multimodal_gen.runtime.utils.mesh3d_utils import mesh_uv_wrap - mesh = batch.extra["shape_meshes"] - if isinstance(mesh, list): - mesh = mesh[0] - - try: - start_time = time.time() - mesh = mesh_uv_wrap(mesh) - elapsed = time.time() - start_time - logger.info(f"UV unwrapping completed in {elapsed:.2f}s") - except Exception as e: - logger.warning(f"UV unwrapping failed: {e}") - - batch.extra["paint_mesh"] = mesh - return batch - - # --- Delight --- - - def _load_delight_model(self, server_args: ServerArgs): - if self._delight_loaded: - return - - from diffusers import ( - EulerAncestralDiscreteScheduler, - StableDiffusionInstructPix2PixPipeline, - ) - from huggingface_hub import snapshot_download - - model_path = server_args.model_path - delight_subfolder = getattr( - self.config, "delight_subfolder", "hunyuan3d-delight-v2-0" - ) - - local_path = os.path.join(model_path, delight_subfolder) - if not os.path.exists(local_path): - local_path = os.path.expanduser(local_path) - - if not os.path.exists(local_path): - try: - downloaded = snapshot_download( - repo_id=model_path, - allow_patterns=[f"{delight_subfolder}/*"], - ) - local_path = os.path.join(downloaded, delight_subfolder) - except Exception as e: - logger.warning("Could not download delight model: %s", e) - local_path = None - - if local_path and os.path.exists(local_path): - # Resolve precision from config with a simple fallback for CPU/MPS - dit_dtype = PRECISION_TO_TYPE.get( - getattr(self.config, "dit_precision", "fp16"), torch.float16 - ) - if self.device.type in ("cpu", "mps") and dit_dtype in ( - torch.float16, - torch.bfloat16, - ): - # Avoid half/bfloat on CPU/MPS to be safe - dit_dtype = torch.float32 - pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained( - local_path, - torch_dtype=dit_dtype, - safety_checker=None, - ) - pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( - pipeline.scheduler.config - ) - pipeline.set_progress_bar_config(disable=True) - self._delight_pipeline = pipeline.to(self.device, dit_dtype) - logger.info("Delight model loaded successfully") - else: - logger.warning( - "Delight model not available, skipping delight preprocessing" - ) - - self._delight_loaded = True - - @torch.no_grad() - def _run_delight(self, image): - import cv2 - from PIL import Image as PILImage - - image = image.resize((512, 512)) - - if image.mode == "RGBA": - image_array = np.array(image) - alpha_channel = image_array[:, :, 3] - erosion_size = 3 - kernel = np.ones((erosion_size, erosion_size), np.uint8) - alpha_channel = cv2.erode(alpha_channel, kernel, iterations=1) - image_array[alpha_channel == 0, :3] = 255 - image_array[:, :, 3] = alpha_channel - image = PILImage.fromarray(image_array) - - image_tensor = torch.tensor(np.array(image) / 255.0).float().to(self.device) - alpha = image_tensor[:, :, 3:] - rgb_target = image_tensor[:, :, :3] - else: - image_tensor = torch.tensor(np.array(image) / 255.0).float().to(self.device) - alpha = torch.ones_like(image_tensor)[:, :, :1] - rgb_target = image_tensor[:, :, :3] - - image = image.convert("RGB") - - image = self._delight_pipeline( - prompt=self.config.delight_prompt, - negative_prompt=getattr(self.config, "delight_negative_prompt", ""), - image=image, - generator=torch.manual_seed(42), - height=512, - width=512, - num_inference_steps=self.config.delight_num_inference_steps, - image_guidance_scale=self.config.delight_cfg_image, - guidance_scale=self.config.delight_guidance_scale, - ).images[0] - - image_tensor = torch.tensor(np.array(image) / 255.0).float().to(self.device) - rgb_src = image_tensor[:, :, :3] - image = _recorrect_rgb(rgb_src, rgb_target, alpha) - image = image[:, :, :3] * image[:, :, 3:] + torch.ones_like(image[:, :, :3]) * ( - 1.0 - image[:, :, 3:] - ) - image = PILImage.fromarray((image.cpu().numpy() * 255).astype(np.uint8)) - - return image - - def _do_delight(self, batch: Req, server_args: ServerArgs) -> Req: - from PIL import Image + return mesh_uv_wrap(mesh) + @staticmethod + def _load_input_image(image_path: str) -> Image.Image: from sglang.multimodal_gen.runtime.utils.mesh3d_utils import recenter_image - image = Image.open(batch.image_path) - image = recenter_image(image) + with Image.open(image_path) as input_image: + return recenter_image(input_image.copy()) - if not self.config.delight_enable: - logger.info("Delight preprocessing disabled, using original image") - batch.extra["delighted_image"] = image - return batch + @staticmethod + def _prepare_delight_target( + image: Image.Image, device: torch.device + ) -> tuple[Image.Image, torch.Tensor, torch.Tensor]: + image = image.resize((512, 512), Image.Resampling.BICUBIC) + if image.mode == "RGBA": + pixels = np.asarray(image).copy() + kernel = np.ones((3, 3), np.uint8) + pixels[..., 3] = cv2.erode(pixels[..., 3], kernel, iterations=1) + pixels[pixels[..., 3] == 0, :3] = 255 + image = Image.fromarray(pixels, mode="RGBA") + target = torch.from_numpy(pixels.astype(np.float32) / 255).to(device) + return image.convert("RGB"), target[..., :3], target[..., 3:] - self._load_delight_model(server_args) - if self._delight_pipeline is not None: - try: - image = self._run_delight(image) - logger.info("Image delight completed") - except Exception as e: - logger.warning(f"Image delight failed: {e}") + image = image.convert("RGB") + target = torch.from_numpy(np.asarray(image, dtype=np.float32) / 255).to(device) + return image, target, torch.ones_like(target[..., :1]) - batch.extra["delighted_image"] = image - return batch - - # --- Multi-view rendering --- - - def _init_renderer(self): - if self._renderer_loaded: - return - - from sglang.multimodal_gen.runtime.utils.mesh3d_utils import MeshRender - - self._renderer = MeshRender( - default_resolution=self.config.paint_render_size, - texture_size=self.config.paint_texture_size, - device=self.device, + def _encode_delight_prompts( + self, use_cfg: bool, device: torch.device + ) -> torch.Tensor: + assert self.delight_tokenizer is not None + assert self.delight_text_encoder is not None + prompts = [self.config.delight_prompt] + if use_cfg: + prompts.append(self.config.delight_negative_prompt) + text_inputs = self.delight_tokenizer( + prompts, + padding="max_length", + max_length=self.delight_tokenizer.model_max_length, + truncation=True, + return_tensors="pt", ) - self._renderer_loaded = True - logger.info("Mesh renderer initialized") + with self.use_declared_component( + component_name="delight_text_encoder", + module=self.delight_text_encoder, + phase="prompt", + ) as text_encoder: + assert isinstance(text_encoder, CLIPTextModel) + self.delight_text_encoder = text_encoder + output: BaseEncoderOutput = text_encoder( + input_ids=text_inputs.input_ids.to(device), + attention_mask=None, + ) + prompt_embeds = output.last_hidden_state + if prompt_embeds is None: + raise RuntimeError("The delight text encoder returned no hidden states.") + if not use_cfg: + return prompt_embeds + positive, negative = prompt_embeds.chunk(2) + return torch.cat([positive, negative, negative]) - def _render_multiview(self, mesh) -> tuple: - self._init_renderer() + @torch.no_grad() + def _run_delight(self, image: Image.Image) -> Image.Image: + assert self.delight_transformer is not None + assert self.delight_vae is not None + assert self.delight_scheduler is not None + + device = self.device + image, target_rgb, alpha = self._prepare_delight_target(image, device) + use_cfg = ( + self.config.delight_guidance_scale > 1 + and self.config.delight_cfg_image >= 1 + ) + prompt_embeds = self._encode_delight_prompts(use_cfg, device) + + processed_image = self._delight_image_processor.preprocess(image) + with self.use_declared_component( + component_name="delight_vae", + module=self.delight_vae, + phase="encode", + ) as vae: + assert isinstance(vae, AutoencoderKL) + self.delight_vae = vae + vae_dtype = _module_dtype(vae) + image_latents = vae.encode( + processed_image.to(device=device, dtype=vae_dtype) + ).latent_dist.mode() + + if use_cfg: + image_latents = torch.cat( + [image_latents, image_latents, torch.zeros_like(image_latents)] + ) + + scheduler = self.delight_scheduler + scheduler.set_timesteps(self.config.delight_num_inference_steps, device=device) + generator = torch.Generator(device="cpu").manual_seed(42) + latent_channels = self.delight_transformer.config.out_channels + latents = randn_tensor( + (1, latent_channels, image_latents.shape[-2], image_latents.shape[-1]), + generator=generator, + device=device, + dtype=prompt_embeds.dtype, + ) + latents *= scheduler.init_noise_sigma + step_kwargs = _scheduler_step_kwargs(scheduler, generator) + + with self.use_declared_component( + component_name="delight_transformer", + module=self.delight_transformer, + phase="denoise", + ) as transformer: + assert isinstance(transformer, StableDiffusionUNet2DConditionModel) + self.delight_transformer = transformer + for step_index, timestep in enumerate(scheduler.timesteps): + latent_input = torch.cat([latents] * 3) if use_cfg else latents + latent_input = scheduler.scale_model_input(latent_input, timestep) + latent_input = torch.cat([latent_input, image_latents], dim=1) + with set_forward_context( + current_timestep=step_index, + attn_metadata=None, + ): + noise_prediction = transformer( + latent_input, + timestep, + encoder_hidden_states=prompt_embeds, + return_dict=False, + )[0] + if use_cfg: + text, image_only, unconditioned = noise_prediction.chunk(3) + noise_prediction = ( + unconditioned + + self.config.delight_guidance_scale * (text - image_only) + + self.config.delight_cfg_image * (image_only - unconditioned) + ) + latents = scheduler.step( + noise_prediction, + timestep, + latents, + **step_kwargs, + return_dict=False, + )[0] + + with self.use_declared_component( + component_name="delight_vae", + module=self.delight_vae, + phase="decode", + ) as vae: + assert isinstance(vae, AutoencoderKL) + self.delight_vae = vae + scaling_factor = vae.config.arch_config.scaling_factor + decoded = vae.decode(latents / scaling_factor) + result = self._delight_image_processor.postprocess(decoded, output_type="pil")[ + 0 + ] + + source_rgb = torch.from_numpy(np.asarray(result, dtype=np.float32) / 255).to( + device + ) + corrected = _recorrect_rgb(source_rgb, target_rgb, alpha) + composited = corrected[..., :3] * corrected[..., 3:] + 1.0 * ( + 1.0 - corrected[..., 3:] + ) + return Image.fromarray( + (composited.clamp(0, 1).cpu().numpy() * 255).astype(np.uint8) + ) + + def _prepare_reference_image(self, image_path: str) -> Image.Image: + image = self._load_input_image(image_path) + if not self.config.delight_enable: + return image + return self._run_delight(image) + + def _render_multiview( + self, mesh: Any + ) -> tuple[list[Image.Image], list[Image.Image]]: + if self._renderer is None: + from sglang.multimodal_gen.runtime.utils.mesh3d_utils import MeshRender + + self._renderer = MeshRender( + default_resolution=self.config.paint_render_size, + texture_size=self.config.paint_texture_size, + device=self.device, + ) self._renderer.load_mesh(mesh) - normal_maps = self._renderer.render_normal_multiview( self.CAMERA_ELEVS, self.CAMERA_AZIMS, use_abs_coor=True ) position_maps = self._renderer.render_position_multiview( self.CAMERA_ELEVS, self.CAMERA_AZIMS ) - - logger.info(f"Rendered {len(normal_maps)} views for texture generation") return normal_maps, position_maps - # --- Forward --- - def forward(self, batch: Req, server_args: ServerArgs) -> Req: + del server_args if batch.extra.get("_mesh_failed"): - logger.warning("Mesh generation failed, skipping paint preprocessing") - batch.extra["paint_mesh"] = None - batch.extra["delighted_image"] = None - batch.extra["normal_maps"] = [] - batch.extra["position_maps"] = [] - batch.extra["camera_azims"] = self.CAMERA_AZIMS - batch.extra["camera_elevs"] = self.CAMERA_ELEVS - batch.extra["view_weights"] = self.VIEW_WEIGHTS - batch.extra["renderer"] = None + batch.extra.update( + { + "paint_mesh": None, + "delighted_image": None, + "normal_maps": [], + "position_maps": [], + "camera_azims": self.CAMERA_AZIMS, + "camera_elevs": self.CAMERA_ELEVS, + "view_weights": self.VIEW_WEIGHTS, + "renderer": None, + } + ) return batch - import concurrent.futures - import copy - - # 1. UV unwrap + delight in parallel - batch_for_uv = batch - batch_for_delight = copy.copy(batch) - batch_for_delight.extra = batch.extra.copy() + mesh = batch.extra["shape_meshes"] + if isinstance(mesh, list): + mesh = mesh[0] + if isinstance(mesh, list): + mesh = mesh[0] + image_path = batch.image_path + if not isinstance(image_path, str): + raise TypeError("Hunyuan3D Paint expects one image path per request.") with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: - uv_future = executor.submit(self._do_uv_unwrap, batch_for_uv, server_args) - delight_future = executor.submit( - self._do_delight, batch_for_delight, server_args - ) - uv_future.result() - delight_future.result() - - batch.extra["paint_mesh"] = batch_for_uv.extra.get("paint_mesh") - batch.extra["delighted_image"] = batch_for_delight.extra.get("delighted_image") - - # 2. Multi-view rendering - normal_maps, position_maps = self._render_multiview(batch.extra["paint_mesh"]) - batch.extra["normal_maps"] = normal_maps - batch.extra["position_maps"] = position_maps - batch.extra["camera_azims"] = self.CAMERA_AZIMS - batch.extra["camera_elevs"] = self.CAMERA_ELEVS - batch.extra["view_weights"] = self.VIEW_WEIGHTS - batch.extra["renderer"] = self._renderer + mesh_future = executor.submit(self._unwrap_mesh, mesh) + image_future = executor.submit(self._prepare_reference_image, image_path) + paint_mesh = mesh_future.result() + delighted_image = image_future.result() + normal_maps, position_maps = self._render_multiview(paint_mesh) + batch.extra.update( + { + "paint_mesh": paint_mesh, + "delighted_image": delighted_image, + "normal_maps": normal_maps, + "position_maps": position_maps, + "camera_azims": self.CAMERA_AZIMS, + "camera_elevs": self.CAMERA_ELEVS, + "view_weights": self.VIEW_WEIGHTS, + "renderer": self._renderer, + } + ) return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + del server_args result = VerificationResult() result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) result.add_check("image_path", batch.image_path, V.not_none) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + del server_args + if batch.extra.get("_mesh_failed"): + return VerificationResult() result = VerificationResult() result.add_check("paint_mesh", batch.extra.get("paint_mesh"), V.not_none) result.add_check( @@ -495,494 +445,301 @@ class Hunyuan3DPaintPreprocessStage(PipelineStage): return result -# Stage 2: TexGen (model loading + input prep + denoising + decode) class Hunyuan3DPaintTexGenStage(PipelineStage): + """Generate consistent multi-view textures from geometry controls.""" + def __init__( self, config: Hunyuan3D2PipelineConfig, - paint_dir: str | None = None, - transformer: Any = None, - scheduler: Any = None, - vae: Any = None, - vae_scale_factor: int = 8, - image_processor: Any = None, - solver: Any = None, - is_turbo: bool = False, + transformer: Hunyuan3DPaintUNet, + scheduler: Any, + vae: AutoencoderKL, ) -> None: super().__init__() self.config = config - self.paint_dir = paint_dir self.transformer = transformer self.scheduler = scheduler self.vae = vae - self.vae_scale_factor = vae_scale_factor - self.image_processor = image_processor - self.solver = solver - self.is_turbo = is_turbo - self._loaded = transformer is not None + block_channels = vae.config.arch_config.block_out_channels + self.vae_scale_factor = 2 ** (len(block_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) @property def parallelism_type(self) -> StageParallelismType: return StageParallelismType.MAIN_RANK_ONLY - def _load_paint_models(self, server_args: ServerArgs) -> None: - """Load paint models from pre-resolved local path (no network).""" - if self._loaded: - return - if self.paint_dir is None: - logger.warning("No paint model directory resolved, skipping") - self._loaded = True - return - try: - self._do_load_paint(server_args) - logger.info("Paint pipeline loaded successfully") - except Exception as e: - logger.warning("Failed to load paint pipeline: %s", e) - self.vae = None - self.transformer = None - self.scheduler = None - self._loaded = True - - def _do_load_paint(self, server_args: ServerArgs) -> None: - import json - - from diffusers import AutoencoderKL - from diffusers.image_processor import VaeImageProcessor - - from sglang.multimodal_gen.runtime.models.dits.hunyuan3d import ( - UNet2p5DConditionModel, - ) - - local_path = self.paint_dir - logger.info("Loading paint model from %s", local_path) - vae_dir = os.path.join(local_path, "vae") - with open(os.path.join(vae_dir, "config.json"), "r") as f: - vae_config = json.load(f) - vae_config = {k: v for k, v in vae_config.items() if not k.startswith("_")} - self.vae = AutoencoderKL(**vae_config) - st_path = os.path.join(vae_dir, "diffusion_pytorch_model.safetensors") - bin_path = os.path.join(vae_dir, "diffusion_pytorch_model.bin") - if os.path.exists(st_path): - from safetensors.torch import load_file - - state_dict = load_file(st_path) - elif os.path.exists(bin_path): - state_dict = torch.load(bin_path, map_location="cpu", weights_only=True) - else: - raise FileNotFoundError(f"No VAE weights in {vae_dir}") - self.vae.load_state_dict(state_dict) - # Resolve the DiT (multiview UNet) dtype from config, with CPU/MPS fallback. - dit_dtype = PRECISION_TO_TYPE.get( - getattr(self.config, "dit_precision", "fp16"), torch.float16 - ) - if self.device.type in ("cpu", "mps") and dit_dtype in ( - torch.float16, - torch.bfloat16, - ): - dit_dtype = torch.float32 - # The multiview (Stable-Diffusion) AutoencoderKL must share the UNet dtype. - # Reference attention feeds its VAE-encoded ref_latents straight into the - # fp16 UNet, and the official HunyuanPaint pipeline runs VAE+UNet entirely - # in fp16. The `vae_precision` knob targets the 3D ShapeVAE (geometry - # precision) — applying it to this 2D texture VAE produces an - # fp32-input / fp16-weight mismatch that crashes the paint UNet. - vae_dtype = dit_dtype - - self.vae = self.vae.to(device=self.device, dtype=vae_dtype).eval() - self.transformer = UNet2p5DConditionModel.from_pretrained( - os.path.join(local_path, "unet"), - torch_dtype=dit_dtype, - ).to(self.device) - self.is_turbo = bool(getattr(self.config, "paint_turbo_mode", False)) - sched_path = os.path.join(local_path, "scheduler", "scheduler_config.json") - with open(sched_path, "r") as f: - sched_cfg = json.load(f) - if self.is_turbo: - from diffusers import LCMScheduler - - self.scheduler = LCMScheduler.from_config(sched_cfg) - else: - from diffusers import EulerAncestralDiscreteScheduler - - self.scheduler = EulerAncestralDiscreteScheduler.from_config( - sched_cfg, timestep_spacing="trailing" - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.solver = DDIMSolver( - self.scheduler.alphas_cumprod.cpu().numpy(), - timesteps=self.scheduler.config.num_train_timesteps, - ddim_timesteps=30, - ).to(self.device) - if server_args.enable_torch_compile: - dit_config = getattr(server_args.pipeline_config, "dit_config", None) - compile_mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE") or getattr( - dit_config, - "torch_compile_mode", - "max-autotune-no-cudagraphs", - ) - logger.info("Compiling paint transformer with mode: %s", compile_mode) - self.transformer.compile(mode=compile_mode, fullgraph=False, dynamic=None) - - def _convert_pil_list_to_tensor( - self, images: list, device: torch.device - ) -> torch.Tensor: - bg_c = [1.0, 1.0, 1.0] - images_tensor = [] - for batch_imgs in images: - view_imgs = [] - for pil_img in batch_imgs: - if pil_img.mode == "L": - pil_img = pil_img.point( - lambda x: 255 if x > 1 else 0, mode="1" - ).convert("RGB") - img = np.asarray(pil_img, dtype=np.float32) / 255.0 - if img.shape[2] > 3: - alpha = img[:, :, 3:] - img = img[:, :, :3] * alpha + bg_c * (1 - alpha) - img = ( - torch.from_numpy(img) - .permute(2, 0, 1) - .unsqueeze(0) - .contiguous() - .to(device=device, dtype=self.vae.dtype) - ) - view_imgs.append(img) - view_imgs = torch.cat(view_imgs, dim=0) - images_tensor.append(view_imgs.unsqueeze(0)) - return torch.cat(images_tensor, dim=0) - - @torch.no_grad() - def _encode_images(self, images: torch.Tensor) -> torch.Tensor: - batch_size = images.shape[0] - images = rearrange(images, "b n c h w -> (b n) c h w") - dtype = next(self.vae.parameters()).dtype - images = (images - 0.5) * 2.0 - posterior = self.vae.encode(images.to(dtype)).latent_dist - latents = posterior.sample() * self.vae.config.scaling_factor - return rearrange(latents, "(b n) c h w -> b n c h w", b=batch_size) + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + del server_args + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse(stage_name, "paint_vae", phase="encode"), + ComponentUse( + stage_name, + "paint_transformer", + phase="denoise", + memory_intensive=True, + ), + ComponentUse(stage_name, "paint_vae", phase="decode"), + ] @staticmethod - def _compute_camera_index(azim: float, elev: float) -> int: - base_idx = int(((azim // 30) + 9) % 12) - if elev == 0: + def _pil_views_to_tensor( + images: list[Image.Image], + size: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + tensors = [] + for image in images: + image = image.resize((size, size), Image.Resampling.BICUBIC) + if image.mode == "L": + image = image.point(lambda value: 255 if value > 1 else 0).convert( + "RGB" + ) + pixels = np.asarray(image, dtype=np.float32) / 255 + if pixels.shape[-1] == 4: + alpha = pixels[..., 3:] + pixels = pixels[..., :3] * alpha + (1 - alpha) + tensor = torch.from_numpy(pixels).permute(2, 0, 1).contiguous() + tensors.append(tensor) + return torch.stack(tensors).unsqueeze(0).to(device=device, dtype=dtype) + + @staticmethod + def _encode_images( + vae: AutoencoderKL, + images: torch.Tensor, + generator: torch.Generator, + ) -> torch.Tensor: + batch_size, num_images = images.shape[:2] + images = rearrange(images, "b n c h w -> (b n) c h w") + images = images.mul(2).sub(1) + posterior = vae.encode(images).latent_dist + scaling_factor = vae.config.arch_config.scaling_factor + latents = posterior.sample(generator=generator) * scaling_factor + return rearrange( + latents, + "(b n) c h w -> b n c h w", + b=batch_size, + n=num_images, + ) + + @staticmethod + def _camera_index(azimuth: float, elevation: float) -> int: + base_index = int(((azimuth // 30) + 9) % 12) + if elevation == 0: base, divisor = 12, 1 - elif elev == 20: + elif elevation == 20: base, divisor = 24, 1 - elif elev == -20: + elif elevation == -20: base, divisor = 0, 1 - elif elev == 90: + elif elevation == 90: base, divisor = 40, 3 - elif elev == -90: + elif elevation == -90: base, divisor = 36, 3 else: base, divisor = 12, 1 - return base + (base_idx // divisor) + return base + base_index // divisor - def _prepare_denoising_inputs( - self, - batch: Req, - server_args: ServerArgs, - ) -> dict[str, Any]: - import random + def _timesteps(self, device: torch.device) -> torch.Tensor: + if not self.config.paint_turbo_mode: + self.scheduler.set_timesteps( + self.config.paint_num_inference_steps, device=device + ) + return self.scheduler.timesteps - from diffusers.utils.torch_utils import randn_tensor + self.scheduler.set_timesteps( + num_inference_steps=10, + original_inference_steps=30, + device=device, + ) + return self.scheduler.timesteps + def _prepare_denoising_inputs(self, batch: Req) -> PaintDenoisingInputs: device = self.device + render_size = self.config.paint_resolution normal_maps = batch.extra["normal_maps"] position_maps = batch.extra["position_maps"] - camera_azims = batch.extra["camera_azims"] - camera_elevs = batch.extra["camera_elevs"] + if not isinstance(normal_maps, list) or not isinstance(position_maps, list): + raise TypeError("Hunyuan3D Paint geometry controls must be image lists.") - num_steps = self.config.paint_num_inference_steps - guidance_scale = self.config.paint_guidance_scale - render_size = self.config.paint_resolution - num_in_batch = len(normal_maps) - - seed = 0 - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - generator = torch.Generator(device=device).manual_seed(seed) - - image = batch.extra["delighted_image"] - if not isinstance(image, list): - image = [image] - image = [to_rgb_image(img) for img in image] - - image_vae = [ - torch.tensor(np.array(img, dtype=np.float32) / 255.0) for img in image + reference = batch.extra["delighted_image"] + references = reference if isinstance(reference, list) else [reference] + reference_images = [ + _to_rgb_image(image).resize( + (render_size, render_size), Image.Resampling.BICUBIC + ) + for image in references ] - image_vae = [ - iv.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(0) for iv in image_vae - ] - image_vae = torch.cat(image_vae, dim=1).to(device=device, dtype=self.vae.dtype) - ref_latents = self._encode_images(image_vae) + vae_generator = torch.Generator(device=device).manual_seed(0) + latent_generator = torch.Generator(device=device).manual_seed(0) - target_size = render_size - if isinstance(normal_maps, list): - normal_maps = [ - ( - img.resize((target_size, target_size)) - if hasattr(img, "resize") - else img - ) - for img in normal_maps - ] - normal_maps = self._convert_pil_list_to_tensor([normal_maps], device) - if isinstance(position_maps, list): - position_maps = [ - ( - img.resize((target_size, target_size)) - if hasattr(img, "resize") - else img - ) - for img in position_maps - ] - position_maps = self._convert_pil_list_to_tensor([position_maps], device) - - normal_imgs = ( - self._encode_images(normal_maps) if normal_maps is not None else None - ) - position_imgs = ( - self._encode_images(position_maps) if position_maps is not None else None - ) + with self.use_declared_component( + component_name="paint_vae", module=self.vae, phase="encode" + ) as vae: + assert isinstance(vae, AutoencoderKL) + self.vae = vae + vae_dtype = _module_dtype(vae) + reference_tensor = self._pil_views_to_tensor( + reference_images, render_size, device, vae_dtype + ) + normal_tensor = self._pil_views_to_tensor( + normal_maps, render_size, device, vae_dtype + ) + position_tensor = self._pil_views_to_tensor( + position_maps, render_size, device, vae_dtype + ) + reference_latents = self._encode_images( + vae, reference_tensor, vae_generator + ) + normal_latents = self._encode_images(vae, normal_tensor, vae_generator) + position_latents = self._encode_images(vae, position_tensor, vae_generator) camera_info = [ - self._compute_camera_index(azim, elev) - for azim, elev in zip(camera_azims, camera_elevs) + self._camera_index(azimuth, elevation) + for azimuth, elevation in zip( + batch.extra["camera_azims"], batch.extra["camera_elevs"] + ) ] camera_info_gen = torch.tensor([camera_info], device=device, dtype=torch.int64) - camera_info_ref = torch.tensor([[0]], device=device, dtype=torch.int64) - - do_cfg = guidance_scale > 1 and not self.is_turbo - - if self.is_turbo and position_maps is not None: - from sglang.multimodal_gen.runtime.models.dits.hunyuan3d import ( - compute_multi_resolution_discrete_voxel_indice, - compute_multi_resolution_mask, - ) - - position_attn_mask = compute_multi_resolution_mask(position_maps) - position_voxel_indices = compute_multi_resolution_discrete_voxel_indice( - position_maps - ) - else: - position_attn_mask = None - position_voxel_indices = None - - if do_cfg: - negative_ref_latents = torch.zeros_like(ref_latents) - ref_latents = torch.cat([negative_ref_latents, ref_latents]) - ref_scale = torch.as_tensor([0.0, 1.0]).to(ref_latents) - if normal_imgs is not None: - normal_imgs = torch.cat((normal_imgs, normal_imgs)) - if position_imgs is not None: - position_imgs = torch.cat((position_imgs, position_imgs)) - if position_maps is not None: - position_maps = torch.cat((position_maps, position_maps)) - camera_info_gen = torch.cat((camera_info_gen, camera_info_gen)) - camera_info_ref = torch.cat((camera_info_ref, camera_info_ref)) - else: - ref_scale = None - - model_kwargs = { - "ref_latents": ref_latents, - "num_in_batch": num_in_batch, - } - if ref_scale is not None: - model_kwargs["ref_scale"] = ref_scale - if normal_imgs is not None: - model_kwargs["normal_imgs"] = normal_imgs - if position_imgs is not None: - model_kwargs["position_imgs"] = position_imgs - if position_maps is not None: - model_kwargs["position_maps"] = position_maps - model_kwargs["camera_info_gen"] = camera_info_gen - model_kwargs["camera_info_ref"] = camera_info_ref - if position_attn_mask is not None: - model_kwargs["position_attn_mask"] = position_attn_mask - if position_voxel_indices is not None: - model_kwargs["position_voxel_indices"] = position_voxel_indices - - prompt_embeds = self.transformer.learned_text_clip_gen.repeat(1, 1, 1) - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - scheduler = self.scheduler - - if self.is_turbo: - bsz = 3 - index = torch.arange(29, -1, -bsz, device=device).long() - timesteps = self.solver.ddim_timesteps[index] - scheduler.set_timesteps(timesteps=timesteps.cpu(), device=device) - timesteps = scheduler.timesteps - else: - timesteps, num_steps = retrieve_timesteps( - scheduler, num_steps, device, None, None - ) - - num_channels_latents = self.transformer.config.in_channels - latent_shape = ( - num_in_batch, - num_channels_latents, - render_size // self.vae_scale_factor, - render_size // self.vae_scale_factor, + camera_info_ref = torch.zeros((1, 1), device=device, dtype=torch.int64) + use_cfg = ( + self.config.paint_guidance_scale > 1 and not self.config.paint_turbo_mode ) + + position_attention_mask = None + if self.config.paint_turbo_mode: + position_attention_mask = compute_multi_resolution_mask(position_tensor) + + reference_scale: torch.Tensor | float = 1.0 + if use_cfg: + reference_latents = torch.cat( + [torch.zeros_like(reference_latents), reference_latents] + ) + reference_scale = torch.as_tensor( + [0.0, 1.0], device=device, dtype=reference_latents.dtype + ) + normal_latents = torch.cat([normal_latents, normal_latents]) + position_latents = torch.cat([position_latents, position_latents]) + camera_info_gen = torch.cat([camera_info_gen, camera_info_gen]) + camera_info_ref = torch.cat([camera_info_ref, camera_info_ref]) + + num_views = len(normal_maps) + model_kwargs: dict[str, Any] = { + "ref_latents": reference_latents, + "num_in_batch": num_views, + "condition_embed_dict": {}, + "normal_imgs": normal_latents, + "position_imgs": position_latents, + "camera_info_gen": camera_info_gen, + "camera_info_ref": camera_info_ref, + "ref_scale": reference_scale, + } + if position_attention_mask is not None: + model_kwargs["position_attn_mask"] = position_attention_mask + + timesteps = self._timesteps(device) + latent_channels = self.transformer.config.in_channels + latent_size = render_size // self.vae_scale_factor latents = randn_tensor( - latent_shape, generator=generator, device=device, dtype=prompt_embeds.dtype + (num_views, latent_channels, latent_size, latent_size), + generator=latent_generator, + device=device, + dtype=_module_dtype(self.transformer), + ) + latents *= self.scheduler.init_noise_sigma + return PaintDenoisingInputs( + timesteps=timesteps, + latents=latents, + model_kwargs=model_kwargs, + num_views=num_views, + guidance_scale=self.config.paint_guidance_scale, + use_cfg=use_cfg, + generator=latent_generator, + latent_channels=latent_channels, ) - latents = latents * scheduler.init_noise_sigma - - return { - "scheduler": scheduler, - "timesteps": timesteps, - "latents": latents, - "prompt_embeds": prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "model_kwargs": model_kwargs, - "num_in_batch": num_in_batch, - "num_inference_steps": num_steps, - "guidance_scale": guidance_scale, - "do_cfg": do_cfg, - "generator": generator, - "num_channels_latents": num_channels_latents, - } @torch.no_grad() - def _denoise_loop( - self, - timesteps: torch.Tensor, - latents: torch.Tensor, - prompt_embeds: torch.Tensor, - negative_prompt_embeds: torch.Tensor, - model_kwargs: dict[str, Any], - num_in_batch: int, - guidance_scale: float, - do_cfg: bool, - generator: torch.Generator, - num_channels_latents: int, - scheduler: Any, - ) -> torch.Tensor: - import inspect - - if do_cfg: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - extra_step_kwargs = {} - if "eta" in inspect.signature(scheduler.step).parameters: - extra_step_kwargs["eta"] = 0.0 - if "generator" in inspect.signature(scheduler.step).parameters: - extra_step_kwargs["generator"] = generator - - for step_idx, t in enumerate(timesteps): - latents = rearrange(latents, "(b n) c h w -> b n c h w", n=num_in_batch) - latent_model_input = torch.cat([latents] * 2) if do_cfg else latents - latent_model_input = rearrange( - latent_model_input, "b n c h w -> (b n) c h w" - ) - latent_model_input = scheduler.scale_model_input(latent_model_input, t) - latent_model_input = rearrange( - latent_model_input, "(b n) c h w -> b n c h w", n=num_in_batch - ) - - with set_forward_context( - current_timestep=step_idx, - attn_metadata=None, - ): - noise_pred = self.transformer( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=None, - cross_attention_kwargs=None, - added_cond_kwargs=None, - return_dict=False, - **model_kwargs, - )[0] - - latents = rearrange(latents, "b n c h w -> (b n) c h w") - - if do_cfg: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * ( - noise_pred_text - noise_pred_uncond + def _denoise(self, inputs: PaintDenoisingInputs) -> torch.Tensor: + scheduler = self.scheduler + latents = inputs.latents + step_kwargs = _scheduler_step_kwargs(scheduler, inputs.generator) + with self.use_declared_component( + component_name="paint_transformer", + module=self.transformer, + phase="denoise", + ) as transformer: + assert isinstance(transformer, Hunyuan3DPaintUNet) + self.transformer = transformer + prompt_embeds = transformer.learned_text_clip_gen + if inputs.use_cfg: + prompt_embeds = torch.cat( + [torch.zeros_like(prompt_embeds), prompt_embeds] ) - latents = scheduler.step( - noise_pred, - t, - latents[:, :num_channels_latents, :, :], - **extra_step_kwargs, - return_dict=False, - )[0] - + for step_index, timestep in enumerate(inputs.timesteps): + latents = rearrange( + latents, "(b n) c h w -> b n c h w", n=inputs.num_views + ) + latent_input = ( + torch.cat([latents, latents]) if inputs.use_cfg else latents + ) + latent_input = rearrange(latent_input, "b n c h w -> (b n) c h w") + latent_input = scheduler.scale_model_input(latent_input, timestep) + latent_input = rearrange( + latent_input, + "(b n) c h w -> b n c h w", + n=inputs.num_views, + ) + with set_forward_context( + current_timestep=step_index, + attn_metadata=None, + ): + noise_prediction = transformer( + latent_input, + timestep, + encoder_hidden_states=prompt_embeds, + return_dict=False, + **inputs.model_kwargs, + )[0] + latents = rearrange(latents, "b n c h w -> (b n) c h w") + if inputs.use_cfg: + unconditioned, conditioned = noise_prediction.chunk(2) + noise_prediction = unconditioned + inputs.guidance_scale * ( + conditioned - unconditioned + ) + latents = scheduler.step( + noise_prediction, + timestep, + latents[:, : inputs.latent_channels], + **step_kwargs, + return_dict=False, + )[0] return latents @torch.no_grad() - def _decode_latents(self, latents: torch.Tensor) -> list: - image = self.vae.decode( - latents / self.vae.config.scaling_factor, return_dict=False - )[0] - return self.image_processor.postprocess(image, output_type="pil") + def _decode(self, latents: torch.Tensor) -> list[Image.Image]: + with self.use_declared_component( + component_name="paint_vae", module=self.vae, phase="decode" + ) as vae: + assert isinstance(vae, AutoencoderKL) + self.vae = vae + scaling_factor = vae.config.arch_config.scaling_factor + decoded = vae.decode(latents / scaling_factor) + return self.image_processor.postprocess(decoded, output_type="pil") def forward(self, batch: Req, server_args: ServerArgs) -> Req: + del server_args if batch.extra.get("_mesh_failed"): - logger.warning("Mesh generation failed, skipping paint texgen") batch.extra["multiview_textures"] = [] return batch - - self._load_paint_models(server_args) - - delighted_image = batch.extra["delighted_image"] - normal_maps = batch.extra["normal_maps"] - - if self.transformer is not None: - try: - prepared = self._prepare_denoising_inputs(batch, server_args) - - latents = self._denoise_loop( - timesteps=prepared["timesteps"], - latents=prepared["latents"], - prompt_embeds=prepared["prompt_embeds"], - negative_prompt_embeds=prepared["negative_prompt_embeds"], - model_kwargs=prepared["model_kwargs"], - num_in_batch=prepared["num_in_batch"], - guidance_scale=prepared["guidance_scale"], - do_cfg=prepared["do_cfg"], - generator=prepared["generator"], - num_channels_latents=prepared["num_channels_latents"], - scheduler=prepared["scheduler"], - ) - - multiview_textures = self._decode_latents(latents) - logger.info( - "Paint pipeline generated %d textures", len(multiview_textures) - ) - - except Exception as e: - logger.error(f"Paint pipeline execution failed: {e}") - import traceback - - traceback.print_exc() - render_size = self.config.paint_resolution - multiview_textures = [ - delighted_image.resize((render_size, render_size)) - for _ in range(len(normal_maps)) - ] - else: - logger.warning( - "Paint pipeline not available, using reference image for all views" - ) - render_size = self.config.paint_resolution - multiview_textures = [ - delighted_image.resize((render_size, render_size)) - for _ in range(len(normal_maps)) - ] - - batch.extra["multiview_textures"] = multiview_textures - logger.info(f"Generated {len(multiview_textures)} texture views") + inputs = self._prepare_denoising_inputs(batch) + batch.extra["multiview_textures"] = self._decode(self._denoise(inputs)) return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + del server_args if batch.extra.get("_mesh_failed"): return VerificationResult() result = VerificationResult() @@ -996,6 +753,7 @@ class Hunyuan3DPaintTexGenStage(PipelineStage): return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + del server_args result = VerificationResult() result.add_check( "multiview_textures", batch.extra.get("multiview_textures"), V.is_list @@ -1003,86 +761,72 @@ class Hunyuan3DPaintTexGenStage(PipelineStage): return result -# Stage 3: Postprocess (texture baking + mesh export) class Hunyuan3DPaintPostprocessStage(PipelineStage): - """Texture baking from multi-view images and final mesh export.""" - - @property - def parallelism_type(self) -> StageParallelismType: - return StageParallelismType.MAIN_RANK_ONLY + """Bake generated views into a texture and export the final mesh.""" def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: super().__init__() self.config = config + @property + def parallelism_type(self) -> StageParallelismType: + return StageParallelismType.MAIN_RANK_ONLY + + @staticmethod + def _cleanup_obj_artifacts(obj_path: str, files_before_export: set[str]) -> None: + obj_dir = os.path.dirname(obj_path) or "." + generated_files = set(os.listdir(obj_dir)) - files_before_export + cleanup_paths = {obj_path} + cleanup_paths.update( + os.path.join(obj_dir, filename) + for filename in generated_files + if filename.endswith(".mtl") or filename.endswith(".png") + ) + for path in cleanup_paths: + try: + os.remove(path) + except FileNotFoundError: + continue + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: - if batch.extra.get("_mesh_failed"): - logger.warning("Mesh generation failed, skipping paint postprocess") + del server_args + if batch.is_warmup or batch.extra.get("_mesh_failed"): return OutputBatch(output_file_paths=[], metrics=batch.metrics) renderer = batch.extra["renderer"] - multiview_textures = batch.extra["multiview_textures"] - camera_elevs = batch.extra["camera_elevs"] - camera_azims = batch.extra["camera_azims"] - view_weights = batch.extra["view_weights"] - - render_size = getattr(self.config, "paint_render_size", 2048) - resized_textures = [] - for tex in multiview_textures: - if hasattr(tex, "resize"): - resized_textures.append(tex.resize((render_size, render_size))) - else: - resized_textures.append(tex) - - try: - texture, mask = renderer.bake_from_multiview( - resized_textures, - camera_elevs, - camera_azims, - view_weights, - method="fast", + textures = [ + image.resize( + (self.config.paint_render_size, self.config.paint_render_size), + Image.Resampling.BICUBIC, ) - - mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype("uint8") - texture = renderer.texture_inpaint(texture, mask_np) - - renderer.set_texture(texture) - textured_mesh = renderer.save_mesh() - logger.info("Texture baking completed") - except Exception as e: - logger.error(f"Texture baking failed: {e}") - textured_mesh = batch.extra["paint_mesh"] + for image in batch.extra["multiview_textures"] + ] + texture, mask = renderer.bake_from_multiview( + textures, + batch.extra["camera_elevs"], + batch.extra["camera_azims"], + batch.extra["view_weights"], + method="fast", + ) + mask_array = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8) + texture = renderer.texture_inpaint(texture, mask_array) + renderer.set_texture(texture) + textured_mesh = renderer.save_mesh() obj_path = batch.extra["shape_obj_path"] return_path = batch.extra["shape_return_path"] - - try: - textured_mesh.export(obj_path) - if self.config.paint_save_glb: - glb_path = obj_path[:-4] + ".glb" - textured_mesh.export(glb_path) - return_path = glb_path - self._cleanup_obj_artifacts(obj_path) - except Exception as e: - logger.error(f"Mesh export failed: {e}") + obj_dir = os.path.dirname(obj_path) or "." + files_before_export = set(os.listdir(obj_dir)) + textured_mesh.export(obj_path) + if self.config.paint_save_glb: + return_path = os.path.splitext(obj_path)[0] + ".glb" + textured_mesh.export(return_path) + self._cleanup_obj_artifacts(obj_path, files_before_export) return OutputBatch(output_file_paths=[return_path], metrics=batch.metrics) - @staticmethod - def _cleanup_obj_artifacts(obj_path: str) -> None: - """Remove OBJ file and trimesh-generated material artifacts.""" - obj_dir = os.path.dirname(obj_path) or "." - targets = [obj_path] - for f in os.listdir(obj_dir): - if f.endswith(".mtl") or (f.startswith("material") and f.endswith(".png")): - targets.append(os.path.join(obj_dir, f)) - for path in targets: - try: - os.remove(path) - except OSError: - pass - def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + del server_args if batch.extra.get("_mesh_failed"): return VerificationResult() result = VerificationResult() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py index 7eabec492..79ffd6d02 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py @@ -551,6 +551,11 @@ class Hunyuan3DShapeSaveStage(PipelineStage): "The surface level may be outside the volume data range." ) + if batch.is_warmup: + if self.config.paint_enable: + return batch + return OutputBatch(output_file_paths=[], metrics=batch.metrics) + obj_path, return_path = self._get_output_paths(batch) output_dir = os.path.dirname(obj_path) if output_dir: diff --git a/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py b/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py new file mode 100644 index 000000000..e3103dda0 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace + +import torch +from diffusers import AutoencoderKL as DiffusersAutoencoderKL +from diffusers import LCMScheduler, UNet2DConditionModel + +from sglang.multimodal_gen.configs.models.vaes.stable_diffusion import ( + StableDiffusionVAEConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, +) +from sglang.multimodal_gen.runtime.models.dits.hunyuan3d_paint import ( + Hunyuan3DPaintUNet, +) +from sglang.multimodal_gen.runtime.models.dits.stable_diffusion import ( + StableDiffusionUNet2DConditionModel, + StableDiffusionUNetConfig, +) +from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.paint import ( + Hunyuan3DPaintPostprocessStage, + Hunyuan3DPaintTexGenStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape import ( + Hunyuan3DShapeSaveStage, +) + + +def _unet_config() -> dict: + return { + "sample_size": 8, + "in_channels": 4, + "out_channels": 4, + "center_input_sample": False, + "flip_sin_to_cos": True, + "freq_shift": 0, + "down_block_types": ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + "up_block_types": ( + "UpBlock2D", + "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", + ), + "block_out_channels": (32, 32, 32, 32), + "layers_per_block": 2, + "downsample_padding": 1, + "dropout": 0.0, + "norm_num_groups": 8, + "norm_eps": 1e-5, + "cross_attention_dim": 16, + "attention_head_dim": (4, 4, 4, 4), + "transformer_layers_per_block": 1, + "use_linear_projection": True, + } + + +class TestNativeStableDiffusionUNet(unittest.TestCase): + def test_matches_diffusers_sd21_layout_and_forward(self): + torch.manual_seed(0) + raw_config = _unet_config() + reference = UNet2DConditionModel(**raw_config).eval() + native = StableDiffusionUNet2DConditionModel( + StableDiffusionUNetConfig.from_dict(raw_config) + ).eval() + native.load_state_dict(reference.state_dict(), strict=True) + + sample = torch.randn(1, 4, 8, 8) + timestep = torch.tensor([10]) + encoder_hidden_states = torch.randn(1, 5, 16) + class_labels = torch.tensor([0]) + with torch.inference_mode(): + expected = reference( + sample, + timestep, + encoder_hidden_states, + class_labels=class_labels, + ).sample + actual = native( + sample, + timestep, + encoder_hidden_states, + class_labels=class_labels, + ).sample + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + def test_paint_reference_branch_and_layer_groups(self): + config = StableDiffusionUNetConfig.from_dict(_unet_config()) + model = Hunyuan3DPaintUNet(config).eval() + modules = dict(model.named_modules()) + self.assertTrue(all(name in modules for name in model.layer_names)) + + sample = torch.randn(1, 2, 4, 8, 8) + prompt = model.learned_text_clip_gen + condition_cache: dict[str, torch.Tensor] = {} + with torch.inference_mode(): + output = model( + sample, + torch.tensor(10), + prompt, + ref_latents=torch.randn(1, 1, 4, 8, 8), + num_in_batch=2, + condition_embed_dict=condition_cache, + normal_imgs=torch.randn(1, 2, 4, 8, 8), + position_imgs=torch.randn(1, 2, 4, 8, 8), + camera_info_gen=torch.tensor([[12, 15]]), + camera_info_ref=torch.tensor([[0]]), + ).sample + + self.assertEqual(output.shape, (2, 4, 8, 8)) + self.assertTrue(condition_cache) + + +class TestNativeStableDiffusionVAE(unittest.TestCase): + def test_old_diffusers_config_defaults_and_forward(self): + raw_config = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 4, + "sample_size": 8, + "block_out_channels": (32, 32), + "layers_per_block": 1, + "act_fn": "silu", + "norm_num_groups": 8, + "down_block_types": ("DownEncoderBlock2D", "DownEncoderBlock2D"), + "up_block_types": ("UpDecoderBlock2D", "UpDecoderBlock2D"), + } + reference = DiffusersAutoencoderKL(**raw_config).eval() + config = StableDiffusionVAEConfig() + config.update_model_arch(raw_config) + native = AutoencoderKL(config).eval() + native.load_state_dict(reference.state_dict(), strict=True) + + image = torch.randn(1, 3, 8, 8) + latent = torch.randn(1, 4, 4, 4) + with torch.inference_mode(): + expected_posterior = reference.encode(image).latent_dist + actual_posterior = native.encode(image).latent_dist + expected_decoded = reference.decode(latent).sample + actual_decoded = native.decode(latent) + + torch.testing.assert_close( + actual_posterior.parameters, + expected_posterior.parameters, + rtol=1e-5, + atol=1e-5, + ) + torch.testing.assert_close( + actual_decoded, expected_decoded, rtol=1e-5, atol=1e-5 + ) + + +class TestHunyuan3DWarmupOutput(unittest.TestCase): + @staticmethod + def _batch(): + return SimpleNamespace( + extra={"shape_meshes": [object()]}, + is_warmup=True, + metrics=None, + ) + + def test_shape_save_does_not_require_output_path_during_paint_warmup(self): + batch = self._batch() + stage = Hunyuan3DShapeSaveStage(Hunyuan3D2PipelineConfig(paint_enable=True)) + + self.assertIs(stage.forward(batch, SimpleNamespace()), batch) + + def test_shape_only_warmup_returns_no_files(self): + stage = Hunyuan3DShapeSaveStage(Hunyuan3D2PipelineConfig(paint_enable=False)) + + output = stage.forward(self._batch(), SimpleNamespace()) + + self.assertEqual(output.output_file_paths, []) + + def test_paint_postprocess_skips_export_during_warmup(self): + stage = Hunyuan3DPaintPostprocessStage(Hunyuan3D2PipelineConfig()) + + output = stage.forward(self._batch(), SimpleNamespace()) + + self.assertEqual(output.output_file_paths, []) + + +class TestHunyuan3DPaintTurboSchedule(unittest.TestCase): + def test_uses_standard_lcm_schedule_without_custom_timesteps(self): + stage = Hunyuan3DPaintTexGenStage.__new__(Hunyuan3DPaintTexGenStage) + stage.config = Hunyuan3D2PipelineConfig(paint_turbo_mode=True) + stage.scheduler = LCMScheduler( + num_train_timesteps=1000, + original_inference_steps=50, + ) + + timesteps = stage._timesteps(torch.device("cpu")) + + self.assertEqual( + timesteps.tolist(), + [989, 890, 791, 692, 593, 494, 395, 296, 197, 98], + ) + self.assertFalse(stage.scheduler.custom_timesteps) + + +if __name__ == "__main__": + unittest.main()