minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] chore: use native hunyuan3d paint and delight models (#34980)

Browse Source
This commit is contained in:
Mick
2026-08-16 10:03:48 +08:00
committed by GitHub
parent d106e8b23a
commit 4f9da62547
11 changed files with 2442 additions and 1767 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/vaes/stable_diffusion.py
+37
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py
+7
View File
@@ -4,6 +4,7 @@ from typing import Optional
from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig 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.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.models.vaes.hunyuan3d import Hunyuan3DVAEConfig
from sglang.multimodal_gen.configs.pipeline_configs.base import ( from sglang.multimodal_gen.configs.pipeline_configs.base import (
ModelTaskType, ModelTaskType,
@@ -30,6 +31,12 @@ class Hunyuan3D2PipelineConfig(PipelineConfig):
vae_config: VAEConfig = field(default_factory=Hunyuan3DVAEConfig) vae_config: VAEConfig = field(default_factory=Hunyuan3DVAEConfig)
vae_precision: str = "fp32" 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 configuration
shape_model_path: Optional[str] = None shape_model_path: Optional[str] = None
shape_use_safetensors: bool = True shape_use_safetensors: bool = True
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py
+6
View File
@@ -29,6 +29,8 @@ DIT_COMPONENT_NAMES = frozenset(
"video_dit_2", "video_dit_2",
"audio_dit", "audio_dit",
"dual_tower_bridge", "dual_tower_bridge",
"delight_transformer",
"paint_transformer",
} }
) )
VAE_COMPONENT_NAMES = frozenset( VAE_COMPONENT_NAMES = frozenset(
@@ -40,6 +42,8 @@ VAE_COMPONENT_NAMES = frozenset(
"spatial_upsampler", "spatial_upsampler",
"condition_image_encoder", "condition_image_encoder",
"diffusion_decoder", "diffusion_decoder",
"delight_vae",
"paint_vae",
} }
) )
DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset( DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
@@ -47,6 +51,8 @@ DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
"vae", "vae",
"video_vae", "video_vae",
"condition_image_encoder", "condition_image_encoder",
"delight_vae",
"paint_vae",
} }
) )
CPU_OFFLOAD_FLAG_NAMES = ( CPU_OFFLOAD_FLAG_NAMES = (
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py
+2 -841
View File
@@ -3,12 +3,11 @@ from __future__ import annotations
import math import math
from dataclasses import dataclass from dataclasses import dataclass
from typing import List, Optional, Tuple from typing import Optional, Tuple
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from einops import rearrange
from sglang.multimodal_gen.configs.models.dits.hunyuan3d import ( from sglang.multimodal_gen.configs.models.dits.hunyuan3d import (
Hunyuan3DDiTArchConfig, Hunyuan3DDiTArchConfig,
@@ -612,842 +611,4 @@ class Hunyuan3D2DiT(CachableDiT, LayerwiseOffloadableModuleMixin):
return latent return latent
import copy EntryClass = Hunyuan3D2DiT
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]
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d_paint.py
+386
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/stable_diffusion.py
+918
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py
+2 -2
View File
@@ -21,7 +21,7 @@ from diffusers.models.autoencoders.vae import (
from diffusers.models.modeling_outputs import AutoencoderKLOutput from diffusers.models.modeling_outputs import AutoencoderKLOutput
from torch import nn 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 ( from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
LayerwiseOffloadableModuleMixin, LayerwiseOffloadableModuleMixin,
) )
@@ -76,7 +76,7 @@ class AutoencoderKL(nn.Module, LayerwiseOffloadableModuleMixin):
def __init__( def __init__(
self, self,
config: FluxVAEConfig, config: VAEConfig,
): ):
super().__init__() super().__init__()
self.config = config self.config = config
python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py
+232 -31
View File
@@ -9,22 +9,45 @@ from __future__ import annotations
import glob import glob
import importlib import importlib
import json
import os import os
from itertools import chain from itertools import chain
from typing import Any from typing import Any
import torch import torch
import torch.nn as nn 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 ( from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
Hunyuan3D2PipelineConfig, Hunyuan3D2PipelineConfig,
) )
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType 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 ( from sglang.multimodal_gen.runtime.loader.fsdp_load import (
load_model_from_full_model_state_dict, load_model_from_full_model_state_dict,
set_default_torch_dtype, set_default_torch_dtype,
) )
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping 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 ( from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase, 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.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -132,11 +156,9 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
"Local path %s not found, downloading from HuggingFace Hub", "Local path %s not found, downloading from HuggingFace Hub",
local_path, local_path,
) )
from huggingface_hub import snapshot_download
downloaded = snapshot_download( downloaded = snapshot_download(
repo_id=model_path, repo_id=model_path,
allow_patterns=[f"{subfolder}/*"], allow_patterns=[f"{subfolder}/**"],
) )
local_path = os.path.join(downloaded, subfolder) local_path = os.path.join(downloaded, subfolder)
@@ -167,8 +189,11 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
return config_path, ckpt_path return config_path, ckpt_path
@staticmethod @staticmethod
def _resolve_paint_dir(model_path: str, subfolder: str) -> str: def _resolve_model_subfolder(
"""Locate (or download) the paint subfolder and return its local path.""" model_path: str,
subfolder: str,
required_files: tuple[str, ...],
) -> str:
local_path = os.path.join(model_path, subfolder) local_path = os.path.join(model_path, subfolder)
if not os.path.exists(local_path): if not os.path.exists(local_path):
local_path = os.path.expanduser(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 %s not found, downloading from HuggingFace Hub",
local_path, local_path,
) )
from huggingface_hub import snapshot_download
downloaded = snapshot_download( downloaded = snapshot_download(
repo_id=model_path, repo_id=model_path,
allow_patterns=[f"{subfolder}/*"], allow_patterns=[f"{subfolder}/**"],
) )
local_path = os.path.join(downloaded, subfolder) local_path = os.path.join(downloaded, subfolder)
for subdir in ("vae", "unet"): for relative_path in required_files:
config_file = os.path.join(local_path, subdir, "config.json") required_file = os.path.join(local_path, relative_path)
if not os.path.exists(config_file): if not os.path.exists(required_file):
raise FileNotFoundError( 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." "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 return local_path
@staticmethod @staticmethod
@@ -203,9 +226,7 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
) -> dict[str, dict[str, torch.Tensor]]: ) -> dict[str, dict[str, torch.Tensor]]:
"""Load a bundled checkpoint and split by the first '.' in each key.""" """Load a bundled checkpoint and split by the first '.' in each key."""
if use_safetensors: if use_safetensors:
import safetensors.torch flat = load_safetensors(ckpt_path, device="cpu")
flat = safetensors.torch.load_file(ckpt_path, device="cpu")
ckpt: dict[str, dict[str, torch.Tensor]] = {} ckpt: dict[str, dict[str, torch.Tensor]] = {}
for key, value in flat.items(): for key, value in flat.items():
component = key.split(".")[0] component = key.split(".")[0]
@@ -291,16 +312,193 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
params = cfg.get("params", {}) params = cfg.get("params", {})
return target_cls(**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 # Module loading override
def load_modules( def load_modules(
self, self,
server_args: ServerArgs, server_args: ServerArgs,
loaded_modules: dict[str, torch.nn.Module] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Load all Hunyuan3D shape components from a bundled checkpoint.""" """Load Hunyuan3D shape and optional texture components."""
import yaml del loaded_modules
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
config = server_args.pipeline_config config = server_args.pipeline_config
if not isinstance(config, Hunyuan3D2PipelineConfig): if not isinstance(config, Hunyuan3D2PipelineConfig):
@@ -348,16 +546,10 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
model_config["image_processor"] model_config["image_processor"]
) )
logger.info("All Hunyuan3D shape components loaded successfully")
if config.paint_enable: if config.paint_enable:
try: components.update(self._load_texture_components(server_args, config))
paint_dir = self._resolve_paint_dir(
server_args.model_path, config.paint_subfolder logger.info("Loaded Hunyuan3D pipeline components: %s", sorted(components))
)
components["hy3dpaint_dir"] = paint_dir
except Exception as e:
logger.warning("Failed to resolve paint model path: %s", e)
return components return components
@@ -411,13 +603,22 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
if config.paint_enable: if config.paint_enable:
self.add_stage( self.add_stage(
stage_name="paint_preprocess", 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( self.add_stage(
stage_name="paint_texgen", stage_name="paint_texgen",
stage=Hunyuan3DPaintTexGenStage( stage=Hunyuan3DPaintTexGenStage(
config=config, 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( self.add_stage(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py
+637 -893
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py
+5
View File
@@ -551,6 +551,11 @@ class Hunyuan3DShapeSaveStage(PipelineStage):
"The surface level may be outside the volume data range." "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) obj_path, return_path = self._get_output_paths(batch)
output_dir = os.path.dirname(obj_path) output_dir = os.path.dirname(obj_path)
if output_dir: if output_dir:
python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py
+210
View File
@@ -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()