# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 from __future__ import annotations import math from dataclasses import dataclass from typing import List, 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, Hunyuan3DDiTConfig, ) from sglang.multimodal_gen.runtime.distributed import divide from sglang.multimodal_gen.runtime.distributed.parallel_state import get_tp_world_size from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, ScaleResidualLayerNormScaleShift, apply_qk_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def _fused_add_gate( residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor ) -> torch.Tensor: return torch.addcmul(residual, x, gate) class MixedRowParallelLinear(RowParallelLinear): """RowParallel for inputs concatenated from multiple separately-sharded sources.""" def __init__(self, input_sizes: list[int], output_size: int, **kwargs): self.input_sizes = input_sizes super().__init__(sum(input_sizes), output_size, **kwargs) def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): input_dim = getattr(param, "input_dim", None) if input_dim is not None: shards = [] offset = 0 for sz in self.input_sizes: part = loaded_weight.narrow(input_dim, offset, sz) per_rank = sz // self.tp_size shard = part.narrow(input_dim, self.tp_rank * per_rank, per_rank) shards.append(shard) offset += sz param.data.copy_(torch.cat(shards, dim=input_dim)) else: param.data.copy_(loaded_weight) def _flux_timestep_embedding( t: torch.Tensor, dim, max_period=10000, time_factor: float = 1000.0 ): """Create sinusoidal timestep embeddings for Flux-style model.""" t = time_factor * t half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half ).to(t.device) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) if torch.is_floating_point(t): embedding = embedding.to(t) return embedding class _FluxGELU(nn.Module): def __init__(self, approximate="tanh"): super().__init__() self.approximate = approximate def forward(self, x: torch.Tensor) -> torch.Tensor: return F.gelu(x, approximate=self.approximate) class _FluxMLPEmbedder(nn.Module): def __init__(self, in_dim: int, hidden_dim: int): super().__init__() self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) self.silu = nn.SiLU() self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.out_layer(self.silu(self.in_layer(x))) class _FluxRMSNorm(nn.Module): def __init__(self, dim: int): super().__init__() self.scale = nn.Parameter(torch.ones(dim)) self.variance_epsilon = 1e-6 self.hidden_size = dim @property def weight(self) -> nn.Parameter: # Keep the original checkpoint key (`scale`) while exposing the # interface expected by the fused QK-norm helper. return self.scale def forward(self, x: torch.Tensor): x_dtype = x.dtype x = x.float() rrms = torch.rsqrt( torch.mean(x**2, dim=-1, keepdim=True) + self.variance_epsilon ) return (x * rrms).to(dtype=x_dtype) * self.scale class _FluxQKNorm(nn.Module): def __init__(self, dim: int): super().__init__() self.dim = dim self.query_norm = _FluxRMSNorm(dim) self.key_norm = _FluxRMSNorm(dim) def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: q, k = apply_qk_norm( q=q.contiguous(), k=k.contiguous(), q_norm=self.query_norm, k_norm=self.key_norm, head_dim=self.dim, allow_inplace=True, ) return q.to(v), k.to(v) class _FluxSelfAttention(nn.Module): def __init__( self, dim: int, num_heads: int = 8, qkv_bias: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, ): super().__init__() tp_size = get_tp_world_size() self.num_heads = num_heads self.local_num_heads = divide(num_heads, tp_size) self.head_dim = dim // num_heads self.qkv = MergedColumnParallelLinear( dim, [dim, dim, dim], bias=qkv_bias, gather_output=False ) self.norm = _FluxQKNorm(self.head_dim) self.proj = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) if supported_attention_backends is None: supported_attention_backends = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } self.local_attn = LocalAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, supported_attention_backends=supported_attention_backends, ) def forward(self, x: torch.Tensor, pe: torch.Tensor) -> torch.Tensor: qkv, _ = self.qkv(x) B, L, _ = qkv.shape qkv = qkv.view(B, L, 3, self.local_num_heads, self.head_dim) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] q = q.transpose(1, 2) k = k.transpose(1, 2) v_for_norm = v.transpose(1, 2) q, k = self.norm(q, k, v_for_norm) q = q.transpose(1, 2) k = k.transpose(1, 2) x = self.local_attn(q, k, v) x = x.flatten(2) x, _ = self.proj(x) return x @dataclass class _FluxModulationOut: shift: torch.Tensor scale: torch.Tensor gate: torch.Tensor class _FluxModulation(nn.Module): def __init__(self, dim: int, double: bool): super().__init__() self.is_double = double self.multiplier = 6 if double else 3 self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) def forward( self, vec: torch.Tensor ) -> Tuple[_FluxModulationOut, Optional[_FluxModulationOut]]: out = self.lin(F.silu(vec))[:, None, :] out = out.chunk(self.multiplier, dim=-1) return ( _FluxModulationOut(*out[:3]), _FluxModulationOut(*out[3:]) if self.is_double else None, ) class _FluxDoubleStreamBlock(nn.Module): def __init__( self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, ): super().__init__() mlp_hidden_dim = int(hidden_size * mlp_ratio) tp_size = get_tp_world_size() self.num_heads = num_heads self.local_num_heads = divide(num_heads, tp_size) self.hidden_size = hidden_size self.head_dim = hidden_size // num_heads self.img_mod = _FluxModulation(hidden_size, double=True) self.img_norm1 = LayerNormScaleShift( hidden_size, elementwise_affine=False, eps=1e-6 ) self.img_attn = _FluxSelfAttention( dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, supported_attention_backends=supported_attention_backends, ) self.img_norm2 = ScaleResidualLayerNormScaleShift( hidden_size, elementwise_affine=False, eps=1e-6 ) self.img_mlp = MLP(hidden_size, mlp_hidden_dim, act_type="gelu_pytorch_tanh") self.txt_mod = _FluxModulation(hidden_size, double=True) self.txt_norm1 = LayerNormScaleShift( hidden_size, elementwise_affine=False, eps=1e-6 ) self.txt_attn = _FluxSelfAttention( dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, supported_attention_backends=supported_attention_backends, ) self.txt_norm2 = ScaleResidualLayerNormScaleShift( hidden_size, elementwise_affine=False, eps=1e-6 ) self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, act_type="gelu_pytorch_tanh") if supported_attention_backends is None: supported_attention_backends = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } self.local_attn_joint = LocalAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, supported_attention_backends=supported_attention_backends, ) def forward( self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: img_mod1, img_mod2 = self.img_mod(vec) txt_mod1, txt_mod2 = self.txt_mod(vec) img_modulated = self.img_norm1(img, shift=img_mod1.shift, scale=img_mod1.scale) B, img_L, _ = img_modulated.shape img_qkv, _ = self.img_attn.qkv(img_modulated) img_qkv = img_qkv.view(B, img_L, 3, self.local_num_heads, self.head_dim) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] img_q_t = img_q.transpose(1, 2) img_k_t = img_k.transpose(1, 2) img_v_t = img_v.transpose(1, 2) img_q_t, img_k_t = self.img_attn.norm(img_q_t, img_k_t, img_v_t) img_q = img_q_t.transpose(1, 2) img_k = img_k_t.transpose(1, 2) txt_modulated = self.txt_norm1(txt, shift=txt_mod1.shift, scale=txt_mod1.scale) txt_L = txt_modulated.shape[1] txt_qkv, _ = self.txt_attn.qkv(txt_modulated) txt_qkv = txt_qkv.view(B, txt_L, 3, self.local_num_heads, self.head_dim) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] txt_q_t = txt_q.transpose(1, 2) txt_k_t = txt_k.transpose(1, 2) txt_v_t = txt_v.transpose(1, 2) txt_q_t, txt_k_t = self.txt_attn.norm(txt_q_t, txt_k_t, txt_v_t) txt_q = txt_q_t.transpose(1, 2) txt_k = txt_k_t.transpose(1, 2) q = torch.cat((txt_q, img_q), dim=1) k = torch.cat((txt_k, img_k), dim=1) v = torch.cat((txt_v, img_v), dim=1) attn = self.local_attn_joint(q, k, v) attn = attn.flatten(2) txt_attn, img_attn = attn[:, :txt_L], attn[:, txt_L:] img_proj, _ = self.img_attn.proj(img_attn) img_modulated, img = self.img_norm2( residual=img, x=img_proj, gate=img_mod1.gate, shift=img_mod2.shift, scale=img_mod2.scale, ) img = _fused_add_gate(img, self.img_mlp(img_modulated), img_mod2.gate) txt_proj, _ = self.txt_attn.proj(txt_attn) txt_modulated, txt = self.txt_norm2( residual=txt, x=txt_proj, gate=txt_mod1.gate, shift=txt_mod2.shift, scale=txt_mod2.scale, ) txt = _fused_add_gate(txt, self.txt_mlp(txt_modulated), txt_mod2.gate) return img, txt class _FluxSingleStreamBlock(nn.Module): """ A DiT block with parallel linear layers as described in https://arxiv.org/abs/2302.05442 and adapted modulation interface. """ def __init__( self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: Optional[float] = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, ): super().__init__() tp_size = get_tp_world_size() self.hidden_dim = hidden_size self.num_heads = num_heads self.local_num_heads = divide(num_heads, tp_size) self.head_dim = hidden_size // num_heads self.tp_size = tp_size self.mlp_hidden_dim = int(hidden_size * mlp_ratio) self.linear1 = MergedColumnParallelLinear( hidden_size, [hidden_size, hidden_size, hidden_size, self.mlp_hidden_dim], bias=True, gather_output=False, ) self.linear2 = MixedRowParallelLinear( [hidden_size, self.mlp_hidden_dim], hidden_size, bias=True, input_is_parallel=True, ) self.norm = _FluxQKNorm(self.head_dim) self.hidden_size = hidden_size self.pre_norm = LayerNormScaleShift( hidden_size, elementwise_affine=False, eps=1e-6 ) self.mlp_act = _FluxGELU(approximate="tanh") self.modulation = _FluxModulation(hidden_size, double=False) if supported_attention_backends is None: supported_attention_backends = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } self.local_attn = LocalAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, supported_attention_backends=supported_attention_backends, ) def forward( self, x: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor ) -> torch.Tensor: mod, _ = self.modulation(vec) x_mod = self.pre_norm(x, shift=mod.shift, scale=mod.scale) linear1_out, _ = self.linear1(x_mod) local_qkv_dim = 3 * self.head_dim * self.local_num_heads local_mlp_dim = self.mlp_hidden_dim // self.tp_size qkv, mlp = torch.split(linear1_out, [local_qkv_dim, local_mlp_dim], dim=-1) B, L, _ = qkv.shape qkv = qkv.view(B, L, 3, self.local_num_heads, self.head_dim) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] q_t = q.transpose(1, 2) k_t = k.transpose(1, 2) v_t = v.transpose(1, 2) q_t, k_t = self.norm(q_t, k_t, v_t) q = q_t.transpose(1, 2) k = k_t.transpose(1, 2) attn = self.local_attn(q, k, v) attn = attn.flatten(2) output, _ = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) return _fused_add_gate(x, output, mod.gate) class _FluxLastLayer(nn.Module): def __init__(self, hidden_size: int, patch_size: int, out_channels: int): super().__init__() self.norm_final = LayerNormScaleShift( hidden_size, elementwise_affine=False, eps=1e-6 ) self.linear = nn.Linear( hidden_size, patch_size * patch_size * out_channels, bias=True ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) ) def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) x = self.norm_final(x, shift=shift[:, None, :], scale=scale[:, None, :]) x = self.linear(x) return x class Hunyuan3D2DiT(CachableDiT, LayerwiseOffloadableModuleMixin): """Hunyuan3D DiT model (Flux-style architecture for Hunyuan3D-2.0).""" _aliases = ["hy3dgen.shapegen.models.Hunyuan3DDiT"] param_names_mapping = Hunyuan3DDiTConfig().param_names_mapping @classmethod def build_config_from_params(cls, params: dict) -> Hunyuan3DDiTConfig: """Build a DiTConfig from YAML-style parameter dict.""" field_mapping = { "num_heads": "num_attention_heads", "depth": "num_layers", "depth_single_blocks": "num_single_layers", } arch_kwargs = {} for k, v in params.items(): if k in ("ckpt_path", "supported_attention_backends"): continue mapped = field_mapping.get(k, k) if k == "axes_dim" and isinstance(v, list): v = tuple(v) arch_kwargs[mapped] = v return Hunyuan3DDiTConfig(arch_config=Hunyuan3DDiTArchConfig(**arch_kwargs)) def __init__( self, config: Hunyuan3DDiTConfig, hf_config: dict | None = None, **kwargs, ): super().__init__(config=config, hf_config=hf_config or {}, **kwargs) arch = config.arch_config in_channels = arch.in_channels context_in_dim = arch.context_in_dim hidden_size = arch.hidden_size mlp_ratio = arch.mlp_ratio num_heads = arch.num_attention_heads depth = arch.num_layers depth_single_blocks = arch.num_single_layers axes_dim = list(arch.axes_dim) theta = arch.theta qkv_bias = arch.qkv_bias time_factor = arch.time_factor guidance_embed = arch.guidance_embed supported_attention_backends = arch._supported_attention_backends self.in_channels = in_channels self.context_in_dim = context_in_dim self.hidden_size = hidden_size self.mlp_ratio = mlp_ratio self.num_heads = num_heads self.num_attention_heads = num_heads self.depth = depth self.depth_single_blocks = depth_single_blocks self.axes_dim = axes_dim self.theta = theta self.qkv_bias = qkv_bias self.time_factor = time_factor self.out_channels = self.in_channels self.num_channels_latents = self.in_channels self.guidance_embed = guidance_embed if hidden_size % num_heads != 0: raise ValueError( f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}" ) pe_dim = hidden_size // num_heads if sum(axes_dim) != pe_dim: raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}") self.latent_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) self.time_in = _FluxMLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) self.cond_in = nn.Linear(context_in_dim, self.hidden_size) self.guidance_in = ( _FluxMLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if guidance_embed else nn.Identity() ) self.double_blocks = nn.ModuleList( [ _FluxDoubleStreamBlock( self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, supported_attention_backends=supported_attention_backends, ) for _ in range(depth) ] ) self.single_blocks = nn.ModuleList( [ _FluxSingleStreamBlock( self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, supported_attention_backends=supported_attention_backends, ) for _ in range(depth_single_blocks) ] ) self.final_layer = _FluxLastLayer(self.hidden_size, 1, self.out_channels) # LayerwiseOffloadableModuleMixin self.layer_names = ["double_blocks", "single_blocks"] def forward( self, x, t, contexts, **kwargs, ) -> torch.Tensor: """Forward pass for denoising.""" cond = contexts["main"] latent = self.latent_in(x) t_emb = _flux_timestep_embedding(t, 256, self.time_factor).to( dtype=latent.dtype ) vec = self.time_in(t_emb) if self.guidance_embed: guidance = kwargs.get("guidance", None) if guidance is None: raise ValueError( "Didn't get guidance strength for guidance distilled model." ) vec = vec + self.guidance_in( _flux_timestep_embedding(guidance, 256, self.time_factor) ) cond = self.cond_in(cond) pe = None # Double blocks for i, block in enumerate(self.double_blocks): latent, cond = block(img=latent, txt=cond, vec=vec, pe=pe) latent = torch.cat((cond, latent), 1) # Single blocks for i, block in enumerate(self.single_blocks): latent = block(latent, vec=vec, pe=pe) latent = latent[:, cond.shape[1] :, ...] latent = self.final_layer(latent, vec) 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]