[diffusion] chore: use native hunyuan3d paint and delight models (#34980)
This commit is contained in:
Mick
@@ -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
|
||||
@@ -4,6 +4,7 @@ from typing import Optional
|
||||
|
||||
from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig
|
||||
from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
|
||||
from sglang.multimodal_gen.configs.models.encoders.clip import CLIPTextConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.base import (
|
||||
ModelTaskType,
|
||||
@@ -30,6 +31,12 @@ class Hunyuan3D2PipelineConfig(PipelineConfig):
|
||||
vae_config: VAEConfig = field(default_factory=Hunyuan3DVAEConfig)
|
||||
vae_precision: str = "fp32"
|
||||
|
||||
text_encoder_configs: tuple[CLIPTextConfig, ...] = field(
|
||||
default_factory=lambda: (CLIPTextConfig(),)
|
||||
)
|
||||
text_encoder_precisions: tuple[str, ...] = ("fp16",)
|
||||
native_only_components = ("delight_text_encoder",)
|
||||
|
||||
# Shape model configuration
|
||||
shape_model_path: Optional[str] = None
|
||||
shape_use_safetensors: bool = True
|
||||
|
||||
+6
@@ -29,6 +29,8 @@ DIT_COMPONENT_NAMES = frozenset(
|
||||
"video_dit_2",
|
||||
"audio_dit",
|
||||
"dual_tower_bridge",
|
||||
"delight_transformer",
|
||||
"paint_transformer",
|
||||
}
|
||||
)
|
||||
VAE_COMPONENT_NAMES = frozenset(
|
||||
@@ -40,6 +42,8 @@ VAE_COMPONENT_NAMES = frozenset(
|
||||
"spatial_upsampler",
|
||||
"condition_image_encoder",
|
||||
"diffusion_decoder",
|
||||
"delight_vae",
|
||||
"paint_vae",
|
||||
}
|
||||
)
|
||||
DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
|
||||
@@ -47,6 +51,8 @@ DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
|
||||
"vae",
|
||||
"video_vae",
|
||||
"condition_image_encoder",
|
||||
"delight_vae",
|
||||
"paint_vae",
|
||||
}
|
||||
)
|
||||
CPU_OFFLOAD_FLAG_NAMES = (
|
||||
|
||||
@@ -3,12 +3,11 @@ from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.hunyuan3d import (
|
||||
Hunyuan3DDiTArchConfig,
|
||||
@@ -612,842 +611,4 @@ class Hunyuan3D2DiT(CachableDiT, LayerwiseOffloadableModuleMixin):
|
||||
return latent
|
||||
|
||||
|
||||
import copy
|
||||
import json
|
||||
import os as _os
|
||||
|
||||
from diffusers.models import UNet2DConditionModel
|
||||
from diffusers.models.attention_processor import Attention as DiffusersAttention
|
||||
from diffusers.models.transformers.transformer_2d import BasicTransformerBlock
|
||||
|
||||
|
||||
def _chunked_feed_forward(
|
||||
ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int
|
||||
):
|
||||
"""Feed forward with chunking to save memory."""
|
||||
if hidden_states.shape[chunk_dim] % chunk_size != 0:
|
||||
raise ValueError(
|
||||
f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]}"
|
||||
f"has to be divisible by chunk size: {chunk_size}."
|
||||
f" Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
|
||||
)
|
||||
|
||||
num_chunks = hidden_states.shape[chunk_dim] // chunk_size
|
||||
ff_output = torch.cat(
|
||||
[ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)],
|
||||
dim=chunk_dim,
|
||||
)
|
||||
return ff_output
|
||||
|
||||
|
||||
class SGLangAttentionWrapper(torch.nn.Module):
|
||||
"""Drop-in replacement for DiffusersAttention that uses sglang's attention backend."""
|
||||
|
||||
_SUPPORTED_BACKENDS = {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
query_dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = False,
|
||||
cross_attention_dim: int | None = None,
|
||||
out_bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.inner_dim = dim_head * heads
|
||||
self.heads = heads
|
||||
self.dim_head = dim_head
|
||||
self.query_dim = query_dim
|
||||
cross_attention_dim = cross_attention_dim or query_dim
|
||||
|
||||
self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
|
||||
self.to_k = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias)
|
||||
self.to_v = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias)
|
||||
self.to_out = nn.ModuleList(
|
||||
[nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)]
|
||||
)
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
wrap_attention_impl_forward,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.selector import (
|
||||
get_attn_backend,
|
||||
)
|
||||
|
||||
attn_backend = get_attn_backend(
|
||||
dim_head, torch.float16, self._SUPPORTED_BACKENDS
|
||||
)
|
||||
impl_cls = attn_backend.get_impl_cls()
|
||||
self.attn_impl = impl_cls(
|
||||
num_heads=heads,
|
||||
head_size=dim_head,
|
||||
softmax_scale=dim_head**-0.5,
|
||||
num_kv_heads=heads,
|
||||
causal=False,
|
||||
)
|
||||
wrap_attention_impl_forward(self.attn_impl)
|
||||
self._attn_backend_name = attn_backend.get_enum().name
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor | None = None,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
|
||||
B, N_q, _ = hidden_states.shape
|
||||
_, N_kv, _ = encoder_hidden_states.shape
|
||||
|
||||
q = self.to_q(hidden_states).view(B, N_q, self.heads, self.dim_head)
|
||||
k = self.to_k(encoder_hidden_states).view(B, N_kv, self.heads, self.dim_head)
|
||||
v = self.to_v(encoder_hidden_states).view(B, N_kv, self.heads, self.dim_head)
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import (
|
||||
get_forward_context,
|
||||
)
|
||||
|
||||
ctx = get_forward_context()
|
||||
out = self.attn_impl.forward(q, k, v, attn_metadata=ctx.attn_metadata)
|
||||
out = out.reshape(B, N_q, self.inner_dim)
|
||||
|
||||
out = self.to_out[0](out)
|
||||
out = self.to_out[1](out)
|
||||
return out
|
||||
|
||||
|
||||
class Basic2p5DTransformerBlock(torch.nn.Module):
|
||||
"""2.5D Transformer block with Multiview Attention (MVA) and Reference View Attention (RVA)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transformer: BasicTransformerBlock,
|
||||
layer_name: str,
|
||||
use_ma: bool = True,
|
||||
use_ra: bool = True,
|
||||
is_turbo: bool = False,
|
||||
use_sglang_attn: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.transformer = transformer
|
||||
self.layer_name = layer_name
|
||||
self.use_ma = use_ma
|
||||
self.use_ra = use_ra
|
||||
self.is_turbo = is_turbo
|
||||
self.use_sglang_attn = use_sglang_attn and not is_turbo
|
||||
|
||||
attn_cls = (
|
||||
SGLangAttentionWrapper if self.use_sglang_attn else DiffusersAttention
|
||||
)
|
||||
attn_kwargs = dict(
|
||||
query_dim=self.dim,
|
||||
heads=self.num_attention_heads,
|
||||
dim_head=self.attention_head_dim,
|
||||
dropout=self.dropout,
|
||||
bias=self.attention_bias,
|
||||
cross_attention_dim=None,
|
||||
upcast_attention=self.attn1.upcast_attention,
|
||||
out_bias=True,
|
||||
)
|
||||
if self.use_sglang_attn:
|
||||
attn_kwargs.pop("upcast_attention")
|
||||
|
||||
if self.use_ma:
|
||||
self.attn_multiview = attn_cls(**attn_kwargs)
|
||||
|
||||
if self.use_ra:
|
||||
self.attn_refview = attn_cls(**attn_kwargs)
|
||||
|
||||
if self.is_turbo:
|
||||
self._initialize_attn_weights()
|
||||
|
||||
def _initialize_attn_weights(self):
|
||||
"""Initialize attention weights for turbo mode."""
|
||||
if self.use_ma:
|
||||
self.attn_multiview.load_state_dict(self.attn1.state_dict())
|
||||
with torch.no_grad():
|
||||
for layer in self.attn_multiview.to_out:
|
||||
for param in layer.parameters():
|
||||
param.zero_()
|
||||
if self.use_ra:
|
||||
self.attn_refview.load_state_dict(self.attn1.state_dict())
|
||||
with torch.no_grad():
|
||||
for layer in self.attn_refview.to_out:
|
||||
for param in layer.parameters():
|
||||
param.zero_()
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self.transformer, name)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: dict = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
added_cond_kwargs: Optional[dict] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Forward pass with MVA and RVA support."""
|
||||
batch_size = hidden_states.shape[0]
|
||||
|
||||
cross_attention_kwargs = (
|
||||
cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
|
||||
)
|
||||
num_in_batch = cross_attention_kwargs.pop("num_in_batch", 1)
|
||||
mode = cross_attention_kwargs.pop("mode", None)
|
||||
|
||||
if not self.is_turbo:
|
||||
mva_scale = cross_attention_kwargs.pop("mva_scale", 1.0)
|
||||
ref_scale = cross_attention_kwargs.pop("ref_scale", 1.0)
|
||||
else:
|
||||
position_attn_mask = cross_attention_kwargs.pop("position_attn_mask", None)
|
||||
position_voxel_indices = cross_attention_kwargs.pop(
|
||||
"position_voxel_indices", None
|
||||
)
|
||||
mva_scale = 1.0
|
||||
ref_scale = 1.0
|
||||
|
||||
condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None)
|
||||
|
||||
# Normalization
|
||||
if self.norm_type == "ada_norm":
|
||||
norm_hidden_states = self.norm1(hidden_states, timestep)
|
||||
elif self.norm_type == "ada_norm_zero":
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
|
||||
hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]:
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
elif self.norm_type == "ada_norm_continuous":
|
||||
norm_hidden_states = self.norm1(
|
||||
hidden_states, added_cond_kwargs["pooled_text_emb"]
|
||||
)
|
||||
elif self.norm_type == "ada_norm_single":
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
|
||||
).chunk(6, dim=1)
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
|
||||
else:
|
||||
raise ValueError("Incorrect norm used")
|
||||
|
||||
if self.pos_embed is not None:
|
||||
norm_hidden_states = self.pos_embed(norm_hidden_states)
|
||||
|
||||
# Prepare GLIGEN inputs
|
||||
cross_attention_kwargs = (
|
||||
cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
|
||||
)
|
||||
gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
|
||||
|
||||
# Self-attention
|
||||
attn_output = self.attn1(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=(
|
||||
encoder_hidden_states if self.only_cross_attention else None
|
||||
),
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
if self.norm_type == "ada_norm_zero":
|
||||
attn_output = gate_msa.unsqueeze(1) * attn_output
|
||||
elif self.norm_type == "ada_norm_single":
|
||||
attn_output = gate_msa * attn_output
|
||||
|
||||
hidden_states = attn_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
# Reference Attention - Write mode
|
||||
if mode is not None and "w" in mode:
|
||||
condition_embed_dict[self.layer_name] = rearrange(
|
||||
norm_hidden_states, "(b n) l c -> b (n l) c", n=num_in_batch
|
||||
)
|
||||
|
||||
# Reference Attention - Read mode
|
||||
if mode is not None and "r" in mode and self.use_ra:
|
||||
condition_embed = (
|
||||
condition_embed_dict[self.layer_name]
|
||||
.unsqueeze(1)
|
||||
.repeat(1, num_in_batch, 1, 1)
|
||||
)
|
||||
condition_embed = rearrange(condition_embed, "b n l c -> (b n) l c")
|
||||
|
||||
attn_output = self.attn_refview(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=condition_embed,
|
||||
attention_mask=None,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
if not self.is_turbo:
|
||||
ref_scale_timing = ref_scale
|
||||
if isinstance(ref_scale, torch.Tensor):
|
||||
ref_scale_timing = (
|
||||
ref_scale.unsqueeze(1).repeat(1, num_in_batch).view(-1)
|
||||
)
|
||||
for _ in range(attn_output.ndim - 1):
|
||||
ref_scale_timing = ref_scale_timing.unsqueeze(-1)
|
||||
|
||||
hidden_states = ref_scale_timing * attn_output + hidden_states
|
||||
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
# Multiview Attention
|
||||
if num_in_batch > 1 and self.use_ma:
|
||||
multivew_hidden_states = rearrange(
|
||||
norm_hidden_states, "(b n) l c -> b (n l) c", n=num_in_batch
|
||||
)
|
||||
|
||||
if self.is_turbo:
|
||||
position_mask = None
|
||||
if position_attn_mask is not None:
|
||||
if multivew_hidden_states.shape[1] in position_attn_mask:
|
||||
position_mask = position_attn_mask[
|
||||
multivew_hidden_states.shape[1]
|
||||
]
|
||||
position_indices = None
|
||||
if position_voxel_indices is not None:
|
||||
if multivew_hidden_states.shape[1] in position_voxel_indices:
|
||||
position_indices = position_voxel_indices[
|
||||
multivew_hidden_states.shape[1]
|
||||
]
|
||||
attn_output = self.attn_multiview(
|
||||
multivew_hidden_states,
|
||||
encoder_hidden_states=multivew_hidden_states,
|
||||
attention_mask=position_mask,
|
||||
position_indices=position_indices,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
attn_output = self.attn_multiview(
|
||||
multivew_hidden_states,
|
||||
encoder_hidden_states=multivew_hidden_states,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
attn_output = rearrange(
|
||||
attn_output, "b (n l) c -> (b n) l c", n=num_in_batch
|
||||
)
|
||||
|
||||
hidden_states = mva_scale * attn_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
# GLIGEN Control
|
||||
if gligen_kwargs is not None:
|
||||
hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
|
||||
|
||||
# Cross-Attention
|
||||
if self.attn2 is not None:
|
||||
if self.norm_type == "ada_norm":
|
||||
norm_hidden_states = self.norm2(hidden_states, timestep)
|
||||
elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]:
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
elif self.norm_type == "ada_norm_single":
|
||||
norm_hidden_states = hidden_states
|
||||
elif self.norm_type == "ada_norm_continuous":
|
||||
norm_hidden_states = self.norm2(
|
||||
hidden_states, added_cond_kwargs["pooled_text_emb"]
|
||||
)
|
||||
else:
|
||||
raise ValueError("Incorrect norm")
|
||||
|
||||
if self.pos_embed is not None and self.norm_type != "ada_norm_single":
|
||||
norm_hidden_states = self.pos_embed(norm_hidden_states)
|
||||
|
||||
attn_output = self.attn2(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=encoder_attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# Feed-forward
|
||||
if self.norm_type == "ada_norm_continuous":
|
||||
norm_hidden_states = self.norm3(
|
||||
hidden_states, added_cond_kwargs["pooled_text_emb"]
|
||||
)
|
||||
elif not self.norm_type == "ada_norm_single":
|
||||
norm_hidden_states = self.norm3(hidden_states)
|
||||
|
||||
if self.norm_type == "ada_norm_zero":
|
||||
norm_hidden_states = (
|
||||
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
)
|
||||
|
||||
if self.norm_type == "ada_norm_single":
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
|
||||
|
||||
if self._chunk_size is not None:
|
||||
ff_output = _chunked_feed_forward(
|
||||
self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size
|
||||
)
|
||||
else:
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
if self.norm_type == "ada_norm_zero":
|
||||
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
||||
elif self.norm_type == "ada_norm_single":
|
||||
ff_output = gate_mlp * ff_output
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_voxel_grid_mask(position: torch.Tensor, grid_resolution: int = 8):
|
||||
"""Compute voxel grid mask for position-aware attention."""
|
||||
position = position.half()
|
||||
B, N, _, H, W = position.shape
|
||||
assert H % grid_resolution == 0 and W % grid_resolution == 0
|
||||
|
||||
valid_mask = (position != 1).all(dim=2, keepdim=True)
|
||||
valid_mask = valid_mask.expand_as(position)
|
||||
position[valid_mask == False] = 0
|
||||
|
||||
position = rearrange(
|
||||
position,
|
||||
"b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w",
|
||||
num_h=grid_resolution,
|
||||
num_w=grid_resolution,
|
||||
)
|
||||
valid_mask = rearrange(
|
||||
valid_mask,
|
||||
"b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w",
|
||||
num_h=grid_resolution,
|
||||
num_w=grid_resolution,
|
||||
)
|
||||
|
||||
grid_position = position.sum(dim=(-2, -1))
|
||||
count_masked = valid_mask.sum(dim=(-2, -1))
|
||||
|
||||
grid_position = grid_position / count_masked.clamp(min=1)
|
||||
grid_position[count_masked < 5] = 0
|
||||
|
||||
grid_position = grid_position.permute(0, 1, 4, 2, 3)
|
||||
grid_position = rearrange(grid_position, "b n c h w -> b n (h w) c")
|
||||
|
||||
grid_position_expanded_1 = grid_position.unsqueeze(2).unsqueeze(4)
|
||||
grid_position_expanded_2 = grid_position.unsqueeze(1).unsqueeze(3)
|
||||
|
||||
distances = torch.norm(grid_position_expanded_1 - grid_position_expanded_2, dim=-1)
|
||||
|
||||
weights = distances
|
||||
grid_distance = 1.73 / grid_resolution
|
||||
|
||||
weights = weights < grid_distance
|
||||
|
||||
return weights
|
||||
|
||||
|
||||
def compute_multi_resolution_mask(
|
||||
position_maps: torch.Tensor, grid_resolutions: List[int] = [32, 16, 8]
|
||||
) -> dict:
|
||||
"""Compute multi-resolution position attention masks."""
|
||||
position_attn_mask = {}
|
||||
with torch.no_grad():
|
||||
for grid_resolution in grid_resolutions:
|
||||
position_mask = compute_voxel_grid_mask(position_maps, grid_resolution)
|
||||
position_mask = rearrange(
|
||||
position_mask, "b ni nj li lj -> b (ni li) (nj lj)"
|
||||
)
|
||||
position_attn_mask[position_mask.shape[1]] = position_mask
|
||||
return position_attn_mask
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_discrete_voxel_indice(
|
||||
position: torch.Tensor, grid_resolution: int = 8, voxel_resolution: int = 128
|
||||
):
|
||||
"""Compute discrete voxel indices for position encoding."""
|
||||
position = position.half()
|
||||
B, N, _, H, W = position.shape
|
||||
assert H % grid_resolution == 0 and W % grid_resolution == 0
|
||||
|
||||
valid_mask = (position != 1).all(dim=2, keepdim=True)
|
||||
valid_mask = valid_mask.expand_as(position)
|
||||
position[valid_mask == False] = 0
|
||||
|
||||
position = rearrange(
|
||||
position,
|
||||
"b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w",
|
||||
num_h=grid_resolution,
|
||||
num_w=grid_resolution,
|
||||
)
|
||||
valid_mask = rearrange(
|
||||
valid_mask,
|
||||
"b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w",
|
||||
num_h=grid_resolution,
|
||||
num_w=grid_resolution,
|
||||
)
|
||||
|
||||
grid_position = position.sum(dim=(-2, -1))
|
||||
count_masked = valid_mask.sum(dim=(-2, -1))
|
||||
|
||||
grid_position = grid_position / count_masked.clamp(min=1)
|
||||
grid_position[count_masked < 5] = 0
|
||||
|
||||
grid_position = grid_position.permute(0, 1, 4, 2, 3).clamp(0, 1)
|
||||
voxel_indices = grid_position * (voxel_resolution - 1)
|
||||
voxel_indices = torch.round(voxel_indices).long()
|
||||
return voxel_indices
|
||||
|
||||
|
||||
def compute_multi_resolution_discrete_voxel_indice(
|
||||
position_maps: torch.Tensor,
|
||||
grid_resolutions: List[int] = [64, 32, 16, 8],
|
||||
voxel_resolutions: List[int] = [512, 256, 128, 64],
|
||||
) -> dict:
|
||||
"""Compute multi-resolution discrete voxel indices."""
|
||||
voxel_indices = {}
|
||||
with torch.no_grad():
|
||||
for grid_resolution, voxel_resolution in zip(
|
||||
grid_resolutions, voxel_resolutions
|
||||
):
|
||||
voxel_indice = compute_discrete_voxel_indice(
|
||||
position_maps, grid_resolution, voxel_resolution
|
||||
)
|
||||
voxel_indice = rearrange(voxel_indice, "b n c h w -> b (n h w) c")
|
||||
voxel_indices[voxel_indice.shape[1]] = {
|
||||
"voxel_indices": voxel_indice,
|
||||
"voxel_resolution": voxel_resolution,
|
||||
}
|
||||
return voxel_indices
|
||||
|
||||
|
||||
class UNet2p5DConditionModel(torch.nn.Module):
|
||||
"""2.5D UNet for multi-view texture generation."""
|
||||
|
||||
def __init__(self, unet: UNet2DConditionModel) -> None:
|
||||
super().__init__()
|
||||
self.unet = unet
|
||||
|
||||
self.use_ma = True
|
||||
self.use_ra = True
|
||||
self.use_camera_embedding = True
|
||||
self.use_dual_stream = True
|
||||
self.is_turbo = False
|
||||
|
||||
if self.use_dual_stream:
|
||||
self.unet_dual = copy.deepcopy(unet)
|
||||
self.init_attention(self.unet_dual)
|
||||
self.init_attention(
|
||||
self.unet, use_ma=self.use_ma, use_ra=self.use_ra, is_turbo=self.is_turbo
|
||||
)
|
||||
self.init_condition()
|
||||
self.init_camera_embedding()
|
||||
|
||||
@staticmethod
|
||||
def from_pretrained(pretrained_model_name_or_path: str, **kwargs):
|
||||
"""Load a pretrained UNet2p5DConditionModel."""
|
||||
torch_dtype = kwargs.pop("dtype", kwargs.pop("torch_dtype", torch.float32))
|
||||
config_path = _os.path.join(pretrained_model_name_or_path, "config.json")
|
||||
unet_ckpt_path = _os.path.join(
|
||||
pretrained_model_name_or_path, "diffusion_pytorch_model.bin"
|
||||
)
|
||||
|
||||
with open(config_path, "r", encoding="utf-8") as file:
|
||||
config = json.load(file)
|
||||
|
||||
unet = UNet2DConditionModel(**config)
|
||||
unet = UNet2p5DConditionModel(unet)
|
||||
unet_ckpt = torch.load(unet_ckpt_path, map_location="cpu", weights_only=True)
|
||||
unet.load_state_dict(unet_ckpt, strict=True)
|
||||
unet = unet.to(torch_dtype)
|
||||
return unet
|
||||
|
||||
def init_condition(self):
|
||||
"""Initialize condition-related modules."""
|
||||
self.unet.conv_in = torch.nn.Conv2d(
|
||||
12, # 4 (latent) + 4 (normal) + 4 (position)
|
||||
self.unet.conv_in.out_channels,
|
||||
kernel_size=self.unet.conv_in.kernel_size,
|
||||
stride=self.unet.conv_in.stride,
|
||||
padding=self.unet.conv_in.padding,
|
||||
dilation=self.unet.conv_in.dilation,
|
||||
groups=self.unet.conv_in.groups,
|
||||
bias=self.unet.conv_in.bias is not None,
|
||||
)
|
||||
|
||||
self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1, 77, 1024))
|
||||
self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1, 77, 1024))
|
||||
|
||||
def init_camera_embedding(self):
|
||||
"""Initialize camera embedding module."""
|
||||
if self.use_camera_embedding:
|
||||
time_embed_dim = 1280
|
||||
self.max_num_ref_image = 5
|
||||
self.max_num_gen_image = 12 * 3 + 4 * 2
|
||||
self.unet.class_embedding = nn.Embedding(
|
||||
self.max_num_ref_image + self.max_num_gen_image, time_embed_dim
|
||||
)
|
||||
|
||||
def init_attention(
|
||||
self,
|
||||
unet: UNet2DConditionModel,
|
||||
use_ma: bool = False,
|
||||
use_ra: bool = False,
|
||||
is_turbo: bool = False,
|
||||
use_sglang_attn: bool = True,
|
||||
):
|
||||
"""Initialize attention blocks with MVA and RVA support."""
|
||||
block_kwargs = dict(
|
||||
use_ma=use_ma,
|
||||
use_ra=use_ra,
|
||||
is_turbo=is_turbo,
|
||||
use_sglang_attn=use_sglang_attn,
|
||||
)
|
||||
|
||||
# Down blocks
|
||||
for down_block_i, down_block in enumerate(unet.down_blocks):
|
||||
if (
|
||||
hasattr(down_block, "has_cross_attention")
|
||||
and down_block.has_cross_attention
|
||||
):
|
||||
for attn_i, attn in enumerate(down_block.attentions):
|
||||
for transformer_i, transformer in enumerate(
|
||||
attn.transformer_blocks
|
||||
):
|
||||
if isinstance(transformer, BasicTransformerBlock):
|
||||
attn.transformer_blocks[transformer_i] = (
|
||||
Basic2p5DTransformerBlock(
|
||||
transformer,
|
||||
f"down_{down_block_i}_{attn_i}_{transformer_i}",
|
||||
**block_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
# Mid block
|
||||
if (
|
||||
hasattr(unet.mid_block, "has_cross_attention")
|
||||
and unet.mid_block.has_cross_attention
|
||||
):
|
||||
for attn_i, attn in enumerate(unet.mid_block.attentions):
|
||||
for transformer_i, transformer in enumerate(attn.transformer_blocks):
|
||||
if isinstance(transformer, BasicTransformerBlock):
|
||||
attn.transformer_blocks[transformer_i] = (
|
||||
Basic2p5DTransformerBlock(
|
||||
transformer,
|
||||
f"mid_{attn_i}_{transformer_i}",
|
||||
**block_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
# Up blocks
|
||||
for up_block_i, up_block in enumerate(unet.up_blocks):
|
||||
if (
|
||||
hasattr(up_block, "has_cross_attention")
|
||||
and up_block.has_cross_attention
|
||||
):
|
||||
for attn_i, attn in enumerate(up_block.attentions):
|
||||
for transformer_i, transformer in enumerate(
|
||||
attn.transformer_blocks
|
||||
):
|
||||
if isinstance(transformer, BasicTransformerBlock):
|
||||
attn.transformer_blocks[transformer_i] = (
|
||||
Basic2p5DTransformerBlock(
|
||||
transformer,
|
||||
f"up_{up_block_i}_{attn_i}_{transformer_i}",
|
||||
**block_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
if use_sglang_attn and (use_ma or use_ra):
|
||||
backend = "unknown"
|
||||
for block in self._iter_2p5d_blocks(unet):
|
||||
for attr in ("attn_multiview", "attn_refview"):
|
||||
wrapper = getattr(block, attr, None)
|
||||
if isinstance(wrapper, SGLangAttentionWrapper):
|
||||
backend = wrapper._attn_backend_name
|
||||
break
|
||||
if backend != "unknown":
|
||||
break
|
||||
count = sum(1 for _ in self._iter_2p5d_blocks(unet))
|
||||
logger.info(
|
||||
"Initialized %d Basic2p5DTransformerBlocks with sglang %s attention",
|
||||
count,
|
||||
backend,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _iter_2p5d_blocks(unet):
|
||||
"""Yield all Basic2p5DTransformerBlock instances in a UNet."""
|
||||
for block_group in (unet.down_blocks, [unet.mid_block], unet.up_blocks):
|
||||
for block in block_group:
|
||||
if not hasattr(block, "attentions"):
|
||||
continue
|
||||
for attn in block.attentions:
|
||||
for tb in attn.transformer_blocks:
|
||||
if isinstance(tb, Basic2p5DTransformerBlock):
|
||||
yield tb
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self.unet, name)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
*args,
|
||||
down_intrablock_additional_residuals=None,
|
||||
down_block_res_samples=None,
|
||||
mid_block_res_sample=None,
|
||||
**cached_condition,
|
||||
):
|
||||
"""Forward pass for multi-view texture generation."""
|
||||
B, N_gen, _, H, W = sample.shape
|
||||
assert H == W
|
||||
|
||||
if self.use_camera_embedding:
|
||||
camera_info_gen = (
|
||||
cached_condition["camera_info_gen"] + self.max_num_ref_image
|
||||
)
|
||||
camera_info_gen = rearrange(camera_info_gen, "b n -> (b n)")
|
||||
else:
|
||||
camera_info_gen = None
|
||||
|
||||
# Concatenate latents with normal and position maps
|
||||
sample = [sample]
|
||||
if "normal_imgs" in cached_condition:
|
||||
sample.append(cached_condition["normal_imgs"])
|
||||
if "position_imgs" in cached_condition:
|
||||
sample.append(cached_condition["position_imgs"])
|
||||
sample = torch.cat(sample, dim=2)
|
||||
|
||||
sample = rearrange(sample, "b n c h w -> (b n) c h w")
|
||||
|
||||
encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(1).repeat(
|
||||
1, N_gen, 1, 1
|
||||
)
|
||||
encoder_hidden_states_gen = rearrange(
|
||||
encoder_hidden_states_gen, "b n l c -> (b n) l c"
|
||||
)
|
||||
|
||||
# Process reference images for RVA
|
||||
if self.use_ra:
|
||||
if "condition_embed_dict" in cached_condition:
|
||||
condition_embed_dict = cached_condition["condition_embed_dict"]
|
||||
else:
|
||||
condition_embed_dict = {}
|
||||
ref_latents = cached_condition["ref_latents"]
|
||||
N_ref = ref_latents.shape[1]
|
||||
|
||||
if self.use_camera_embedding:
|
||||
camera_info_ref = cached_condition["camera_info_ref"]
|
||||
camera_info_ref = rearrange(camera_info_ref, "b n -> (b n)")
|
||||
else:
|
||||
camera_info_ref = None
|
||||
|
||||
ref_latents = rearrange(ref_latents, "b n c h w -> (b n) c h w")
|
||||
|
||||
encoder_hidden_states_ref = self.unet.learned_text_clip_ref.unsqueeze(
|
||||
1
|
||||
).repeat(B, N_ref, 1, 1)
|
||||
encoder_hidden_states_ref = rearrange(
|
||||
encoder_hidden_states_ref, "b n l c -> (b n) l c"
|
||||
)
|
||||
|
||||
noisy_ref_latents = ref_latents
|
||||
timestep_ref = 0
|
||||
|
||||
if self.use_dual_stream:
|
||||
unet_ref = self.unet_dual
|
||||
else:
|
||||
unet_ref = self.unet
|
||||
|
||||
unet_ref(
|
||||
noisy_ref_latents,
|
||||
timestep_ref,
|
||||
encoder_hidden_states=encoder_hidden_states_ref,
|
||||
class_labels=camera_info_ref,
|
||||
return_dict=False,
|
||||
cross_attention_kwargs={
|
||||
"mode": "w",
|
||||
"num_in_batch": N_ref,
|
||||
"condition_embed_dict": condition_embed_dict,
|
||||
},
|
||||
)
|
||||
cached_condition["condition_embed_dict"] = condition_embed_dict
|
||||
else:
|
||||
condition_embed_dict = None
|
||||
|
||||
mva_scale = cached_condition.get("mva_scale", 1.0)
|
||||
ref_scale = cached_condition.get("ref_scale", 1.0)
|
||||
|
||||
if self.is_turbo:
|
||||
position_attn_mask = cached_condition.get("position_attn_mask", None)
|
||||
position_voxel_indices = cached_condition.get(
|
||||
"position_voxel_indices", None
|
||||
)
|
||||
cross_attention_kwargs_ = {
|
||||
"mode": "r",
|
||||
"num_in_batch": N_gen,
|
||||
"condition_embed_dict": condition_embed_dict,
|
||||
"position_attn_mask": position_attn_mask,
|
||||
"position_voxel_indices": position_voxel_indices,
|
||||
"mva_scale": mva_scale,
|
||||
"ref_scale": ref_scale,
|
||||
}
|
||||
else:
|
||||
cross_attention_kwargs_ = {
|
||||
"mode": "r",
|
||||
"num_in_batch": N_gen,
|
||||
"condition_embed_dict": condition_embed_dict,
|
||||
"mva_scale": mva_scale,
|
||||
"ref_scale": ref_scale,
|
||||
}
|
||||
|
||||
return self.unet(
|
||||
sample,
|
||||
timestep,
|
||||
encoder_hidden_states_gen,
|
||||
*args,
|
||||
class_labels=camera_info_gen,
|
||||
down_intrablock_additional_residuals=(
|
||||
[
|
||||
s.to(dtype=self.unet.dtype)
|
||||
for s in down_intrablock_additional_residuals
|
||||
]
|
||||
if down_intrablock_additional_residuals is not None
|
||||
else None
|
||||
),
|
||||
down_block_additional_residuals=(
|
||||
[s.to(dtype=self.unet.dtype) for s in down_block_res_samples]
|
||||
if down_block_res_samples is not None
|
||||
else None
|
||||
),
|
||||
mid_block_additional_residual=(
|
||||
mid_block_res_sample.to(dtype=self.unet.dtype)
|
||||
if mid_block_res_sample is not None
|
||||
else None
|
||||
),
|
||||
return_dict=False,
|
||||
cross_attention_kwargs=cross_attention_kwargs_,
|
||||
)
|
||||
|
||||
|
||||
# Entry class for model registry
|
||||
EntryClass = [Hunyuan3D2DiT, UNet2p5DConditionModel]
|
||||
EntryClass = Hunyuan3D2DiT
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -21,7 +21,7 @@ from diffusers.models.autoencoders.vae import (
|
||||
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
||||
from torch import nn
|
||||
|
||||
from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
@@ -76,7 +76,7 @@ class AutoencoderKL(nn.Module, LayerwiseOffloadableModuleMixin):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: FluxVAEConfig,
|
||||
config: VAEConfig,
|
||||
):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
@@ -9,22 +9,45 @@ from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
from itertools import chain
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import yaml
|
||||
from diffusers import EulerAncestralDiscreteScheduler, LCMScheduler
|
||||
from huggingface_hub import snapshot_download
|
||||
from safetensors.torch import load_file as load_safetensors
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from sglang.multimodal_gen.configs.models.vaes.stable_diffusion import (
|
||||
StableDiffusionVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
|
||||
Hunyuan3D2PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
|
||||
TextEncoderLoader,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.fsdp_load import (
|
||||
load_model_from_full_model_state_dict,
|
||||
set_default_torch_dtype,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
|
||||
from sglang.multimodal_gen.runtime.models.dits.hunyuan3d_paint import (
|
||||
Hunyuan3DPaintUNet,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.stable_diffusion import (
|
||||
StableDiffusionUNet2DConditionModel,
|
||||
StableDiffusionUNetConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.vaes.autoencoder import (
|
||||
AutoencoderKL as StableDiffusionAutoencoderKL,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
@@ -39,6 +62,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.h
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -132,11 +156,9 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
"Local path %s not found, downloading from HuggingFace Hub",
|
||||
local_path,
|
||||
)
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
downloaded = snapshot_download(
|
||||
repo_id=model_path,
|
||||
allow_patterns=[f"{subfolder}/*"],
|
||||
allow_patterns=[f"{subfolder}/**"],
|
||||
)
|
||||
local_path = os.path.join(downloaded, subfolder)
|
||||
|
||||
@@ -167,8 +189,11 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
return config_path, ckpt_path
|
||||
|
||||
@staticmethod
|
||||
def _resolve_paint_dir(model_path: str, subfolder: str) -> str:
|
||||
"""Locate (or download) the paint subfolder and return its local path."""
|
||||
def _resolve_model_subfolder(
|
||||
model_path: str,
|
||||
subfolder: str,
|
||||
required_files: tuple[str, ...],
|
||||
) -> str:
|
||||
local_path = os.path.join(model_path, subfolder)
|
||||
if not os.path.exists(local_path):
|
||||
local_path = os.path.expanduser(local_path)
|
||||
@@ -178,23 +203,21 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
"Local path %s not found, downloading from HuggingFace Hub",
|
||||
local_path,
|
||||
)
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
downloaded = snapshot_download(
|
||||
repo_id=model_path,
|
||||
allow_patterns=[f"{subfolder}/*"],
|
||||
allow_patterns=[f"{subfolder}/**"],
|
||||
)
|
||||
local_path = os.path.join(downloaded, subfolder)
|
||||
|
||||
for subdir in ("vae", "unet"):
|
||||
config_file = os.path.join(local_path, subdir, "config.json")
|
||||
if not os.path.exists(config_file):
|
||||
for relative_path in required_files:
|
||||
required_file = os.path.join(local_path, relative_path)
|
||||
if not os.path.exists(required_file):
|
||||
raise FileNotFoundError(
|
||||
f"Paint model incomplete: {config_file} not found. "
|
||||
f"Hunyuan3D model incomplete: {required_file} not found. "
|
||||
"Download the model or check network connectivity."
|
||||
)
|
||||
|
||||
logger.info("Resolved paint model directory: %s", local_path)
|
||||
logger.debug("Resolved Hunyuan3D model directory: %s", local_path)
|
||||
return local_path
|
||||
|
||||
@staticmethod
|
||||
@@ -203,9 +226,7 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
) -> dict[str, dict[str, torch.Tensor]]:
|
||||
"""Load a bundled checkpoint and split by the first '.' in each key."""
|
||||
if use_safetensors:
|
||||
import safetensors.torch
|
||||
|
||||
flat = safetensors.torch.load_file(ckpt_path, device="cpu")
|
||||
flat = load_safetensors(ckpt_path, device="cpu")
|
||||
ckpt: dict[str, dict[str, torch.Tensor]] = {}
|
||||
for key, value in flat.items():
|
||||
component = key.split(".")[0]
|
||||
@@ -291,16 +312,193 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
params = cfg.get("params", {})
|
||||
return target_cls(**params)
|
||||
|
||||
@staticmethod
|
||||
def _read_json(path: str) -> dict[str, Any]:
|
||||
with open(path, encoding="utf-8") as file:
|
||||
return json.load(file)
|
||||
|
||||
@staticmethod
|
||||
def _load_component_weights(component_dir: str) -> dict[str, torch.Tensor]:
|
||||
safetensors_path = os.path.join(
|
||||
component_dir, "diffusion_pytorch_model.safetensors"
|
||||
)
|
||||
pytorch_path = os.path.join(component_dir, "diffusion_pytorch_model.bin")
|
||||
if os.path.isfile(safetensors_path):
|
||||
return load_safetensors(safetensors_path, device="cpu")
|
||||
if os.path.isfile(pytorch_path):
|
||||
return torch.load(pytorch_path, map_location="cpu", weights_only=True)
|
||||
raise FileNotFoundError(
|
||||
f"No diffusion_pytorch_model weights found in {component_dir}."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _component_device(server_args: ServerArgs, component_name: str) -> torch.device:
|
||||
if server_args.should_cpu_offload_component(component_name):
|
||||
return torch.device("cpu")
|
||||
return get_local_torch_device()
|
||||
|
||||
@staticmethod
|
||||
def _freeze(module: nn.Module) -> nn.Module:
|
||||
for parameter in module.parameters():
|
||||
parameter.requires_grad = False
|
||||
return module.eval()
|
||||
|
||||
@staticmethod
|
||||
def _maybe_compile_texture_transformer(
|
||||
module: nn.Module,
|
||||
server_args: ServerArgs,
|
||||
config: Hunyuan3D2PipelineConfig,
|
||||
) -> None:
|
||||
if not server_args.enable_torch_compile:
|
||||
return
|
||||
compile_mode = (
|
||||
os.environ.get("SGLANG_TORCH_COMPILE_MODE")
|
||||
or config.dit_config.torch_compile_mode
|
||||
)
|
||||
logger.info(
|
||||
"Compiling %s with mode: %s",
|
||||
module.__class__.__name__,
|
||||
compile_mode,
|
||||
)
|
||||
module.compile(mode=compile_mode, fullgraph=False, dynamic=None)
|
||||
|
||||
@classmethod
|
||||
def _load_stable_diffusion_vae(
|
||||
cls,
|
||||
component_dir: str,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> StableDiffusionAutoencoderKL:
|
||||
config_data = cls._read_json(os.path.join(component_dir, "config.json"))
|
||||
vae_config = StableDiffusionVAEConfig()
|
||||
vae_config.update_model_arch(config_data)
|
||||
with set_default_torch_dtype(dtype):
|
||||
vae = StableDiffusionAutoencoderKL(vae_config)
|
||||
weights = cls._load_component_weights(component_dir)
|
||||
vae.load_state_dict(weights, strict=True)
|
||||
vae.to(device=device, dtype=dtype)
|
||||
return cls._freeze(vae)
|
||||
|
||||
@classmethod
|
||||
def _load_stable_diffusion_unet(
|
||||
cls,
|
||||
component_dir: str,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
*,
|
||||
paint: bool,
|
||||
is_turbo: bool = False,
|
||||
) -> nn.Module:
|
||||
config_data = cls._read_json(os.path.join(component_dir, "config.json"))
|
||||
unet_config = StableDiffusionUNetConfig.from_dict(config_data)
|
||||
with set_default_torch_dtype(dtype), torch.device("meta"):
|
||||
unet: nn.Module
|
||||
if paint:
|
||||
unet = Hunyuan3DPaintUNet(unet_config, is_turbo=is_turbo)
|
||||
else:
|
||||
unet = StableDiffusionUNet2DConditionModel(unet_config)
|
||||
weights = cls._load_component_weights(component_dir)
|
||||
unet.load_state_dict(weights, strict=True, assign=True)
|
||||
unet.to(device=device, dtype=dtype)
|
||||
return cls._freeze(unet)
|
||||
|
||||
@classmethod
|
||||
def _load_texture_components(
|
||||
cls,
|
||||
server_args: ServerArgs,
|
||||
config: Hunyuan3D2PipelineConfig,
|
||||
) -> dict[str, Any]:
|
||||
dtype = PRECISION_TO_TYPE[config.dit_precision]
|
||||
components: dict[str, Any] = {}
|
||||
|
||||
paint_dir = cls._resolve_model_subfolder(
|
||||
server_args.model_path,
|
||||
config.paint_subfolder,
|
||||
(
|
||||
"vae/config.json",
|
||||
"unet/config.json",
|
||||
"scheduler/scheduler_config.json",
|
||||
),
|
||||
)
|
||||
components["paint_vae"] = cls._load_stable_diffusion_vae(
|
||||
os.path.join(paint_dir, "vae"),
|
||||
dtype,
|
||||
cls._component_device(server_args, "paint_vae"),
|
||||
)
|
||||
components["paint_transformer"] = cls._load_stable_diffusion_unet(
|
||||
os.path.join(paint_dir, "unet"),
|
||||
dtype,
|
||||
cls._component_device(server_args, "paint_transformer"),
|
||||
paint=True,
|
||||
is_turbo=config.paint_turbo_mode,
|
||||
)
|
||||
cls._maybe_compile_texture_transformer(
|
||||
components["paint_transformer"], server_args, config
|
||||
)
|
||||
paint_scheduler_config = cls._read_json(
|
||||
os.path.join(paint_dir, "scheduler", "scheduler_config.json")
|
||||
)
|
||||
scheduler_class = (
|
||||
LCMScheduler if config.paint_turbo_mode else EulerAncestralDiscreteScheduler
|
||||
)
|
||||
components["paint_scheduler"] = scheduler_class.from_config(
|
||||
paint_scheduler_config,
|
||||
**({} if config.paint_turbo_mode else {"timestep_spacing": "trailing"}),
|
||||
)
|
||||
|
||||
if config.delight_enable:
|
||||
delight_dir = cls._resolve_model_subfolder(
|
||||
server_args.model_path,
|
||||
config.delight_subfolder,
|
||||
(
|
||||
"vae/config.json",
|
||||
"unet/config.json",
|
||||
"scheduler/scheduler_config.json",
|
||||
"text_encoder/config.json",
|
||||
"tokenizer/tokenizer_config.json",
|
||||
),
|
||||
)
|
||||
components["delight_vae"] = cls._load_stable_diffusion_vae(
|
||||
os.path.join(delight_dir, "vae"),
|
||||
dtype,
|
||||
cls._component_device(server_args, "delight_vae"),
|
||||
)
|
||||
components["delight_transformer"] = cls._load_stable_diffusion_unet(
|
||||
os.path.join(delight_dir, "unet"),
|
||||
dtype,
|
||||
cls._component_device(server_args, "delight_transformer"),
|
||||
paint=False,
|
||||
)
|
||||
cls._maybe_compile_texture_transformer(
|
||||
components["delight_transformer"], server_args, config
|
||||
)
|
||||
delight_text_encoder, _ = TextEncoderLoader().load(
|
||||
os.path.join(delight_dir, "text_encoder"),
|
||||
server_args,
|
||||
"delight_text_encoder",
|
||||
"transformers",
|
||||
)
|
||||
components["delight_text_encoder"] = delight_text_encoder
|
||||
components["delight_tokenizer"] = AutoTokenizer.from_pretrained(
|
||||
os.path.join(delight_dir, "tokenizer")
|
||||
)
|
||||
delight_scheduler_config = cls._read_json(
|
||||
os.path.join(delight_dir, "scheduler", "scheduler_config.json")
|
||||
)
|
||||
components["delight_scheduler"] = (
|
||||
EulerAncestralDiscreteScheduler.from_config(delight_scheduler_config)
|
||||
)
|
||||
|
||||
return components
|
||||
|
||||
# Module loading override
|
||||
def load_modules(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
loaded_modules: dict[str, torch.nn.Module] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Load all Hunyuan3D shape components from a bundled checkpoint."""
|
||||
import yaml
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
"""Load Hunyuan3D shape and optional texture components."""
|
||||
del loaded_modules
|
||||
|
||||
config = server_args.pipeline_config
|
||||
if not isinstance(config, Hunyuan3D2PipelineConfig):
|
||||
@@ -348,16 +546,10 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
model_config["image_processor"]
|
||||
)
|
||||
|
||||
logger.info("All Hunyuan3D shape components loaded successfully")
|
||||
|
||||
if config.paint_enable:
|
||||
try:
|
||||
paint_dir = self._resolve_paint_dir(
|
||||
server_args.model_path, config.paint_subfolder
|
||||
)
|
||||
components["hy3dpaint_dir"] = paint_dir
|
||||
except Exception as e:
|
||||
logger.warning("Failed to resolve paint model path: %s", e)
|
||||
components.update(self._load_texture_components(server_args, config))
|
||||
|
||||
logger.info("Loaded Hunyuan3D pipeline components: %s", sorted(components))
|
||||
|
||||
return components
|
||||
|
||||
@@ -411,13 +603,22 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
if config.paint_enable:
|
||||
self.add_stage(
|
||||
stage_name="paint_preprocess",
|
||||
stage=Hunyuan3DPaintPreprocessStage(config=config),
|
||||
stage=Hunyuan3DPaintPreprocessStage(
|
||||
config=config,
|
||||
delight_transformer=self.get_module("delight_transformer"),
|
||||
delight_vae=self.get_module("delight_vae"),
|
||||
delight_text_encoder=self.get_module("delight_text_encoder"),
|
||||
delight_tokenizer=self.get_module("delight_tokenizer"),
|
||||
delight_scheduler=self.get_module("delight_scheduler"),
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="paint_texgen",
|
||||
stage=Hunyuan3DPaintTexGenStage(
|
||||
config=config,
|
||||
paint_dir=self.get_module("hy3dpaint_dir"),
|
||||
transformer=self.get_module("paint_transformer"),
|
||||
scheduler=self.get_module("paint_scheduler"),
|
||||
vae=self.get_module("paint_vae"),
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
|
||||
+637
-893
File diff suppressed because it is too large
Load Diff
+5
@@ -551,6 +551,11 @@ class Hunyuan3DShapeSaveStage(PipelineStage):
|
||||
"The surface level may be outside the volume data range."
|
||||
)
|
||||
|
||||
if batch.is_warmup:
|
||||
if self.config.paint_enable:
|
||||
return batch
|
||||
return OutputBatch(output_file_paths=[], metrics=batch.metrics)
|
||||
|
||||
obj_path, return_path = self._get_output_paths(batch)
|
||||
output_dir = os.path.dirname(obj_path)
|
||||
if output_dir:
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user