From 57c5c343d7907b634ccb6b2344cef49ef133402b Mon Sep 17 00:00:00 2001 From: Prozac614 Date: Mon, 2 Mar 2026 12:28:05 +0800 Subject: [PATCH] [diffusion] model: support Hunyuan3D-2 (#18170) Co-authored-by: yingluosanqian Co-authored-by: daiweitao Co-authored-by: Mick --- .gitignore | 5 + python/pyproject.toml | 3 + python/pyproject_npu.toml | 5 +- python/pyproject_other.toml | 10 +- python/sglang/cli/utils.py | 14 +- .../configs/models/dits/__init__.py | 4 +- .../configs/models/dits/hunyuan3d.py | 44 + .../configs/models/vaes/__init__.py | 2 + .../configs/models/vaes/hunyuan3d.py | 22 + .../configs/pipeline_configs/__init__.py | 4 + .../configs/pipeline_configs/base.py | 10 +- .../configs/pipeline_configs/hunyuan3d.py | 73 + .../configs/sample/hunyuan3d.py | 29 + .../configs/sample/sampling_params.py | 19 +- .../render/hunyuan3d_rasterizer/__init__.py | 80 + .../hunyuan3d_rasterizer/rasterizer.cpp | 140 ++ .../render/hunyuan3d_rasterizer/rasterizer.h | 50 + .../hunyuan3d_rasterizer/rasterizer_gpu.cu | 130 ++ .../csrc/render/mesh_processor/__init__.py | 62 + .../render/mesh_processor/mesh_processor.cpp | 163 ++ python/sglang/multimodal_gen/registry.py | 85 +- .../entrypoints/diffusion_generator.py | 18 +- .../runtime/entrypoints/http_server.py | 3 +- .../runtime/entrypoints/openai/mesh_api.py | 296 ++++ .../runtime/entrypoints/openai/protocol.py | 36 + .../runtime/entrypoints/openai/storage.py | 2 + .../runtime/entrypoints/openai/stores.py | 1 + .../runtime/managers/gpu_worker.py | 33 +- .../runtime/models/dits/hunyuan3d.py | 1401 +++++++++++++++++ .../runtime/models/encoders/hunyuan3d.py | 264 ++++ .../multimodal_gen/runtime/models/registry.py | 62 +- .../models/schedulers/hunyuan3d_scheduler.py | 371 +++++ .../runtime/models/vaes/hunyuan3d_vae.py | 1224 ++++++++++++++ .../runtime/pipelines/hunyuan3d_pipeline.py | 411 +++++ .../runtime/pipelines_core/stages/__init__.py | 24 + .../pipelines_core/stages/hunyuan3d_paint.py | 1066 +++++++++++++ .../pipelines_core/stages/hunyuan3d_shape.py | 527 +++++++ .../pipelines_core/stages/input_validation.py | 50 +- .../runtime/utils/mesh3d_utils.py | 1114 +++++++++++++ .../test/server/perf_baselines.json | 66 + .../test/server/test_server_common.py | 12 + .../test/server/test_server_utils.py | 191 ++- .../test/server/testcase_configs.py | 21 +- 43 files changed, 8082 insertions(+), 65 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/dits/hunyuan3d.py create mode 100644 python/sglang/multimodal_gen/configs/models/vaes/hunyuan3d.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py create mode 100644 python/sglang/multimodal_gen/configs/sample/hunyuan3d.py create mode 100644 python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py create mode 100644 python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp create mode 100644 python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h create mode 100644 python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu create mode 100644 python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py create mode 100644 python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp create mode 100644 python/sglang/multimodal_gen/runtime/entrypoints/openai/mesh_api.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py create mode 100644 python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py create mode 100644 python/sglang/multimodal_gen/runtime/models/schedulers/hunyuan3d_scheduler.py create mode 100644 python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py create mode 100644 python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py diff --git a/.gitignore b/.gitignore index eda02021a..5553dd83e 100644 --- a/.gitignore +++ b/.gitignore @@ -270,3 +270,8 @@ sgl-kernel/csrc/**/*_musa/ # MUSA core dump files *.mudmp + +# Others +*.glb +*.ply +*.npz diff --git a/python/pyproject.toml b/python/pyproject.toml index 7d6d3dc65..0ec5611bf 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -114,6 +114,9 @@ diffusion = [ "cache-dit==1.2.3", "addict==2.4.0", "av==16.1.0", + "scikit-image==0.25.2", + "trimesh>=4.0.0", + "xatlas", ] tracing = [ diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index c62089f44..78a746143 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -79,7 +79,10 @@ diffusion = [ "opencv-python==4.10.0.84", "remote-pdb", "cache-dit==1.2.1", - "addict" + "addict", + "scikit-image==0.25.2", + "trimesh>=4.0.0", + "xatlas", ] tracing = [ diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index f2d1b22aa..718f7886d 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -98,7 +98,10 @@ diffusion_hip = [ "vsa==0.0.4", "runai_model_streamer>=0.15.5", "cache-dit==1.1.8", - "addict" + "addict", + "scikit-image==0.25.2", + "trimesh>=4.0.0", + "xatlas", ] # For Intel Gaudi(device : hpu) follow the installation guide @@ -128,7 +131,10 @@ diffusion_musa = [ "vsa==0.0.4", "runai_model_streamer>=0.15.5", "cache-dit==1.1.8", - "addict" + "addict", + "scikit-image==0.25.2", + "trimesh>=4.0.0", + "xatlas", ] test = [ diff --git a/python/sglang/cli/utils.py b/python/sglang/cli/utils.py index 3871459ec..e3e4d2032 100644 --- a/python/sglang/cli/utils.py +++ b/python/sglang/cli/utils.py @@ -29,8 +29,20 @@ def get_is_diffusion_model(model_path: str) -> bool: Returns False on any failure (network error, 404, offline mode, etc.) so that the caller falls through to the standard LLM server path. """ + try: + from sglang.multimodal_gen.registry import ( + is_known_non_diffusers_multimodal_model, + ) + except ImportError: + is_known_non_diffusers_multimodal_model = lambda _: False + if os.path.isdir(model_path): - return _is_diffusers_model_dir(model_path) + if _is_diffusers_model_dir(model_path): + return True + return is_known_non_diffusers_multimodal_model(model_path) + + if is_known_non_diffusers_multimodal_model(model_path): + return True try: if envs.SGLANG_USE_MODELSCOPE.get(): diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 1e3057352..2c6e5b0c6 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -1,5 +1,6 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig @@ -7,7 +8,8 @@ from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ "HunyuanVideoConfig", + "WanVideoConfig", + "Hunyuan3DDiTConfig", "MOVAAudioConfig", "MOVAVideoConfig", - "WanVideoConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/dits/hunyuan3d.py b/python/sglang/multimodal_gen/configs/models/dits/hunyuan3d.py new file mode 100644 index 000000000..edcac78ca --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/hunyuan3d.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class Hunyuan3DDiTArchConfig(DiTArchConfig): + """Architecture config for Hunyuan3D DiT (Flux-style for Hunyuan3D-2.0).""" + + param_names_mapping: dict = field( + default_factory=lambda: { + r"(.*)\.img_mlp\.0\.(.*)$": r"\1.img_mlp.fc_in.\2", + r"(.*)\.img_mlp\.2\.(.*)$": r"\1.img_mlp.fc_out.\2", + r"(.*)\.txt_mlp\.0\.(.*)$": r"\1.txt_mlp.fc_in.\2", + r"(.*)\.txt_mlp\.2\.(.*)$": r"\1.txt_mlp.fc_out.\2", + } + ) + + in_channels: int = 64 + hidden_size: int = 1024 + num_attention_heads: int = 16 + num_layers: int = 16 + num_single_layers: int = 32 + mlp_ratio: float = 4.0 + context_in_dim: int = 1536 + axes_dim: tuple[int, ...] = (64,) + theta: int = 10000 + qkv_bias: bool = True + guidance_embed: bool = False + time_factor: float = 1000.0 + + def __post_init__(self) -> None: + if self.num_channels_latents == 0: + self.num_channels_latents = self.in_channels + super().__post_init__() + + +@dataclass +class Hunyuan3DDiTConfig(DiTConfig): + """DiT configuration for Hunyuan3D shape generation (Flux-style).""" + + arch_config: Hunyuan3DDiTArchConfig = field(default_factory=Hunyuan3DDiTArchConfig) + subfolder: str = "hunyuan3d-dit-v2-0" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py index d793a6a9d..8b9dd60fe 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -1,6 +1,7 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig +from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig @@ -8,4 +9,5 @@ __all__ = [ "DacVAEConfig", "HunyuanVAEConfig", "WanVAEConfig", + "Hunyuan3DVAEConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/vaes/hunyuan3d.py b/python/sglang/multimodal_gen/configs/models/vaes/hunyuan3d.py new file mode 100644 index 000000000..fd6d72adb --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/hunyuan3d.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class Hunyuan3DVAEArchConfig(VAEArchConfig): + """Architecture config for Hunyuan3D VAE.""" + + latent_shape: tuple[int, ...] = (1024, 64) + scale_factor: float = 1.0 + + +@dataclass +class Hunyuan3DVAEConfig(VAEConfig): + """VAE configuration for Hunyuan3D.""" + + arch_config: Hunyuan3DVAEArchConfig = field(default_factory=Hunyuan3DVAEArchConfig) + subfolder: str = "hunyuan3d-dit-v2-0" + load_encoder: bool = False + load_decoder: bool = True diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index 534caff44..c7005a33f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -19,6 +19,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( FastHunyuanConfig, HunyuanConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import ( @@ -34,6 +37,7 @@ __all__ = [ "DiffusersGenericPipelineConfig", "HunyuanConfig", "FastHunyuanConfig", + "Hunyuan3D2PipelineConfig", "FluxPipelineConfig", "Flux2PipelineConfig", "Flux2KleinPipelineConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index d0b0146ba..05128134b 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -53,6 +53,7 @@ class ModelTaskType(Enum): T2I = auto() # Text to Image I2I = auto() # Image to Image TI2I = auto() # Image to Image or Text-Image to Image + I2M = auto() # Image to Mesh def is_image_gen(self) -> bool: return ( @@ -62,7 +63,11 @@ class ModelTaskType(Enum): ) def requires_image_input(self) -> bool: - return self == ModelTaskType.I2V or self == ModelTaskType.I2I + return ( + self == ModelTaskType.I2V + or self == ModelTaskType.I2I + or self == ModelTaskType.I2M + ) def accepts_image_input(self) -> bool: return ( @@ -70,9 +75,12 @@ class ModelTaskType(Enum): or self == ModelTaskType.I2I or self == ModelTaskType.TI2I or self == ModelTaskType.TI2V + or self == ModelTaskType.I2M ) def data_type(self) -> DataType: + if self == ModelTaskType.I2M: + return DataType.MESH if self.is_image_gen(): return DataType.IMAGE else: diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py new file mode 100644 index 000000000..07ee0a475 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +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.vaes.hunyuan3d import Hunyuan3DVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + PipelineConfig, +) + + +@dataclass +class Hunyuan3D2PipelineConfig(PipelineConfig): + """Pipeline configuration for Hunyuan3D image-to-mesh generation.""" + + task_type: ModelTaskType = ModelTaskType.I2M + + # Subfolder paths + shape_subfolder: str = "hunyuan3d-dit-v2-0" + paint_subfolder: str = "hunyuan3d-paint-v2-0" + delight_subfolder: str = "hunyuan3d-delight-v2-0" + + # DiT configuration + dit_config: DiTConfig = field(default_factory=Hunyuan3DDiTConfig) + dit_precision: str = "fp16" + + # VAE configuration + vae_config: VAEConfig = field(default_factory=Hunyuan3DVAEConfig) + vae_precision: str = "fp32" + + # Shape model configuration + shape_model_path: Optional[str] = None + shape_use_safetensors: bool = True + shape_variant: Optional[str] = "fp16" + shape_num_inference_steps: int = 50 + guidance_scale: float = 5.0 + shape_box_v: float = 1.01 + shape_octree_resolution: int = 384 + shape_mc_level: float = 0.0 + shape_mc_algo: Optional[str] = "mc" + shape_num_chunks: int = 8000 + shape_output_type: str = "trimesh" + + # Delight model configuration + delight_enable: bool = True + delight_prompt: str = "" + delight_negative_prompt: str = "" + delight_strength: float = 1.0 + delight_num_inference_steps: int = 50 + delight_guidance_scale: float = 1.0 + delight_cfg_image: float = 1.5 + + # Paint model configuration + paint_enable: bool = True + paint_num_inference_steps: int = 30 + paint_guidance_scale: float = 2.0 + paint_resolution: int = 512 + paint_render_size: int = 2048 + paint_texture_size: int = 2048 + paint_use_remesh: bool = True + paint_save_glb: bool = True + paint_turbo_mode: bool = False + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + def prepare_latent_shape(self, batch, batch_size, num_frames): + latent_shape = self.vae_config.arch_config.latent_shape + shape = (batch_size, *latent_shape) + return shape diff --git a/python/sglang/multimodal_gen/configs/sample/hunyuan3d.py b/python/sglang/multimodal_gen/configs/sample/hunyuan3d.py new file mode 100644 index 000000000..843ab99c9 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/hunyuan3d.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling parameters for Hunyuan3D generation.""" + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class Hunyuan3DSamplingParams(SamplingParams): + """Sampling parameters for Hunyuan3D image-to-mesh generation.""" + + negative_prompt: str = "" + + shape_num_inference_steps: int = 50 + guidance_scale: float = 5.0 + + paint_num_inference_steps: int = 30 + paint_guidance_scale: float = 2.0 + + def __post_init__(self): + if self.prompt is None: + self.prompt = "" + + if self.num_inference_steps is None: + self.num_inference_steps = self.shape_num_inference_steps + + self.guidance_scale = max(5.0, min(self.guidance_scale, 6.5)) + super().__post_init__() diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 1b411101d..1b7c6b4c7 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -66,12 +66,14 @@ def _sanitize_filename(name: str, replacement: str = "_", max_length: int = 150) class DataType(Enum): IMAGE = auto() VIDEO = auto() + MESH = auto() def get_default_extension(self) -> str: if self == DataType.IMAGE: return "png" - else: + if self == DataType.VIDEO: return "mp4" + return "glb" @dataclass @@ -170,7 +172,7 @@ class SamplingParams: # add extension if needed if not any( self.output_file_name.endswith(ext) - for ext in [".mp4", ".jpg", ".png", ".webp"] + for ext in [".mp4", ".jpg", ".png", ".webp", ".obj", ".glb"] ): self.output_file_name = ( f"{self.output_file_name}.{self.data_type.get_default_extension()}" @@ -357,6 +359,19 @@ class SamplingParams: if not isinstance(self.prompt, str): raise TypeError(f"`prompt` must be a string, but got {type(self.prompt)}") + if self.guidance_scale is None: + try: + from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, + ) + + if isinstance(pipeline_config, Hunyuan3D2PipelineConfig): + self.guidance_scale = pipeline_config.guidance_scale + else: + self.guidance_scale = 1.0 + except ImportError: + self.guidance_scale = 1.0 + self.data_type = server_args.pipeline_config.task_type.data_type() if self.output_path is None and server_args.output_path is not None: diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py new file mode 100644 index 000000000..0d778479e --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Custom CUDA rasterizer for Hunyuan3D texture generation. + +This module provides JIT-compiled CUDA rasterization for fast mesh rendering. +Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 +""" + +from __future__ import annotations + +import os +from typing import List, Tuple + +import torch + +_abs_path = os.path.dirname(os.path.abspath(__file__)) +_custom_rasterizer_kernel = None + + +def _load_custom_rasterizer(): + """JIT compile and load the custom rasterizer kernel.""" + global _custom_rasterizer_kernel + + if _custom_rasterizer_kernel is not None: + return _custom_rasterizer_kernel + + from torch.utils.cpp_extension import load + + _custom_rasterizer_kernel = load( + name="custom_rasterizer_kernel", + sources=[ + f"{_abs_path}/rasterizer.cpp", + f"{_abs_path}/rasterizer_gpu.cu", + ], + extra_cflags=["-O3"], + extra_cuda_cflags=["-O3", "--use_fast_math"], + verbose=False, + ) + return _custom_rasterizer_kernel + + +def rasterize( + pos: torch.Tensor, + tri: torch.Tensor, + resolution: Tuple[int, int], + clamp_depth: torch.Tensor = None, + use_depth_prior: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Rasterize mesh to get face indices and barycentric coordinates.""" + kernel = _load_custom_rasterizer() + + if clamp_depth is None: + clamp_depth = torch.zeros(0, device=pos.device) + + # pos should be [N, 4], remove batch dim if present + if pos.dim() == 3: + pos = pos[0] + + findices, barycentric = kernel.rasterize_image( + pos, tri, clamp_depth, resolution[1], resolution[0], 1e-6, use_depth_prior + ) + return findices, barycentric + + +def interpolate( + col: torch.Tensor, + findices: torch.Tensor, + barycentric: torch.Tensor, + tri: torch.Tensor, +) -> torch.Tensor: + """Interpolate vertex attributes using barycentric coordinates.""" + # Handle zero indices (background) + f = findices - 1 + (findices == 0) + vcol = col[0, tri.long()[f.long()]] + result = barycentric.view(*barycentric.shape, 1) * vcol + result = torch.sum(result, axis=-2) + return result.view(1, *result.shape) + + +__all__ = ["rasterize", "interpolate"] diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp new file mode 100644 index 000000000..72aa005b3 --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 +// Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT + +#include "rasterizer.h" + +void rasterizeTriangleCPU(int idx, float* vt0, float* vt1, float* vt2, int width, int height, INT64* zbuffer, float* d, float occlusion_truncation) { + float x_min = std::min(vt0[0], std::min(vt1[0],vt2[0])); + float x_max = std::max(vt0[0], std::max(vt1[0],vt2[0])); + float y_min = std::min(vt0[1], std::min(vt1[1],vt2[1])); + float y_max = std::max(vt0[1], std::max(vt1[1],vt2[1])); + + for (int px = x_min; px < x_max + 1; ++px) { + if (px < 0 || px >= width) + continue; + for (int py = y_min; py < y_max + 1; ++py) { + if (py < 0 || py >= height) + continue; + float vt[2] = {px + 0.5f, py + 0.5f}; + float baryCentricCoordinate[3]; + calculateBarycentricCoordinate(vt0, vt1, vt2, vt, baryCentricCoordinate); + if (isBarycentricCoordInBounds(baryCentricCoordinate)) { + int pixel = py * width + px; + if (zbuffer == 0) { + zbuffer[pixel] = (INT64)(idx + 1); + continue; + } + + float depth = baryCentricCoordinate[0] * vt0[2] + baryCentricCoordinate[1] * vt1[2] + baryCentricCoordinate[2] * vt2[2]; + float depth_thres = 0; + if (d) { + depth_thres = d[pixel] * 0.49999f + 0.5f + occlusion_truncation; + } + + int z_quantize = depth * (2<<17); + INT64 token = (INT64)z_quantize * MAXINT + (INT64)(idx + 1); + if (depth < depth_thres) + continue; + zbuffer[pixel] = std::min(zbuffer[pixel], token); + } + } + } +} + +void barycentricFromImgcoordCPU(float* V, int* F, int* findices, INT64* zbuffer, int width, int height, int num_vertices, int num_faces, + float* barycentric_map, int pix) +{ + INT64 f = zbuffer[pix] % MAXINT; + if (f == (MAXINT-1)) { + findices[pix] = 0; + barycentric_map[pix * 3] = 0; + barycentric_map[pix * 3 + 1] = 0; + barycentric_map[pix * 3 + 2] = 0; + return; + } + findices[pix] = f; + f -= 1; + float barycentric[3] = {0, 0, 0}; + if (f >= 0) { + float vt[2] = {float(pix % width) + 0.5f, float(pix / width) + 0.5f}; + float* vt0_ptr = V + (F[f * 3] * 4); + float* vt1_ptr = V + (F[f * 3 + 1] * 4); + float* vt2_ptr = V + (F[f * 3 + 2] * 4); + + float vt0[2] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f}; + float vt1[2] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f}; + float vt2[2] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f}; + + calculateBarycentricCoordinate(vt0, vt1, vt2, vt, barycentric); + + barycentric[0] = barycentric[0] / vt0_ptr[3]; + barycentric[1] = barycentric[1] / vt1_ptr[3]; + barycentric[2] = barycentric[2] / vt2_ptr[3]; + float w = 1.0f / (barycentric[0] + barycentric[1] + barycentric[2]); + barycentric[0] *= w; + barycentric[1] *= w; + barycentric[2] *= w; + } + barycentric_map[pix * 3] = barycentric[0]; + barycentric_map[pix * 3 + 1] = barycentric[1]; + barycentric_map[pix * 3 + 2] = barycentric[2]; +} + +void rasterizeImagecoordsKernelCPU(float* V, int* F, float* d, INT64* zbuffer, float occlusion_trunc, int width, int height, int num_vertices, int num_faces, int f) +{ + float* vt0_ptr = V + (F[f * 3] * 4); + float* vt1_ptr = V + (F[f * 3 + 1] * 4); + float* vt2_ptr = V + (F[f * 3 + 2] * 4); + + float vt0[3] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f, vt0_ptr[2] / vt0_ptr[3] * 0.49999f + 0.5f}; + float vt1[3] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f, vt1_ptr[2] / vt1_ptr[3] * 0.49999f + 0.5f}; + float vt2[3] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f, vt2_ptr[2] / vt2_ptr[3] * 0.49999f + 0.5f}; + + rasterizeTriangleCPU(f, vt0, vt1, vt2, width, height, zbuffer, d, occlusion_trunc); +} + +std::vector rasterize_image_cpu(torch::Tensor V, torch::Tensor F, torch::Tensor D, + int width, int height, float occlusion_truncation, int use_depth_prior) +{ + int num_faces = F.size(0); + int num_vertices = V.size(0); + auto options = torch::TensorOptions().dtype(torch::kInt32).requires_grad(false); + auto INT64_options = torch::TensorOptions().dtype(torch::kInt64).requires_grad(false); + auto findices = torch::zeros({height, width}, options); + INT64 maxint = (INT64)MAXINT * (INT64)MAXINT + (MAXINT - 1); + auto z_min = torch::ones({height, width}, INT64_options) * (int64_t)maxint; + + if (!use_depth_prior) { + for (int i = 0; i < num_faces; ++i) { + rasterizeImagecoordsKernelCPU(V.data_ptr(), F.data_ptr(), 0, + (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces, i); + } + } else { + for (int i = 0; i < num_faces; ++i) + rasterizeImagecoordsKernelCPU(V.data_ptr(), F.data_ptr(), D.data_ptr(), + (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces, i); + } + + auto float_options = torch::TensorOptions().dtype(torch::kFloat32).requires_grad(false); + auto barycentric = torch::zeros({height, width, 3}, float_options); + for (int i = 0; i < width * height; ++i) + barycentricFromImgcoordCPU(V.data_ptr(), F.data_ptr(), + findices.data_ptr(), (INT64*)z_min.data_ptr(), width, height, num_vertices, num_faces, barycentric.data_ptr(), i); + + return {findices, barycentric}; +} + +std::vector rasterize_image(torch::Tensor V, torch::Tensor F, torch::Tensor D, + int width, int height, float occlusion_truncation, int use_depth_prior) +{ + int device_id = V.get_device(); + if (device_id == -1) + return rasterize_image_cpu(V, F, D, width, height, occlusion_truncation, use_depth_prior); + else + return rasterize_image_gpu(V, F, D, width, height, occlusion_truncation, use_depth_prior); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("rasterize_image", &rasterize_image, "Custom image rasterization"); +} diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h new file mode 100644 index 000000000..bb1703cf0 --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h @@ -0,0 +1,50 @@ +// SPDX-License-Identifier: Apache-2.0 +// Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 +// Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT + +#ifndef RASTERIZER_H_ +#define RASTERIZER_H_ + +#include +#include +#include +#include + +#define INT64 unsigned long long +#define MAXINT 2147483647 + +__host__ __device__ inline float calculateSignedArea2(float* a, float* b, float* c) { + return ((c[0] - a[0]) * (b[1] - a[1]) - (b[0] - a[0]) * (c[1] - a[1])); +} + +__host__ __device__ inline void calculateBarycentricCoordinate(float* a, float* b, float* c, float* p, + float* barycentric) +{ + float beta_tri = calculateSignedArea2(a, p, c); + float gamma_tri = calculateSignedArea2(a, b, p); + float area = calculateSignedArea2(a, b, c); + if (area == 0) { + barycentric[0] = -1.0; + barycentric[1] = -1.0; + barycentric[2] = -1.0; + return; + } + float tri_inv = 1.0 / area; + float beta = beta_tri * tri_inv; + float gamma = gamma_tri * tri_inv; + float alpha = 1.0 - beta - gamma; + barycentric[0] = alpha; + barycentric[1] = beta; + barycentric[2] = gamma; +} + +__host__ __device__ inline bool isBarycentricCoordInBounds(float* barycentricCoord) { + return barycentricCoord[0] >= 0.0 && barycentricCoord[0] <= 1.0 && + barycentricCoord[1] >= 0.0 && barycentricCoord[1] <= 1.0 && + barycentricCoord[2] >= 0.0 && barycentricCoord[2] <= 1.0; +} + +std::vector rasterize_image_gpu(torch::Tensor V, torch::Tensor F, torch::Tensor D, + int width, int height, float occlusion_truncation, int use_depth_prior); + +#endif diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu new file mode 100644 index 000000000..f1317270d --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 +// Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT + +#include "rasterizer.h" + +__device__ void rasterizeTriangleGPU(int idx, float* vt0, float* vt1, float* vt2, int width, int height, INT64* zbuffer, float* d, float occlusion_truncation) { + float x_min = std::min(vt0[0], std::min(vt1[0],vt2[0])); + float x_max = std::max(vt0[0], std::max(vt1[0],vt2[0])); + float y_min = std::min(vt0[1], std::min(vt1[1],vt2[1])); + float y_max = std::max(vt0[1], std::max(vt1[1],vt2[1])); + + for (int px = x_min; px < x_max + 1; ++px) { + if (px < 0 || px >= width) + continue; + for (int py = y_min; py < y_max + 1; ++py) { + if (py < 0 || py >= height) + continue; + float vt[2] = {px + 0.5f, py + 0.5f}; + float baryCentricCoordinate[3]; + calculateBarycentricCoordinate(vt0, vt1, vt2, vt, baryCentricCoordinate); + if (isBarycentricCoordInBounds(baryCentricCoordinate)) { + int pixel = py * width + px; + if (zbuffer == 0) { + atomicExch(&zbuffer[pixel], (INT64)(idx + 1)); + continue; + } + float depth = baryCentricCoordinate[0] * vt0[2] + baryCentricCoordinate[1] * vt1[2] + baryCentricCoordinate[2] * vt2[2]; + float depth_thres = 0; + if (d) { + depth_thres = d[pixel] * 0.49999f + 0.5f + occlusion_truncation; + } + + int z_quantize = depth * (2<<17); + INT64 token = (INT64)z_quantize * MAXINT + (INT64)(idx + 1); + if (depth < depth_thres) + continue; + atomicMin(&zbuffer[pixel], token); + } + } + } +} + +__global__ void barycentricFromImgcoordGPU(float* V, int* F, int* findices, INT64* zbuffer, int width, int height, int num_vertices, int num_faces, + float* barycentric_map) +{ + int pix = blockIdx.x * blockDim.x + threadIdx.x; + if (pix >= width * height) + return; + INT64 f = zbuffer[pix] % MAXINT; + if (f == (MAXINT-1)) { + findices[pix] = 0; + barycentric_map[pix * 3] = 0; + barycentric_map[pix * 3 + 1] = 0; + barycentric_map[pix * 3 + 2] = 0; + return; + } + findices[pix] = f; + f -= 1; + float barycentric[3] = {0, 0, 0}; + if (f >= 0) { + float vt[2] = {float(pix % width) + 0.5f, float(pix / width) + 0.5f}; + float* vt0_ptr = V + (F[f * 3] * 4); + float* vt1_ptr = V + (F[f * 3 + 1] * 4); + float* vt2_ptr = V + (F[f * 3 + 2] * 4); + + float vt0[2] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f}; + float vt1[2] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f}; + float vt2[2] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f}; + + calculateBarycentricCoordinate(vt0, vt1, vt2, vt, barycentric); + + barycentric[0] = barycentric[0] / vt0_ptr[3]; + barycentric[1] = barycentric[1] / vt1_ptr[3]; + barycentric[2] = barycentric[2] / vt2_ptr[3]; + float w = 1.0f / (barycentric[0] + barycentric[1] + barycentric[2]); + barycentric[0] *= w; + barycentric[1] *= w; + barycentric[2] *= w; + } + barycentric_map[pix * 3] = barycentric[0]; + barycentric_map[pix * 3 + 1] = barycentric[1]; + barycentric_map[pix * 3 + 2] = barycentric[2]; +} + +__global__ void rasterizeImagecoordsKernelGPU(float* V, int* F, float* d, INT64* zbuffer, float occlusion_trunc, int width, int height, int num_vertices, int num_faces) +{ + int f = blockIdx.x * blockDim.x + threadIdx.x; + if (f >= num_faces) + return; + + float* vt0_ptr = V + (F[f * 3] * 4); + float* vt1_ptr = V + (F[f * 3 + 1] * 4); + float* vt2_ptr = V + (F[f * 3 + 2] * 4); + + float vt0[3] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f, vt0_ptr[2] / vt0_ptr[3] * 0.49999f + 0.5f}; + float vt1[3] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f, vt1_ptr[2] / vt1_ptr[3] * 0.49999f + 0.5f}; + float vt2[3] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f, vt2_ptr[2] / vt2_ptr[3] * 0.49999f + 0.5f}; + + rasterizeTriangleGPU(f, vt0, vt1, vt2, width, height, zbuffer, d, occlusion_trunc); +} + +std::vector rasterize_image_gpu(torch::Tensor V, torch::Tensor F, torch::Tensor D, + int width, int height, float occlusion_truncation, int use_depth_prior) +{ + int device_id = V.get_device(); + cudaSetDevice(device_id); + int num_faces = F.size(0); + int num_vertices = V.size(0); + auto options = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA, device_id).requires_grad(false); + auto INT64_options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA, device_id).requires_grad(false); + auto findices = torch::zeros({height, width}, options); + INT64 maxint = (INT64)MAXINT * (INT64)MAXINT + (MAXINT - 1); + auto z_min = torch::ones({height, width}, INT64_options) * (int64_t)maxint; + + if (!use_depth_prior) { + rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr(), F.data_ptr(), 0, + (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces); + } else { + rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr(), F.data_ptr(), D.data_ptr(), + (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces); + } + + auto float_options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, device_id).requires_grad(false); + auto barycentric = torch::zeros({height, width, 3}, float_options); + barycentricFromImgcoordGPU<<<(width * height + 255)/256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(V.data_ptr(), F.data_ptr(), + findices.data_ptr(), (INT64*)z_min.data_ptr(), width, height, num_vertices, num_faces, barycentric.data_ptr()); + + return {findices, barycentric}; +} diff --git a/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py b/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py new file mode 100644 index 000000000..21fc15b96 --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Mesh processor C++ extension for texture inpainting. + +This module provides JIT-compiled C++ mesh processing for fast texture inpainting. +Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 +""" + +from __future__ import annotations + +import os +from typing import Tuple + +import numpy as np + +_abs_path = os.path.dirname(os.path.abspath(__file__)) +_mesh_processor_kernel = None + + +def _load_mesh_processor(): + """JIT compile and load the mesh processor kernel.""" + global _mesh_processor_kernel + + if _mesh_processor_kernel is not None: + return _mesh_processor_kernel + + from torch.utils.cpp_extension import load + + _mesh_processor_kernel = load( + name="mesh_processor_kernel", + sources=[ + f"{_abs_path}/mesh_processor.cpp", + ], + extra_cflags=["-O3"], + verbose=False, + ) + return _mesh_processor_kernel + + +def meshVerticeInpaint( + texture: np.ndarray, + mask: np.ndarray, + vtx_pos: np.ndarray, + vtx_uv: np.ndarray, + pos_idx: np.ndarray, + uv_idx: np.ndarray, + method: str = "smooth", +) -> Tuple[np.ndarray, np.ndarray]: + """Inpaint texture using mesh vertex connectivity.""" + kernel = _load_mesh_processor() + + texture = np.ascontiguousarray(texture, dtype=np.float32) + mask = np.ascontiguousarray(mask, dtype=np.uint8) + vtx_pos = np.ascontiguousarray(vtx_pos, dtype=np.float32) + vtx_uv = np.ascontiguousarray(vtx_uv, dtype=np.float32) + pos_idx = np.ascontiguousarray(pos_idx, dtype=np.int32) + uv_idx = np.ascontiguousarray(uv_idx, dtype=np.int32) + + return kernel.meshVerticeInpaint(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx, method) + + +__all__ = ["meshVerticeInpaint"] diff --git a/python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp b/python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp new file mode 100644 index 000000000..1ce0d35c2 --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 +// Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; +using namespace std; + +std::pair, + py::array_t> meshVerticeInpaint_smooth(py::array_t texture, +py::array_t mask, + py::array_t vtx_pos, py::array_t vtx_uv, + py::array_t pos_idx, py::array_t uv_idx) { + auto texture_buf = texture.request(); + auto mask_buf = mask.request(); + auto vtx_pos_buf = vtx_pos.request(); + auto vtx_uv_buf = vtx_uv.request(); + auto pos_idx_buf = pos_idx.request(); + auto uv_idx_buf = uv_idx.request(); + + int texture_height = texture_buf.shape[0]; + int texture_width = texture_buf.shape[1]; + int texture_channel = texture_buf.shape[2]; + float* texture_ptr = static_cast(texture_buf.ptr); + uint8_t* mask_ptr = static_cast(mask_buf.ptr); + + int vtx_num = vtx_pos_buf.shape[0]; + float* vtx_pos_ptr = static_cast(vtx_pos_buf.ptr); + float* vtx_uv_ptr = static_cast(vtx_uv_buf.ptr); + int* pos_idx_ptr = static_cast(pos_idx_buf.ptr); + int* uv_idx_ptr = static_cast(uv_idx_buf.ptr); + + vector vtx_mask(vtx_num, 0.0f); + vector> vtx_color(vtx_num, vector(texture_channel, 0.0f)); + vector uncolored_vtxs; + + vector> G(vtx_num); + + for (int i = 0; i < uv_idx_buf.shape[0]; ++i) { + for (int k = 0; k < 3; ++k) { + int vtx_uv_idx = uv_idx_ptr[i * 3 + k]; + int vtx_idx = pos_idx_ptr[i * 3 + k]; + int uv_v = round(vtx_uv_ptr[vtx_uv_idx * 2] * (texture_width - 1)); + int uv_u = round((1.0 - vtx_uv_ptr[vtx_uv_idx * 2 + 1]) * (texture_height - 1)); + + if (mask_ptr[uv_u * texture_width + uv_v] > 0) { + vtx_mask[vtx_idx] = 1.0f; + for (int c = 0; c < texture_channel; ++c) { + vtx_color[vtx_idx][c] = texture_ptr[(uv_u * texture_width + uv_v) * texture_channel + c]; + } + }else{ + uncolored_vtxs.push_back(vtx_idx); + } + + G[pos_idx_ptr[i * 3 + k]].push_back(pos_idx_ptr[i * 3 + (k + 1) % 3]); + } + } + + int smooth_count = 2; + int last_uncolored_vtx_count = 0; + while (smooth_count>0) { + int uncolored_vtx_count = 0; + + for (int vtx_idx : uncolored_vtxs) { + + vector sum_color(texture_channel, 0.0f); + float total_weight = 0.0f; + + array vtx_0 = {vtx_pos_ptr[vtx_idx * 3], +vtx_pos_ptr[vtx_idx * 3 + 1], vtx_pos_ptr[vtx_idx * 3 + 2]}; + for (int connected_idx : G[vtx_idx]) { + if (vtx_mask[connected_idx] > 0) { + array vtx1 = {vtx_pos_ptr[connected_idx * 3], + vtx_pos_ptr[connected_idx * 3 + 1], vtx_pos_ptr[connected_idx * 3 + 2]}; + float dist_weight = 1.0f / max(sqrt(pow(vtx_0[0] - vtx1[0], 2) + pow(vtx_0[1] - vtx1[1], 2) + \ + pow(vtx_0[2] - vtx1[2], 2)), 1E-4); + dist_weight = dist_weight * dist_weight; + for (int c = 0; c < texture_channel; ++c) { + sum_color[c] += vtx_color[connected_idx][c] * dist_weight; + } + total_weight += dist_weight; + } + } + + if (total_weight > 0.0f) { + for (int c = 0; c < texture_channel; ++c) { + vtx_color[vtx_idx][c] = sum_color[c] / total_weight; + } + vtx_mask[vtx_idx] = 1.0f; + } else { + uncolored_vtx_count++; + } + + } + + if(last_uncolored_vtx_count==uncolored_vtx_count){ + smooth_count--; + }else{ + smooth_count++; + } + last_uncolored_vtx_count = uncolored_vtx_count; + } + + py::array_t new_texture(texture_buf.size); + py::array_t new_mask(mask_buf.size); + + auto new_texture_buf = new_texture.request(); + auto new_mask_buf = new_mask.request(); + + float* new_texture_ptr = static_cast(new_texture_buf.ptr); + uint8_t* new_mask_ptr = static_cast(new_mask_buf.ptr); + std::copy(texture_ptr, texture_ptr + texture_buf.size, new_texture_ptr); + std::copy(mask_ptr, mask_ptr + mask_buf.size, new_mask_ptr); + + for (int face_idx = 0; face_idx < uv_idx_buf.shape[0]; ++face_idx) { + for (int k = 0; k < 3; ++k) { + int vtx_uv_idx = uv_idx_ptr[face_idx * 3 + k]; + int vtx_idx = pos_idx_ptr[face_idx * 3 + k]; + + if (vtx_mask[vtx_idx] == 1.0f) { + int uv_v = round(vtx_uv_ptr[vtx_uv_idx * 2] * (texture_width - 1)); + int uv_u = round((1.0 - vtx_uv_ptr[vtx_uv_idx * 2 + 1]) * (texture_height - 1)); + + for (int c = 0; c < texture_channel; ++c) { + new_texture_ptr[(uv_u * texture_width + uv_v) * texture_channel + c] = vtx_color[vtx_idx][c]; + } + new_mask_ptr[uv_u * texture_width + uv_v] = 255; + } + } + } + + new_texture.resize({texture_height, texture_width, 3}); + new_mask.resize({texture_height, texture_width}); + return std::make_pair(new_texture, new_mask); +} + + +std::pair, py::array_t> meshVerticeInpaint(py::array_t texture, + py::array_t mask, + py::array_t vtx_pos, py::array_t vtx_uv, + py::array_t pos_idx, py::array_t uv_idx, const std::string& method = "smooth") { + if (method == "smooth") { + return meshVerticeInpaint_smooth(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx); + } else { + throw std::invalid_argument("Invalid method. Use 'smooth'."); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("meshVerticeInpaint", &meshVerticeInpaint, "Mesh-aware texture inpainting", + py::arg("texture"), py::arg("mask"), + py::arg("vtx_pos"), py::arg("vtx_uv"), + py::arg("pos_idx"), py::arg("uv_idx"), + py::arg("method") = "smooth"); +} diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index f96eed817..9ae6492ef 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -45,6 +45,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import ( from sglang.multimodal_gen.configs.pipeline_configs.glm_image import ( GlmImagePipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import ( MOVA360PConfig, @@ -75,6 +78,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( FastHunyuanSamplingParam, HunyuanSamplingParams, ) +from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams from sglang.multimodal_gen.configs.sample.mova import ( MOVA_360P_SamplingParams, @@ -367,26 +371,36 @@ def get_model_info( # 1. Discover all available pipeline classes and cache them _discover_and_register_pipelines() - # 2. Get pipeline class from model's model_index.json - try: - if os.path.exists(model_path): - config = verify_model_config_and_directory(model_path) - else: - config = maybe_download_model_index(model_path) - except Exception as e: - logger.error(f"Could not read model config for '{model_path}': {e}") - if backend == Backend.AUTO: - logger.info("Falling back to diffusers backend") - return _get_diffusers_model_info(model_path) - return None + # 2. Get pipeline class - check non-diffusers models first + pipeline_class_name = get_non_diffusers_pipeline_name(model_path) + if pipeline_class_name: + # Known non-diffusers model, skip model_index.json download + logger.debug( + f"Using registered pipeline '{pipeline_class_name}' for non-diffusers model '{model_path}'" + ) + else: + # Try to get from model_index.json + try: + if os.path.exists(model_path): + config = verify_model_config_and_directory(model_path) + else: + config = maybe_download_model_index(model_path) + except Exception as e: + logger.error(f"Could not read model config for '{model_path}': {e}") + if backend == Backend.AUTO: + logger.info("Falling back to diffusers backend") + return _get_diffusers_model_info(model_path) + return None - pipeline_class_name = config.get("_class_name") - if not pipeline_class_name: - logger.error(f"'_class_name' not found in model_index.json for '{model_path}'") - if backend == Backend.AUTO: - logger.info("Falling back to diffusers backend") - return _get_diffusers_model_info(model_path) - return None + pipeline_class_name = config.get("_class_name") + if not pipeline_class_name: + logger.error( + f"'_class_name' not found in model_index.json for '{model_path}'" + ) + if backend == Backend.AUTO: + logger.info("Falling back to diffusers backend") + return _get_diffusers_model_info(model_path) + return None pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name) if not pipeline_cls: @@ -453,7 +467,7 @@ def _register_configs(): hf_model_paths=[ "hunyuanvideo-community/HunyuanVideo", ], - model_detectors=[lambda hf_id: "hunyuan" in hf_id.lower()], + model_detectors=[lambda hf_id: "hunyuanvideo" in hf_id.lower()], ) register_configs( sampling_param_cls=FastHunyuanSamplingParam, @@ -658,6 +672,37 @@ def _register_configs(): pipeline_config_cls=GlmImagePipelineConfig, model_detectors=[lambda hf_id: "glm-image" in hf_id.lower()], ) + register_configs( + sampling_param_cls=Hunyuan3DSamplingParams, + pipeline_config_cls=Hunyuan3D2PipelineConfig, + hf_model_paths=[ + "tencent/Hunyuan3D-2", + ], + model_detectors=[lambda hf_id: "hunyuan3d" in hf_id.lower()], + ) _register_configs() + + +# Known non-diffusers multimodal model patterns +# Maps pattern -> pipeline_name for models that don't have model_index.json +_NON_DIFFUSERS_MULTIMODAL_PATTERNS: Dict[str, str] = { + "hunyuan3d": "Hunyuan3D2Pipeline", +} + + +def is_known_non_diffusers_multimodal_model(model_path: str) -> bool: + model_path_lower = model_path.lower() + return any( + pattern in model_path_lower for pattern in _NON_DIFFUSERS_MULTIMODAL_PATTERNS + ) + + +def get_non_diffusers_pipeline_name(model_path: str) -> Optional[str]: + """Get the pipeline name for a known non-diffusers model.""" + model_path_lower = model_path.lower() + for pattern, pipeline_name in _NON_DIFFUSERS_MULTIMODAL_PATTERNS.items(): + if pattern in model_path_lower: + return pipeline_name + return None diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index ced687e90..ace021b20 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -13,7 +13,10 @@ import os import time from typing import Any, List, Union -from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) from sglang.multimodal_gen.runtime.entrypoints.utils import ( GenerationResult, ListLorasReq, @@ -227,6 +230,19 @@ class DiffGenerator: ) continue + if req.data_type == DataType.MESH: + for output_idx, sample in enumerate( + output_batch.output_file_paths + ): + results.append( + GenerationResult( + **common, + prompt_index=output_idx, + output_file_path=sample, + ) + ) + continue + samples_out: list[Any] = [] audios_out: list[Any] = [] frames_out: list[Any] = [] diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index b2ab9623b..46a29a819 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -214,11 +214,12 @@ def create_app(server_args: ServerArgs): app.include_router(health_router) app.include_router(vertex_router) - from sglang.multimodal_gen.runtime.entrypoints.openai import common_api + from sglang.multimodal_gen.runtime.entrypoints.openai import common_api, mesh_api app.include_router(common_api.router) app.include_router(image_api.router) app.include_router(video_api.router) + app.include_router(mesh_api.router) app.include_router(weights_api.router) app.state.server_args = server_args diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/mesh_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/mesh_api.py new file mode 100644 index 000000000..ab0b90468 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/mesh_api.py @@ -0,0 +1,296 @@ +import asyncio +import os +import time +from typing import Any, Dict, List, Optional + +from fastapi import ( + APIRouter, + File, + Form, + HTTPException, + Path, + Query, + Request, + UploadFile, +) +from fastapi.responses import FileResponse + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + SamplingParams, + generate_request_id, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + MeshGenerationsRequest, + MeshListResponse, + MeshResponse, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage +from sglang.multimodal_gen.runtime.entrypoints.openai.stores import MESH_STORE +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + add_common_data_to_response, + merge_image_input_list, + process_generation_batch, + save_image_to_path, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.server_args import get_global_server_args +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) +router = APIRouter(prefix="/v1/meshes", tags=["meshes"]) + + +def _normalize_format(fmt: Optional[str]) -> str: + fmt = (fmt or "glb").lower() + return fmt if fmt in ("glb", "obj") else "glb" + + +def _build_sampling_params_from_request( + request_id: str, req: MeshGenerationsRequest, image_path: Optional[str] = None +) -> SamplingParams: + ext = _normalize_format(req.output_format) + + server_args = get_global_server_args() + sampling_kwargs: Dict[str, Any] = { + "request_id": request_id, + "prompt": req.prompt, + "num_frames": 1, + "image_path": [image_path] if image_path else None, + "save_output": True, + "output_file_name": f"{request_id}.{ext}", + "seed": req.seed, + "generator_device": req.generator_device, + } + if req.num_inference_steps is not None: + sampling_kwargs["num_inference_steps"] = req.num_inference_steps + if req.guidance_scale is not None: + sampling_kwargs["guidance_scale"] = req.guidance_scale + if req.negative_prompt is not None: + sampling_kwargs["negative_prompt"] = req.negative_prompt + + return SamplingParams.from_user_sampling_params_args( + model_path=server_args.model_path, + server_args=server_args, + **sampling_kwargs, + ) + + +def _mesh_job_from_sampling( + request_id: str, req: MeshGenerationsRequest, sampling: SamplingParams +) -> Dict[str, Any]: + return { + "id": request_id, + "object": "mesh", + "model": req.model or "", + "status": "queued", + "progress": 0, + "created_at": int(time.time()), + "format": _normalize_format(req.output_format), + "file_path": os.path.abspath(sampling.output_file_path()), + } + + +async def _dispatch_job_async(job_id: str, batch: Req) -> None: + from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client + + try: + save_file_path_list, result = await process_generation_batch( + async_scheduler_client, batch + ) + save_file_path = save_file_path_list[0] + + file_size = None + if os.path.exists(save_file_path): + file_size = os.path.getsize(save_file_path) + + cloud_url = await cloud_storage.upload_and_cleanup(save_file_path) + + update_fields: Dict[str, Any] = { + "status": "completed", + "progress": 100, + "completed_at": int(time.time()), + "url": cloud_url, + "file_path": save_file_path if not cloud_url else None, + "file_size_bytes": file_size, + } + update_fields = add_common_data_to_response( + update_fields, request_id=job_id, result=result + ) + await MESH_STORE.update_fields(job_id, update_fields) + except Exception as e: + logger.error(f"{e}") + await MESH_STORE.update_fields( + job_id, {"status": "failed", "error": {"message": str(e)}} + ) + + +@router.post("", response_model=MeshResponse) +async def create_mesh( + request: Request, + image: Optional[List[UploadFile]] = File(None), + image_array: Optional[List[UploadFile]] = File(None, alias="image[]"), + url: Optional[List[str]] = Form(None), + url_array: Optional[List[str]] = Form(None, alias="url[]"), + prompt: Optional[str] = Form("generate 3d mesh"), + model: Optional[str] = Form(None), + seed: Optional[int] = Form(None), + generator_device: Optional[str] = Form("cuda"), + guidance_scale: Optional[float] = Form(None), + num_inference_steps: Optional[int] = Form(None), + negative_prompt: Optional[str] = Form(None), + output_format: Optional[str] = Form("glb"), +): + content_type = request.headers.get("content-type", "").lower() + request_id = generate_request_id() + server_args = get_global_server_args() + + input_path = None + + if "multipart/form-data" in content_type: + images = image or image_array + urls = url or url_array + image_list = merge_image_input_list(images, urls) + + if not image_list: + raise HTTPException( + status_code=422, + detail="Field 'image' or 'url' is required for mesh generation", + ) + + uploads_dir = os.path.join("outputs", "uploads") + os.makedirs(uploads_dir, exist_ok=True) + img = image_list[0] + filename = img.filename if hasattr(img, "filename") else "input_image" + try: + input_path = await save_image_to_path( + img, os.path.join(uploads_dir, f"{request_id}_{filename}") + ) + except Exception as e: + raise HTTPException( + status_code=400, detail=f"Failed to process image source: {str(e)}" + ) + + req = MeshGenerationsRequest( + prompt=prompt or "generate 3d mesh", + model=model, + seed=seed, + generator_device=generator_device, + num_inference_steps=num_inference_steps, + negative_prompt=negative_prompt, + output_format=output_format, + **( + {"guidance_scale": guidance_scale} if guidance_scale is not None else {} + ), + ) + else: + try: + body = await request.json() + except Exception: + body = {} + try: + payload: Dict[str, Any] = dict(body or {}) + + if payload.get("input_image"): + img_src = payload.pop("input_image") + uploads_dir = os.path.join("outputs", "uploads") + os.makedirs(uploads_dir, exist_ok=True) + input_path = await save_image_to_path( + img_src, + os.path.join(uploads_dir, f"{request_id}_input_image"), + ) + + req = MeshGenerationsRequest(**payload) + except Exception as e: + raise HTTPException(status_code=400, detail=f"Invalid request body: {e}") + + if not input_path: + raise HTTPException( + status_code=422, + detail="An input image is required for mesh generation", + ) + + sampling_params = _build_sampling_params_from_request(request_id, req, input_path) + job = _mesh_job_from_sampling(request_id, req, sampling_params) + await MESH_STORE.upsert(request_id, job) + + batch = prepare_request( + server_args=server_args, + sampling_params=sampling_params, + ) + + asyncio.create_task(_dispatch_job_async(request_id, batch)) + return MeshResponse(**job) + + +@router.get("", response_model=MeshListResponse) +async def list_meshes( + after: Optional[str] = Query(None), + limit: Optional[int] = Query(None, ge=1, le=100), + order: Optional[str] = Query("desc"), +): + order = (order or "desc").lower() + if order not in ("asc", "desc"): + order = "desc" + jobs = await MESH_STORE.list_values() + + reverse = order != "asc" + jobs.sort(key=lambda j: j.get("created_at", 0), reverse=reverse) + + if after is not None: + try: + idx = next(i for i, j in enumerate(jobs) if j["id"] == after) + jobs = jobs[idx + 1 :] + except StopIteration: + jobs = [] + + if limit is not None: + jobs = jobs[:limit] + items = [MeshResponse(**j) for j in jobs] + return MeshListResponse(data=items) + + +@router.get("/{mesh_id}", response_model=MeshResponse) +async def retrieve_mesh(mesh_id: str = Path(...)): + job = await MESH_STORE.get(mesh_id) + if not job: + raise HTTPException(status_code=404, detail="Mesh not found") + return MeshResponse(**job) + + +@router.delete("/{mesh_id}", response_model=MeshResponse) +async def delete_mesh(mesh_id: str = Path(...)): + job = await MESH_STORE.pop(mesh_id) + if not job: + raise HTTPException(status_code=404, detail="Mesh not found") + job["status"] = "deleted" + return MeshResponse(**job) + + +@router.get("/{mesh_id}/content") +async def download_mesh_content( + mesh_id: str = Path(...), variant: Optional[str] = Query(None) +): + job = await MESH_STORE.get(mesh_id) + if not job: + raise HTTPException(status_code=404, detail="Mesh not found") + + if job.get("url"): + raise HTTPException( + status_code=400, + detail=f"Mesh has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}", + ) + + file_path = job.get("file_path") + if not file_path or not os.path.exists(file_path): + raise HTTPException(status_code=404, detail="Generation is still in-progress") + + ext = os.path.splitext(file_path)[1].lower() + media_type = { + ".glb": "model/gltf-binary", + ".obj": "text/plain", + }.get(ext, "application/octet-stream") + + return FileResponse( + path=file_path, media_type=media_type, filename=os.path.basename(file_path) + ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index 9d0e89904..c959e2f22 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -110,6 +110,42 @@ class VideoRemixRequest(BaseModel): prompt: str +# Mesh API protocol models +class MeshResponse(BaseModel): + id: str + object: str = "mesh" + model: str = "" + status: str = "queued" + progress: int = 0 + created_at: int = Field(default_factory=lambda: int(time.time())) + format: str = "glb" + url: Optional[str] = None + completed_at: Optional[int] = None + expires_at: Optional[int] = None + error: Optional[Dict[str, Any]] = None + file_path: Optional[str] = None + file_size_bytes: Optional[int] = None + peak_memory_mb: Optional[float] = None + inference_time_s: Optional[float] = None + + +class MeshGenerationsRequest(BaseModel): + prompt: str = "generate 3d mesh" + input_image: Optional[str] = None + model: Optional[str] = None + seed: Optional[int] = None + generator_device: Optional[str] = "cuda" + num_inference_steps: Optional[int] = None + guidance_scale: Optional[float] = None + negative_prompt: Optional[str] = None + output_format: Optional[str] = "glb" + + +class MeshListResponse(BaseModel): + data: List[MeshResponse] + object: str = "list" + + @dataclass class BaseReq(ABC): rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py index c52508f86..68f66827c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py @@ -56,6 +56,8 @@ class CloudStorage: ".jpeg": "image/jpeg", ".webp": "image/webp", ".mp4": "video/mp4", + ".glb": "model/gltf-binary", + ".obj": "text/plain", }.get(ext, "application/octet-stream") # Use the client created once in __init__ diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py index 3fa212b07..29622f651 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py @@ -45,3 +45,4 @@ class AsyncDictStore: # [request_id, dict] VIDEO_STORE = AsyncDictStore() IMAGE_STORE = AsyncDictStore() +MESH_STORE = AsyncDictStore() diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index d1e618d9d..ca51f21fb 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -255,22 +255,23 @@ class GPUWorker: # Save output to file and return file path only if requested. Avoid the serialization # and deserialization overhead between scheduler_client and gpu_worker. if req.save_output and req.return_file_paths_only and self.rank == 0: - output_paths = save_outputs( - output_batch.output, - req.data_type, - req.fps, - True, - lambda idx: req.output_file_path(len(output_batch.output), idx), - audio=output_batch.audio, - audio_sample_rate=output_batch.audio_sample_rate, - output_compression=req.output_compression, - enable_frame_interpolation=req.enable_frame_interpolation, - frame_interpolation_exp=req.frame_interpolation_exp, - frame_interpolation_scale=req.frame_interpolation_scale, - frame_interpolation_model_path=req.frame_interpolation_model_path, - ) - output_batch.output_file_paths = output_paths - output_batch.output = None + if output_batch.output is not None: + output_paths = save_outputs( + output_batch.output, + req.data_type, + req.fps, + True, + lambda idx: req.output_file_path(len(output_batch.output), idx), + audio=output_batch.audio, + audio_sample_rate=output_batch.audio_sample_rate, + output_compression=req.output_compression, + enable_frame_interpolation=req.enable_frame_interpolation, + frame_interpolation_exp=req.frame_interpolation_exp, + frame_interpolation_scale=req.frame_interpolation_scale, + frame_interpolation_model_path=req.frame_interpolation_model_path, + ) + output_batch.output_file_paths = output_paths + output_batch.output = None # TODO: extract to avoid duplication if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py new file mode 100644 index 000000000..034b7718b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py @@ -0,0 +1,1401 @@ +# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from sglang.multimodal_gen.configs.models.dits.hunyuan3d import ( + Hunyuan3DDiTArchConfig, + Hunyuan3DDiTConfig, +) +from sglang.multimodal_gen.runtime.distributed import divide +from sglang.multimodal_gen.runtime.distributed.parallel_state import get_tp_world_size +from sglang.multimodal_gen.runtime.layers.attention import LocalAttention +from sglang.multimodal_gen.runtime.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class MixedRowParallelLinear(RowParallelLinear): + """RowParallel for inputs concatenated from multiple separately-sharded sources.""" + + def __init__(self, input_sizes: list[int], output_size: int, **kwargs): + self.input_sizes = input_sizes + super().__init__(sum(input_sizes), output_size, **kwargs) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + input_dim = getattr(param, "input_dim", None) + if input_dim is not None: + shards = [] + offset = 0 + for sz in self.input_sizes: + part = loaded_weight.narrow(input_dim, offset, sz) + per_rank = sz // self.tp_size + shard = part.narrow(input_dim, self.tp_rank * per_rank, per_rank) + shards.append(shard) + offset += sz + param.data.copy_(torch.cat(shards, dim=input_dim)) + else: + param.data.copy_(loaded_weight) + + +def _flux_timestep_embedding( + t: torch.Tensor, dim, max_period=10000, time_factor: float = 1000.0 +): + """Create sinusoidal timestep embeddings for Flux-style model.""" + t = time_factor * t + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(t.device) + + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(t) + return embedding + + +class _FluxGELU(nn.Module): + def __init__(self, approximate="tanh"): + super().__init__() + self.approximate = approximate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate=self.approximate) + + +class _FluxMLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class _FluxRMSNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor): + x_dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=x_dtype) * self.scale + + +class _FluxQKNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = _FluxRMSNorm(dim) + self.key_norm = _FluxRMSNorm(dim) + + def forward( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + q = self.query_norm(q) + k = self.key_norm(k) + return q.to(v), k.to(v) + + +class _FluxSelfAttention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + ): + super().__init__() + tp_size = get_tp_world_size() + self.num_heads = num_heads + self.local_num_heads = divide(num_heads, tp_size) + self.head_dim = dim // num_heads + + self.qkv = MergedColumnParallelLinear( + dim, [dim, dim, dim], bias=qkv_bias, gather_output=False + ) + self.norm = _FluxQKNorm(self.head_dim) + self.proj = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) + + if supported_attention_backends is None: + supported_attention_backends = { + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + } + self.local_attn = LocalAttention( + num_heads=self.local_num_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + + def forward(self, x: torch.Tensor, pe: torch.Tensor) -> torch.Tensor: + qkv, _ = self.qkv(x) + B, L, _ = qkv.shape + qkv = qkv.view(B, L, 3, self.local_num_heads, self.head_dim) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v_for_norm = v.transpose(1, 2) + q, k = self.norm(q, k, v_for_norm) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + x = self.local_attn(q, k, v) + x = x.flatten(2) + x, _ = self.proj(x) + return x + + +@dataclass +class _FluxModulationOut: + shift: torch.Tensor + scale: torch.Tensor + gate: torch.Tensor + + +class _FluxModulation(nn.Module): + def __init__(self, dim: int, double: bool): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) + + def forward( + self, vec: torch.Tensor + ) -> Tuple[_FluxModulationOut, Optional[_FluxModulationOut]]: + out = self.lin(F.silu(vec))[:, None, :] + out = out.chunk(self.multiplier, dim=-1) + + return ( + _FluxModulationOut(*out[:3]), + _FluxModulationOut(*out[3:]) if self.is_double else None, + ) + + +class _FluxDoubleStreamBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float, + qkv_bias: bool = False, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + ): + super().__init__() + mlp_hidden_dim = int(hidden_size * mlp_ratio) + tp_size = get_tp_world_size() + self.num_heads = num_heads + self.local_num_heads = divide(num_heads, tp_size) + self.hidden_size = hidden_size + self.head_dim = hidden_size // num_heads + self.img_mod = _FluxModulation(hidden_size, double=True) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_attn = _FluxSelfAttention( + dim=hidden_size, + num_heads=num_heads, + qkv_bias=qkv_bias, + supported_attention_backends=supported_attention_backends, + ) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_mlp = MLP(hidden_size, mlp_hidden_dim, act_type="gelu_pytorch_tanh") + + self.txt_mod = _FluxModulation(hidden_size, double=True) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_attn = _FluxSelfAttention( + dim=hidden_size, + num_heads=num_heads, + qkv_bias=qkv_bias, + supported_attention_backends=supported_attention_backends, + ) + + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, act_type="gelu_pytorch_tanh") + + if supported_attention_backends is None: + supported_attention_backends = { + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + } + self.local_attn_joint = LocalAttention( + num_heads=self.local_num_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + + def forward( + self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + img_modulated = self.img_norm1(img) + img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift + + B, img_L, _ = img_modulated.shape + img_qkv, _ = self.img_attn.qkv(img_modulated) + img_qkv = img_qkv.view(B, img_L, 3, self.local_num_heads, self.head_dim) + img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] + img_q_t = img_q.transpose(1, 2) + img_k_t = img_k.transpose(1, 2) + img_v_t = img_v.transpose(1, 2) + img_q_t, img_k_t = self.img_attn.norm(img_q_t, img_k_t, img_v_t) + img_q = img_q_t.transpose(1, 2) + img_k = img_k_t.transpose(1, 2) + + txt_modulated = self.txt_norm1(txt) + txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift + txt_L = txt_modulated.shape[1] + txt_qkv, _ = self.txt_attn.qkv(txt_modulated) + txt_qkv = txt_qkv.view(B, txt_L, 3, self.local_num_heads, self.head_dim) + txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] + txt_q_t = txt_q.transpose(1, 2) + txt_k_t = txt_k.transpose(1, 2) + txt_v_t = txt_v.transpose(1, 2) + txt_q_t, txt_k_t = self.txt_attn.norm(txt_q_t, txt_k_t, txt_v_t) + txt_q = txt_q_t.transpose(1, 2) + txt_k = txt_k_t.transpose(1, 2) + + q = torch.cat((txt_q, img_q), dim=1) + k = torch.cat((txt_k, img_k), dim=1) + v = torch.cat((txt_v, img_v), dim=1) + + attn = self.local_attn_joint(q, k, v) + attn = attn.flatten(2) + + txt_attn, img_attn = attn[:, :txt_L], attn[:, txt_L:] + + img_proj, _ = self.img_attn.proj(img_attn) + img = img + img_mod1.gate * img_proj + img = img + img_mod2.gate * self.img_mlp( + (1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift + ) + + txt_proj, _ = self.txt_attn.proj(txt_attn) + txt = txt + txt_mod1.gate * txt_proj + txt = txt + txt_mod2.gate * self.txt_mlp( + (1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift + ) + return img, txt + + +class _FluxSingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + qk_scale: Optional[float] = None, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + ): + super().__init__() + + tp_size = get_tp_world_size() + self.hidden_dim = hidden_size + self.num_heads = num_heads + self.local_num_heads = divide(num_heads, tp_size) + self.head_dim = hidden_size // num_heads + self.tp_size = tp_size + + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.linear1 = MergedColumnParallelLinear( + hidden_size, + [hidden_size, hidden_size, hidden_size, self.mlp_hidden_dim], + bias=True, + gather_output=False, + ) + self.linear2 = MixedRowParallelLinear( + [hidden_size, self.mlp_hidden_dim], + hidden_size, + bias=True, + input_is_parallel=True, + ) + + self.norm = _FluxQKNorm(self.head_dim) + + self.hidden_size = hidden_size + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + self.mlp_act = _FluxGELU(approximate="tanh") + self.modulation = _FluxModulation(hidden_size, double=False) + + if supported_attention_backends is None: + supported_attention_backends = { + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + } + self.local_attn = LocalAttention( + num_heads=self.local_num_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + + def forward( + self, x: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor + ) -> torch.Tensor: + mod, _ = self.modulation(vec) + + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + linear1_out, _ = self.linear1(x_mod) + local_qkv_dim = 3 * self.head_dim * self.local_num_heads + local_mlp_dim = self.mlp_hidden_dim // self.tp_size + qkv, mlp = torch.split(linear1_out, [local_qkv_dim, local_mlp_dim], dim=-1) + + B, L, _ = qkv.shape + qkv = qkv.view(B, L, 3, self.local_num_heads, self.head_dim) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + q_t = q.transpose(1, 2) + k_t = k.transpose(1, 2) + v_t = v.transpose(1, 2) + q_t, k_t = self.norm(q_t, k_t, v_t) + q = q_t.transpose(1, 2) + k = k_t.transpose(1, 2) + + attn = self.local_attn(q, k, v) + attn = attn.flatten(2) + + output, _ = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output + + +class _FluxLastLayer(nn.Module): + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear( + hidden_size, patch_size * patch_size * out_channels, bias=True + ) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) + ) + + def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + x = self.linear(x) + return x + + +class Hunyuan3D2DiT(CachableDiT, OffloadableDiTMixin): + """Hunyuan3D DiT model (Flux-style architecture for Hunyuan3D-2.0).""" + + _aliases = ["hy3dgen.shapegen.models.Hunyuan3DDiT"] + + param_names_mapping = Hunyuan3DDiTConfig().param_names_mapping + + @classmethod + def build_config_from_params(cls, params: dict) -> Hunyuan3DDiTConfig: + """Build a DiTConfig from YAML-style parameter dict.""" + field_mapping = { + "num_heads": "num_attention_heads", + "depth": "num_layers", + "depth_single_blocks": "num_single_layers", + } + arch_kwargs = {} + for k, v in params.items(): + if k in ("ckpt_path", "supported_attention_backends"): + continue + mapped = field_mapping.get(k, k) + if k == "axes_dim" and isinstance(v, list): + v = tuple(v) + arch_kwargs[mapped] = v + return Hunyuan3DDiTConfig(arch_config=Hunyuan3DDiTArchConfig(**arch_kwargs)) + + def __init__( + self, + config: Hunyuan3DDiTConfig, + hf_config: dict | None = None, + **kwargs, + ): + super().__init__(config=config, hf_config=hf_config or {}, **kwargs) + arch = config.arch_config + + in_channels = arch.in_channels + context_in_dim = arch.context_in_dim + hidden_size = arch.hidden_size + mlp_ratio = arch.mlp_ratio + num_heads = arch.num_attention_heads + depth = arch.num_layers + depth_single_blocks = arch.num_single_layers + axes_dim = list(arch.axes_dim) + theta = arch.theta + qkv_bias = arch.qkv_bias + time_factor = arch.time_factor + guidance_embed = arch.guidance_embed + supported_attention_backends = arch._supported_attention_backends + + self.in_channels = in_channels + self.context_in_dim = context_in_dim + self.hidden_size = hidden_size + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.num_attention_heads = num_heads + self.depth = depth + self.depth_single_blocks = depth_single_blocks + self.axes_dim = axes_dim + self.theta = theta + self.qkv_bias = qkv_bias + self.time_factor = time_factor + self.out_channels = self.in_channels + self.num_channels_latents = self.in_channels + self.guidance_embed = guidance_embed + + if hidden_size % num_heads != 0: + raise ValueError( + f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}" + ) + pe_dim = hidden_size // num_heads + if sum(axes_dim) != pe_dim: + raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}") + self.latent_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) + self.time_in = _FluxMLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.cond_in = nn.Linear(context_in_dim, self.hidden_size) + self.guidance_in = ( + _FluxMLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + if guidance_embed + else nn.Identity() + ) + + self.double_blocks = nn.ModuleList( + [ + _FluxDoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + supported_attention_backends=supported_attention_backends, + ) + for _ in range(depth) + ] + ) + + self.single_blocks = nn.ModuleList( + [ + _FluxSingleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=mlp_ratio, + supported_attention_backends=supported_attention_backends, + ) + for _ in range(depth_single_blocks) + ] + ) + + self.final_layer = _FluxLastLayer(self.hidden_size, 1, self.out_channels) + + # OffloadableDiTMixin + self.layer_names = ["double_blocks", "single_blocks"] + + def forward( + self, + x, + t, + contexts, + **kwargs, + ) -> torch.Tensor: + """Forward pass for denoising.""" + + cond = contexts["main"] + + latent = self.latent_in(x) + + t_emb = _flux_timestep_embedding(t, 256, self.time_factor).to( + dtype=latent.dtype + ) + + vec = self.time_in(t_emb) + + if self.guidance_embed: + guidance = kwargs.get("guidance", None) + if guidance is None: + raise ValueError( + "Didn't get guidance strength for guidance distilled model." + ) + vec = vec + self.guidance_in( + _flux_timestep_embedding(guidance, 256, self.time_factor) + ) + + cond = self.cond_in(cond) + + pe = None + + # Double blocks + for i, block in enumerate(self.double_blocks): + latent, cond = block(img=latent, txt=cond, vec=vec, pe=pe) + latent = torch.cat((cond, latent), 1) + + # Single blocks + for i, block in enumerate(self.single_blocks): + latent = block(latent, vec=vec, pe=pe) + + latent = latent[:, cond.shape[1] :, ...] + latent = self.final_layer(latent, vec) + return latent + + +import copy +import json +import os as _os + +from diffusers.models import UNet2DConditionModel +from diffusers.models.attention_processor import Attention as DiffusersAttention +from diffusers.models.transformers.transformer_2d import BasicTransformerBlock + + +def _chunked_feed_forward( + ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int +): + """Feed forward with chunking to save memory.""" + if hidden_states.shape[chunk_dim] % chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]}" + f"has to be divisible by chunk size: {chunk_size}." + f" Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = hidden_states.shape[chunk_dim] // chunk_size + ff_output = torch.cat( + [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)], + dim=chunk_dim, + ) + return ff_output + + +class SGLangAttentionWrapper(torch.nn.Module): + """Drop-in replacement for DiffusersAttention that uses sglang's attention backend.""" + + _SUPPORTED_BACKENDS = {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA} + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + cross_attention_dim: int | None = None, + out_bias: bool = True, + ) -> None: + super().__init__() + self.inner_dim = dim_head * heads + self.heads = heads + self.dim_head = dim_head + self.query_dim = query_dim + cross_attention_dim = cross_attention_dim or query_dim + + self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias) + self.to_v = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias) + self.to_out = nn.ModuleList( + [nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)] + ) + + from sglang.multimodal_gen.runtime.layers.attention.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, + ) + 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] diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py new file mode 100644 index 000000000..448ffaa8b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py @@ -0,0 +1,264 @@ +# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 + +import numpy as np +import torch +import torch.nn as nn +from torchvision import transforms +from transformers import ( + CLIPVisionConfig, + CLIPVisionModelWithProjection, + Dinov2Config, + Dinov2Model, +) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + return np.concatenate([emb_sin, emb_cos], axis=1) + + +class ImageEncoder(nn.Module): + MODEL_CLASS = None + MODEL_CONFIG_CLASS = None + mean = [] + std = [] + + def __init__( + self, + version=None, + config=None, + use_cls_token=True, + image_size=224, + **kwargs, + ): + super().__init__() + + if config is None: + self.model = self.MODEL_CLASS.from_pretrained(version) + else: + self.model = self.MODEL_CLASS(self.MODEL_CONFIG_CLASS.from_dict(config)) + self.model.eval() + self.model.requires_grad_(False) + self.use_cls_token = use_cls_token + self.size = image_size // 14 + self.num_patches = (image_size // 14) ** 2 + if self.use_cls_token: + self.num_patches += 1 + + self.transform = transforms.Compose( + [ + transforms.Resize( + image_size, transforms.InterpolationMode.BILINEAR, antialias=True + ), + transforms.CenterCrop(image_size), + transforms.Normalize( + mean=self.mean, + std=self.std, + ), + ] + ) + + def forward(self, image, mask=None, value_range=(-1, 1), **kwargs): + if value_range is not None: + low, high = value_range + image = (image - low) / (high - low) + + image = image.to(self.model.device, dtype=self.model.dtype) + inputs = self.transform(image) + outputs = self.model(inputs) + + last_hidden_state = outputs.last_hidden_state + if not self.use_cls_token: + last_hidden_state = last_hidden_state[:, 1:, :] + + return last_hidden_state + + def unconditional_embedding(self, batch_size, **kwargs): + device = next(self.model.parameters()).device + dtype = next(self.model.parameters()).dtype + zero = torch.zeros( + batch_size, + self.num_patches, + self.model.config.hidden_size, + device=device, + dtype=dtype, + ) + + return zero + + +class CLIPImageEncoder(ImageEncoder): + MODEL_CLASS = CLIPVisionModelWithProjection + MODEL_CONFIG_CLASS = CLIPVisionConfig + mean = [0.48145466, 0.4578275, 0.40821073] + std = [0.26862954, 0.26130258, 0.27577711] + + +class DinoImageEncoder(ImageEncoder): + MODEL_CLASS = Dinov2Model + MODEL_CONFIG_CLASS = Dinov2Config + mean = [0.485, 0.456, 0.406] + std = [0.229, 0.224, 0.225] + + +class DinoImageEncoderMV(DinoImageEncoder): + _aliases = [ + "hy3dshape.models.conditioner.DinoImageEncoderMV", + ] + + def __init__( + self, + version=None, + config=None, + use_cls_token=True, + image_size=224, + view_num=4, + **kwargs, + ): + super().__init__(version, config, use_cls_token, image_size, **kwargs) + self.view_num = view_num + self.num_patches = self.num_patches + pos = np.arange(self.view_num, dtype=np.float32) + view_embedding = torch.from_numpy( + get_1d_sincos_pos_embed_from_grid(self.model.config.hidden_size, pos) + ).float() + + view_embedding = view_embedding.unsqueeze(1).repeat(1, self.num_patches, 1) + self.view_embed = view_embedding.unsqueeze(0) + + def forward(self, image, mask=None, value_range=(-1, 1), view_idxs=None, **kwargs): + if value_range is not None: + low, high = value_range + image = (image - low) / (high - low) + + image = image.to(self.model.device, dtype=self.model.dtype) + + bs, num_views, c, h, w = image.shape + image = image.view(bs * num_views, c, h, w) + + inputs = self.transform(image) + outputs = self.model(inputs) + + last_hidden_state = outputs.last_hidden_state + last_hidden_state = last_hidden_state.view( + bs, num_views, last_hidden_state.shape[-2], last_hidden_state.shape[-1] + ) + + view_embedding = self.view_embed.to(last_hidden_state.dtype).to( + last_hidden_state.device + ) + if view_idxs is not None: + assert len(view_idxs) == bs + view_embeddings = [] + for i in range(bs): + view_idx = view_idxs[i] + assert num_views == len(view_idx) + view_embeddings.append(self.view_embed[:, view_idx, ...]) + view_embedding = ( + torch.cat(view_embeddings, 0) + .to(last_hidden_state.dtype) + .to(last_hidden_state.device) + ) + + if num_views != self.view_num: + view_embedding = view_embedding[:, :num_views, ...] + last_hidden_state = last_hidden_state + view_embedding + last_hidden_state = last_hidden_state.view( + bs, num_views * last_hidden_state.shape[-2], last_hidden_state.shape[-1] + ) + return last_hidden_state + + def unconditional_embedding(self, batch_size, view_idxs, **kwargs): + device = next(self.model.parameters()).device + dtype = next(self.model.parameters()).dtype + zero = torch.zeros( + batch_size, + self.num_patches * len(view_idxs[0]), + self.model.config.hidden_size, + device=device, + dtype=dtype, + ) + return zero + + +def build_image_encoder(config): + if config["type"] == "CLIPImageEncoder": + return CLIPImageEncoder(**config["kwargs"]) + elif config["type"] == "DinoImageEncoder": + return DinoImageEncoder(**config["kwargs"]) + elif config["type"] == "DinoImageEncoderMV": + return DinoImageEncoderMV(**config["kwargs"]) + else: + raise ValueError(f'Unknown image encoder type: {config["type"]}') + + +class DualImageEncoder(nn.Module): + def __init__( + self, + main_image_encoder, + additional_image_encoder, + ): + super().__init__() + self.main_image_encoder = build_image_encoder(main_image_encoder) + self.additional_image_encoder = build_image_encoder(additional_image_encoder) + + def forward(self, image, mask=None, **kwargs): + outputs = { + "main": self.main_image_encoder(image, mask=mask, **kwargs), + "additional": self.additional_image_encoder(image, mask=mask, **kwargs), + } + return outputs + + def unconditional_embedding(self, batch_size, **kwargs): + outputs = { + "main": self.main_image_encoder.unconditional_embedding( + batch_size, **kwargs + ), + "additional": self.additional_image_encoder.unconditional_embedding( + batch_size, **kwargs + ), + } + return outputs + + +class SingleImageEncoder(nn.Module): + def __init__( + self, + main_image_encoder, + ): + super().__init__() + self.main_image_encoder = build_image_encoder(main_image_encoder) + + def forward(self, image, mask=None, **kwargs): + outputs = { + "main": self.main_image_encoder(image, mask=mask, **kwargs), + } + return outputs + + def unconditional_embedding(self, batch_size, **kwargs): + outputs = { + "main": self.main_image_encoder.unconditional_embedding( + batch_size, **kwargs + ), + } + return outputs + + +# Entry class for model registry +EntryClass = [ + SingleImageEncoder, + DualImageEncoder, + DinoImageEncoder, + DinoImageEncoderMV, +] diff --git a/python/sglang/multimodal_gen/runtime/models/registry.py b/python/sglang/multimodal_gen/runtime/models/registry.py index 5e6367a40..9eb0a0a32 100644 --- a/python/sglang/multimodal_gen/runtime/models/registry.py +++ b/python/sglang/multimodal_gen/runtime/models/registry.py @@ -37,10 +37,27 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = { "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"), } +# Global alias mapping: external_path -> canonical_class_name +_ALIAS_TO_MODEL: dict[str, str] = {} + + +def _parse_aliases_from_ast(value_node: ast.expr) -> list[str]: + """Parse _aliases list from AST node.""" + aliases = [] + if isinstance(value_node, (ast.List, ast.Tuple)): + for elt in value_node.elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): + aliases.append(elt.value) + return aliases + @lru_cache(maxsize=None) def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: - discovered_models = _IMAGE_ENCODER_MODELS + discovered_models = dict(_IMAGE_ENCODER_MODELS) + + # Collect class definitions with their _aliases + class_aliases: dict[str, list[str]] = {} + for component in COMPONENT_DIRS: component_path = os.path.join(MODELS_PATH, component) for filename in os.listdir(component_path): @@ -57,7 +74,25 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: entry_class_node = None first_class_def = None + # Collect all class definitions and their _aliases + file_class_aliases: dict[str, list[str]] = {} for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + if first_class_def is None: + first_class_def = node + # Look for _aliases in the class body + for class_body_node in node.body: + if isinstance(class_body_node, ast.Assign): + for target in class_body_node.targets: + if ( + isinstance(target, ast.Name) + and target.id == "_aliases" + ): + aliases = _parse_aliases_from_ast( + class_body_node.value + ) + if aliases: + file_class_aliases[node.name] = aliases if isinstance(node, ast.Assign): for target in node.targets: if ( @@ -66,8 +101,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: ): entry_class_node = node break - if first_class_def is None and isinstance(node, ast.ClassDef): - first_class_def = node + if entry_class_node and first_class_def: model_cls_name_list = [] value_node = entry_class_node.value @@ -95,10 +129,25 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: mod_relname, model_cls_str, ) + # Collect aliases for this class + if model_cls_str in file_class_aliases: + class_aliases[model_cls_str] = file_class_aliases[ + model_cls_str + ] except Exception as e: logger.warning(f"Could not parse {filepath} to find models: {e}") + # Build alias -> canonical class name mapping + for class_name, aliases in class_aliases.items(): + for alias in aliases: + if alias in _ALIAS_TO_MODEL: + logger.warning( + f"Alias '{alias}' already registered for '{_ALIAS_TO_MODEL[alias]}', " + f"will be overwritten by '{class_name}'" + ) + _ALIAS_TO_MODEL[alias] = class_name + return discovered_models @@ -242,6 +291,13 @@ class _ModelRegistry: def get_supported_archs(self) -> Set[str]: return self.registered_models.keys() + def resolve_by_alias(self, alias: str) -> type[nn.Module] | None: + """Resolve a model class by its alias (external module path).""" + if alias in _ALIAS_TO_MODEL: + canonical_name = _ALIAS_TO_MODEL[alias] + return self._try_load_model_cls(canonical_name) + return None + def register_model( self, model_arch: str, diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/hunyuan3d_scheduler.py b/python/sglang/multimodal_gen/runtime/models/schedulers/hunyuan3d_scheduler.py new file mode 100644 index 000000000..259c6c52c --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/hunyuan3d_scheduler.py @@ -0,0 +1,371 @@ +# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class Hunyuan3DFlowMatchSchedulerOutput(BaseOutput): + """Output class for the scheduler's step function.""" + + prev_sample: torch.FloatTensor + + +class Hunyuan3DFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """Euler discrete scheduler for flow matching.""" + + # External module path aliases for compatibility with Hunyuan3D configs + _aliases = [ + "hy3dgen.shapegen.schedulers.FlowMatchEulerDiscreteScheduler", + "hy3dshape.schedulers.FlowMatchEulerDiscreteScheduler", + ] + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting: bool = False, + ): + timesteps = np.linspace( + 1, num_train_timesteps, num_train_timesteps, dtype=np.float32 + ).copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + if not use_dynamic_shifting: + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + self._step_index = None + self._begin_index = None + + self.sigmas = sigmas.to("cpu") + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self) -> Optional[int]: + """The index counter for current timestep.""" + return self._step_index + + @property + def begin_index(self) -> Optional[int]: + """The index for the first timestep.""" + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + """Set the begin index for the scheduler. + + Args: + begin_index: The begin index for the scheduler. + """ + self._begin_index = begin_index + + def scale_model_input( + self, + sample: torch.FloatTensor, + timestep: Optional[Union[float, torch.FloatTensor]] = None, + ) -> torch.FloatTensor: + """Identity operation for flow matching (no input scaling needed).""" + return sample + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + noise: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """Forward process in flow-matching (add noise to sample).""" + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + if self.begin_index is None: + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) for t in timestep + ] + elif self.step_index is not None: + step_indices = [self.step_index] * timestep.shape[0] + else: + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + sample = sigma * noise + (1.0 - sigma) * sample + return sample + + def _sigma_to_t(self, sigma: float) -> float: + """Convert sigma to timestep.""" + return sigma * self.config.num_train_timesteps + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: + """Apply time shift transformation.""" + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def set_timesteps( + self, + num_inference_steps: int = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[float] = None, + ): + """Set the discrete timesteps for the diffusion chain.""" + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + "Must pass a value for `mu` when `use_dynamic_shifting` is True" + ) + + if sigmas is None: + self.num_inference_steps = num_inference_steps + timesteps = np.linspace( + self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps, + ) + sigmas = timesteps / self.config.num_train_timesteps + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) + + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) + timesteps = sigmas * self.config.num_train_timesteps + + self.timesteps = timesteps.to(device=device) + self.sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) + + self._step_index = None + self._begin_index = None + + def index_for_timestep( + self, timestep: float, schedule_timesteps: Optional[torch.Tensor] = None + ) -> int: + """Find the index for a given timestep.""" + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + pos = 1 if len(indices) > 1 else 0 + return indices[pos].item() + + def _init_step_index(self, timestep: Union[float, torch.Tensor]): + """Initialize step index from timestep.""" + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[Hunyuan3DFlowMatchSchedulerOutput, Tuple]: + """Predict the sample from the previous timestep.""" + if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): + raise ValueError( + "Passing integer indices as timesteps is not supported. " + "Pass one of `scheduler.timesteps` as a timestep." + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues + sample = sample.to(torch.float32) + + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + + prev_sample = sample + (sigma_next - sigma) * model_output + prev_sample = prev_sample.to(model_output.dtype) + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + + return Hunyuan3DFlowMatchSchedulerOutput(prev_sample=prev_sample) + + def __len__(self) -> int: + return self.config.num_train_timesteps + + +@dataclass +class Hunyuan3DConsistencyFlowMatchSchedulerOutput(BaseOutput): + """Output for consistency flow matching scheduler.""" + + prev_sample: torch.FloatTensor + pred_original_sample: torch.FloatTensor + + +class Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """Consistency Flow Matching Euler Discrete Scheduler.""" + + # External module path aliases for compatibility with Hunyuan3D configs + _aliases = [ + "hy3dshape.schedulers.Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler", + ] + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + pcm_timesteps: int = 50, + ): + sigmas = np.linspace(0, 1, num_train_timesteps) + step_ratio = num_train_timesteps // pcm_timesteps + + euler_timesteps = (np.arange(1, pcm_timesteps) * step_ratio).round().astype( + np.int64 + ) - 1 + euler_timesteps = np.asarray([0] + euler_timesteps.tolist()) + + self.euler_timesteps = euler_timesteps + self.sigmas = sigmas[self.euler_timesteps] + self.sigmas = torch.from_numpy(self.sigmas.copy()).to(dtype=torch.float32) + self.timesteps = self.sigmas * num_train_timesteps + self._step_index = None + self._begin_index = None + self.sigmas = self.sigmas.to("cpu") + + @property + def step_index(self) -> Optional[int]: + return self._step_index + + @property + def begin_index(self) -> Optional[int]: + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + self._begin_index = begin_index + + def scale_model_input( + self, + sample: torch.FloatTensor, + timestep: Optional[Union[float, torch.FloatTensor]] = None, + ) -> torch.FloatTensor: + """Identity operation for flow matching (no input scaling needed).""" + return sample + + def _sigma_to_t(self, sigma: float) -> float: + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + ): + """Set timesteps for inference.""" + self.num_inference_steps = ( + num_inference_steps if num_inference_steps is not None else len(sigmas) + ) + inference_indices = np.linspace( + 0, self.config.pcm_timesteps, num=self.num_inference_steps, endpoint=False + ) + inference_indices = np.floor(inference_indices).astype(np.int64) + inference_indices = torch.from_numpy(inference_indices).long() + + self.sigmas_ = self.sigmas[inference_indices] + timesteps = self.sigmas_ * self.config.num_train_timesteps + self.timesteps = timesteps.to(device=device) + self.sigmas_ = torch.cat( + [self.sigmas_, torch.ones(1, device=self.sigmas_.device)] + ) + + self._step_index = None + self._begin_index = None + + def index_for_timestep( + self, timestep: float, schedule_timesteps: Optional[torch.Tensor] = None + ) -> int: + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + indices = (schedule_timesteps == timestep).nonzero() + pos = 1 if len(indices) > 1 else 0 + return indices[pos].item() + + def _init_step_index(self, timestep: Union[float, torch.Tensor]): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[Hunyuan3DConsistencyFlowMatchSchedulerOutput, Tuple]: + """Perform one step of the consistency flow matching scheduler.""" + if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): + raise ValueError("Passing integer indices as timesteps is not supported.") + + if self.step_index is None: + self._init_step_index(timestep) + + sample = sample.to(torch.float32) + + sigma = self.sigmas_[self.step_index] + sigma_next = self.sigmas_[self.step_index + 1] + + prev_sample = sample + (sigma_next - sigma) * model_output + prev_sample = prev_sample.to(model_output.dtype) + + pred_original_sample = sample + (1.0 - sigma) * model_output + pred_original_sample = pred_original_sample.to(model_output.dtype) + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + + return Hunyuan3DConsistencyFlowMatchSchedulerOutput( + prev_sample=prev_sample, pred_original_sample=pred_original_sample + ) + + def __len__(self) -> int: + return self.config.num_train_timesteps + + +# Entry class for model registry +EntryClass = [ + Hunyuan3DFlowMatchEulerDiscreteScheduler, + Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler, +] diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py new file mode 100644 index 000000000..c2792ff18 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py @@ -0,0 +1,1224 @@ +# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 + + +from __future__ import annotations + +from typing import Callable, List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat +from tqdm import tqdm + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Attention backend selection +scaled_dot_product_attention = F.scaled_dot_product_attention + + +class CrossAttentionProcessor: + def __call__(self, attn, q, k, v): + out = scaled_dot_product_attention(q, k, v) + return out + + +class FlashVDMCrossAttentionProcessor: + def __init__(self, topk=None): + self.topk = topk + + def __call__(self, attn, q, k, v): + if k.shape[-2] == 3072: + topk = 1024 + elif k.shape[-2] == 512: + topk = 256 + else: + topk = k.shape[-2] // 3 + + if self.topk is True: + q1 = q[:, :, ::100, :] + sim = q1 @ k.transpose(-1, -2) + sim = torch.mean(sim, -2) + topk_ind = torch.topk(sim, dim=-1, k=topk).indices.squeeze(-2).unsqueeze(-1) + topk_ind = topk_ind.expand(-1, -1, -1, v.shape[-1]) + v0 = torch.gather(v, dim=-2, index=topk_ind) + k0 = torch.gather(k, dim=-2, index=topk_ind) + out = scaled_dot_product_attention(q, k0, v0) + elif self.topk is False: + out = scaled_dot_product_attention(q, k, v) + else: + idx, counts = self.topk + start = 0 + outs = [] + for grid_coord, count in zip(idx, counts): + end = start + count + q_chunk = q[:, :, start:end, :] + k0, v0 = self.select_topkv(q_chunk, k, v, topk) + out = scaled_dot_product_attention(q_chunk, k0, v0) + outs.append(out) + start += count + out = torch.cat(outs, dim=-2) + self.topk = False + return out + + def select_topkv(self, q_chunk, k, v, topk): + q1 = q_chunk[:, :, ::50, :] + sim = q1 @ k.transpose(-1, -2) + sim = torch.mean(sim, -2) + topk_ind = torch.topk(sim, dim=-1, k=topk).indices.squeeze(-2).unsqueeze(-1) + topk_ind = topk_ind.expand(-1, -1, -1, v.shape[-1]) + v0 = torch.gather(v, dim=-2, index=topk_ind) + k0 = torch.gather(k, dim=-2, index=topk_ind) + return k0, v0 + + +class FlashVDMTopMCrossAttentionProcessor(FlashVDMCrossAttentionProcessor): + def select_topkv(self, q_chunk, k, v, topk): + q1 = q_chunk[:, :, ::30, :] + sim = q1 @ k.transpose(-1, -2) + # sim = sim.to(torch.float32) + sim = sim.softmax(-1) + sim = torch.mean(sim, 1) + activated_token = torch.where(sim > 1e-6)[2] + index = ( + torch.unique(activated_token, return_counts=True)[0] + .unsqueeze(0) + .unsqueeze(0) + .unsqueeze(-1) + ) + index = index.expand(-1, v.shape[1], -1, v.shape[-1]) + v0 = torch.gather(v, dim=-2, index=index) + k0 = torch.gather(k, dim=-2, index=index) + return k0, v0 + + +class FourierEmbedder(nn.Module): + def __init__( + self, + num_freqs: int = 6, + logspace: bool = True, + input_dim: int = 3, + include_input: bool = True, + include_pi: bool = True, + ) -> None: + """The initialization""" + + super().__init__() + + if logspace: + frequencies = 2.0 ** torch.arange(num_freqs, dtype=torch.float32) + else: + frequencies = torch.linspace( + 1.0, 2.0 ** (num_freqs - 1), num_freqs, dtype=torch.float32 + ) + + if include_pi: + frequencies *= torch.pi + + self.register_buffer("frequencies", frequencies, persistent=False) + self.include_input = include_input + self.num_freqs = num_freqs + + self.out_dim = self.get_dims(input_dim) + + def get_dims(self, input_dim): + temp = 1 if self.include_input or self.num_freqs == 0 else 0 + out_dim = input_dim * (self.num_freqs * 2 + temp) + + return out_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward process.""" + + if self.num_freqs > 0: + embed = (x[..., None].contiguous() * self.frequencies).view( + *x.shape[:-1], -1 + ) + if self.include_input: + return torch.cat((x, embed.sin(), embed.cos()), dim=-1) + else: + return torch.cat((embed.sin(), embed.cos()), dim=-1) + else: + return x + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * ( + x.ndim - 1 + ) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and self.scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + def extra_repr(self): + return f"drop_prob={round(self.drop_prob, 3):0.3f}" + + +class MLP(nn.Module): + def __init__( + self, + *, + width: int, + expand_ratio: int = 4, + output_width: int = None, + drop_path_rate: float = 0.0, + ): + super().__init__() + self.width = width + self.c_fc = nn.Linear(width, width * expand_ratio) + self.c_proj = nn.Linear( + width * expand_ratio, output_width if output_width is not None else width + ) + self.gelu = nn.GELU() + self.drop_path = ( + DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + ) + + def forward(self, x): + return self.drop_path(self.c_proj(self.gelu(self.c_fc(x)))) + + +class QKVMultiheadCrossAttention(nn.Module): + def __init__( + self, + *, + heads: int, + n_data: Optional[int] = None, + width=None, + qk_norm=False, + norm_layer=nn.LayerNorm, + ): + super().__init__() + self.heads = heads + self.n_data = n_data + self.q_norm = ( + norm_layer(width // heads, elementwise_affine=True, eps=1e-6) + if qk_norm + else nn.Identity() + ) + self.k_norm = ( + norm_layer(width // heads, elementwise_affine=True, eps=1e-6) + if qk_norm + else nn.Identity() + ) + + self.attn_processor = CrossAttentionProcessor() + + def forward(self, q, kv): + _, n_ctx, _ = q.shape + bs, n_data, width = kv.shape + attn_ch = width // self.heads // 2 + q = q.view(bs, n_ctx, self.heads, -1) + kv = kv.view(bs, n_data, self.heads, -1) + k, v = torch.split(kv, attn_ch, dim=-1) + + q = self.q_norm(q) + k = self.k_norm(k) + q, k, v = map( + lambda t: rearrange(t, "b n h d -> b h n d", h=self.heads), (q, k, v) + ) + out = self.attn_processor(self, q, k, v) + out = out.transpose(1, 2).reshape(bs, n_ctx, -1) + return out + + +class MultiheadCrossAttention(nn.Module): + def __init__( + self, + *, + width: int, + heads: int, + qkv_bias: bool = True, + n_data: Optional[int] = None, + data_width: Optional[int] = None, + norm_layer=nn.LayerNorm, + qk_norm: bool = False, + kv_cache: bool = False, + ): + super().__init__() + self.n_data = n_data + self.width = width + self.heads = heads + self.data_width = width if data_width is None else data_width + self.c_q = nn.Linear(width, width, bias=qkv_bias) + self.c_kv = nn.Linear(self.data_width, width * 2, bias=qkv_bias) + self.c_proj = nn.Linear(width, width) + self.attention = QKVMultiheadCrossAttention( + heads=heads, + n_data=n_data, + width=width, + norm_layer=norm_layer, + qk_norm=qk_norm, + ) + self.kv_cache = kv_cache + self.data = None + + def forward(self, x, data): + x = self.c_q(x) + if self.kv_cache: + if self.data is None: + self.data = self.c_kv(data) + logger.info( + "Save kv cache,this should be called only once for one mesh" + ) + data = self.data + else: + data = self.c_kv(data) + x = self.attention(x, data) + x = self.c_proj(x) + return x + + +class ResidualCrossAttentionBlock(nn.Module): + def __init__( + self, + *, + n_data: Optional[int] = None, + width: int, + heads: int, + mlp_expand_ratio: int = 4, + data_width: Optional[int] = None, + qkv_bias: bool = True, + norm_layer=nn.LayerNorm, + qk_norm: bool = False, + ): + super().__init__() + + if data_width is None: + data_width = width + + self.attn = MultiheadCrossAttention( + n_data=n_data, + width=width, + heads=heads, + data_width=data_width, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + qk_norm=qk_norm, + ) + self.ln_1 = norm_layer(width, elementwise_affine=True, eps=1e-6) + self.ln_2 = norm_layer(data_width, elementwise_affine=True, eps=1e-6) + self.ln_3 = norm_layer(width, elementwise_affine=True, eps=1e-6) + self.mlp = MLP(width=width, expand_ratio=mlp_expand_ratio) + + def forward(self, x: torch.Tensor, data: torch.Tensor): + x = x + self.attn(self.ln_1(x), self.ln_2(data)) + x = x + self.mlp(self.ln_3(x)) + return x + + +class QKVMultiheadAttention(nn.Module): + def __init__( + self, + *, + heads: int, + n_ctx: int, + width=None, + qk_norm=False, + norm_layer=nn.LayerNorm, + ): + super().__init__() + self.heads = heads + self.n_ctx = n_ctx + self.q_norm = ( + norm_layer(width // heads, elementwise_affine=True, eps=1e-6) + if qk_norm + else nn.Identity() + ) + self.k_norm = ( + norm_layer(width // heads, elementwise_affine=True, eps=1e-6) + if qk_norm + else nn.Identity() + ) + + def forward(self, qkv): + bs, n_ctx, width = qkv.shape + attn_ch = width // self.heads // 3 + qkv = qkv.view(bs, n_ctx, self.heads, -1) + q, k, v = torch.split(qkv, attn_ch, dim=-1) + + q = self.q_norm(q) + k = self.k_norm(k) + + q, k, v = map( + lambda t: rearrange(t, "b n h d -> b h n d", h=self.heads), (q, k, v) + ) + out = ( + scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(bs, n_ctx, -1) + ) + return out + + +class MultiheadAttention(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + heads: int, + qkv_bias: bool, + norm_layer=nn.LayerNorm, + qk_norm: bool = False, + drop_path_rate: float = 0.0, + ): + super().__init__() + self.n_ctx = n_ctx + self.width = width + self.heads = heads + self.c_qkv = nn.Linear(width, width * 3, bias=qkv_bias) + self.c_proj = nn.Linear(width, width) + self.attention = QKVMultiheadAttention( + heads=heads, + n_ctx=n_ctx, + width=width, + norm_layer=norm_layer, + qk_norm=qk_norm, + ) + self.drop_path = ( + DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + ) + + def forward(self, x): + x = self.c_qkv(x) + x = self.attention(x) + x = self.drop_path(self.c_proj(x)) + return x + + +class ResidualAttentionBlock(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + heads: int, + qkv_bias: bool = True, + norm_layer=nn.LayerNorm, + qk_norm: bool = False, + drop_path_rate: float = 0.0, + ): + super().__init__() + self.attn = MultiheadAttention( + n_ctx=n_ctx, + width=width, + heads=heads, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + qk_norm=qk_norm, + drop_path_rate=drop_path_rate, + ) + self.ln_1 = norm_layer(width, elementwise_affine=True, eps=1e-6) + self.mlp = MLP(width=width, drop_path_rate=drop_path_rate) + self.ln_2 = norm_layer(width, elementwise_affine=True, eps=1e-6) + + def forward(self, x: torch.Tensor): + x = x + self.attn(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +class Transformer(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + layers: int, + heads: int, + qkv_bias: bool = True, + norm_layer=nn.LayerNorm, + qk_norm: bool = False, + drop_path_rate: float = 0.0, + ): + super().__init__() + self.n_ctx = n_ctx + self.width = width + self.layers = layers + self.resblocks = nn.ModuleList( + [ + ResidualAttentionBlock( + n_ctx=n_ctx, + width=width, + heads=heads, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + qk_norm=qk_norm, + drop_path_rate=drop_path_rate, + ) + for _ in range(layers) + ] + ) + + def forward(self, x: torch.Tensor): + for block in self.resblocks: + x = block(x) + return x + + +class CrossAttentionDecoder(nn.Module): + + def __init__( + self, + *, + num_latents: int, + out_channels: int, + fourier_embedder: FourierEmbedder, + width: int, + heads: int, + mlp_expand_ratio: int = 4, + downsample_ratio: int = 1, + enable_ln_post: bool = True, + qkv_bias: bool = True, + qk_norm: bool = False, + label_type: str = "binary", + ): + super().__init__() + + self.enable_ln_post = enable_ln_post + self.fourier_embedder = fourier_embedder + self.downsample_ratio = downsample_ratio + self.query_proj = nn.Linear(self.fourier_embedder.out_dim, width) + if self.downsample_ratio != 1: + self.latents_proj = nn.Linear(width * downsample_ratio, width) + if self.enable_ln_post == False: + qk_norm = False + self.cross_attn_decoder = ResidualCrossAttentionBlock( + n_data=num_latents, + width=width, + mlp_expand_ratio=mlp_expand_ratio, + heads=heads, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + ) + + if self.enable_ln_post: + self.ln_post = nn.LayerNorm(width) + self.output_proj = nn.Linear(width, out_channels) + self.label_type = label_type + + def set_cross_attention_processor(self, processor): + self.cross_attn_decoder.attn.attention.attn_processor = processor + + def forward(self, queries=None, query_embeddings=None, latents=None): + if query_embeddings is None: + fourier_out = self.fourier_embedder(queries) + query_embeddings = self.query_proj(fourier_out.to(latents.dtype)) + + if self.downsample_ratio != 1: + latents = self.latents_proj(latents) + + x = self.cross_attn_decoder(query_embeddings, latents) + + if self.enable_ln_post: + x = self.ln_post(x) + + occ = self.output_proj(x) + return occ + + +def generate_dense_grid_points( + bbox_min: np.ndarray, + bbox_max: np.ndarray, + octree_resolution: int, + indexing: str = "ij", +): + length = bbox_max - bbox_min + num_cells = octree_resolution + + x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32) + y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32) + z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32) + [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing) + xyz = np.stack((xs, ys, zs), axis=-1) + grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1] + + return xyz, grid_size, length + + +def extract_near_surface_volume_fn(input_tensor: torch.Tensor, alpha: float): + """Extract near-surface voxels for hierarchical decoding.""" + device = input_tensor.device + + val = input_tensor + alpha + valid_mask = val > -9000 + + def get_neighbor(t, shift, axis): + if shift == 0: + return t.clone() + pad_dims = [0, 0, 0, 0, 0, 0] + if axis == 0: + pad_idx = 0 if shift > 0 else 1 + pad_dims[pad_idx] = abs(shift) + elif axis == 1: + pad_idx = 2 if shift > 0 else 3 + pad_dims[pad_idx] = abs(shift) + elif axis == 2: + pad_idx = 4 if shift > 0 else 5 + pad_dims[pad_idx] = abs(shift) + + padded = F.pad(t.unsqueeze(0).unsqueeze(0), pad_dims[::-1], mode="replicate") + + slice_dims = [slice(None)] * 3 + if axis == 0: + slice_dims[0] = slice(shift, None) if shift > 0 else slice(None, shift) + elif axis == 1: + slice_dims[1] = slice(shift, None) if shift > 0 else slice(None, shift) + elif axis == 2: + slice_dims[2] = slice(shift, None) if shift > 0 else slice(None, shift) + + padded = padded.squeeze(0).squeeze(0) + return padded[slice_dims] + + left = get_neighbor(val, 1, axis=0) + right = get_neighbor(val, -1, axis=0) + back = get_neighbor(val, 1, axis=1) + front = get_neighbor(val, -1, axis=1) + down = get_neighbor(val, 1, axis=2) + up = get_neighbor(val, -1, axis=2) + + def safe_where(neighbor): + return torch.where(neighbor > -9000, neighbor, val) + + left, right = safe_where(left), safe_where(right) + back, front = safe_where(back), safe_where(front) + down, up = safe_where(down), safe_where(up) + + sign = torch.sign(val.to(torch.float32)) + neighbors_sign = torch.stack( + [ + torch.sign(left.to(torch.float32)), + torch.sign(right.to(torch.float32)), + torch.sign(back.to(torch.float32)), + torch.sign(front.to(torch.float32)), + torch.sign(down.to(torch.float32)), + torch.sign(up.to(torch.float32)), + ], + dim=0, + ) + + same_sign = torch.all(neighbors_sign == sign, dim=0) + mask = (~same_sign).to(torch.int32) + return mask * valid_mask.to(torch.int32) + + +class VanillaVolumeDecoder: + """Standard volume decoder using dense grid evaluation.""" + + @torch.no_grad() + def __call__( + self, + latents: torch.FloatTensor, + geo_decoder: Callable, + bounds: Union[Tuple[float], List[float], float] = 1.01, + num_chunks: int = 10000, + octree_resolution: int = None, + enable_pbar: bool = True, + **kwargs, + ): + device = latents.device + dtype = latents.dtype + batch_size = latents.shape[0] + + if isinstance(bounds, float): + bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] + + bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6]) + xyz_samples, grid_size, length = generate_dense_grid_points( + bbox_min=bbox_min, + bbox_max=bbox_max, + octree_resolution=octree_resolution, + indexing="ij", + ) + xyz_samples = ( + torch.from_numpy(xyz_samples) + .to(device, dtype=dtype) + .contiguous() + .reshape(-1, 3) + ) + + batch_logits = [] + for start in tqdm( + range(0, xyz_samples.shape[0], num_chunks), + desc="Volume Decoding", + disable=not enable_pbar, + ): + chunk_queries = xyz_samples[start : start + num_chunks, :] + chunk_queries = repeat(chunk_queries, "p c -> b p c", b=batch_size) + logits = geo_decoder(queries=chunk_queries, latents=latents) + batch_logits.append(logits) + + grid_logits = torch.cat(batch_logits, dim=1) + grid_logits = grid_logits.view((batch_size, *grid_size)).float() + + return grid_logits + + +class HierarchicalVolumeDecoding: + """Hierarchical volume decoder with multi-resolution refinement.""" + + @torch.no_grad() + def __call__( + self, + latents: torch.FloatTensor, + geo_decoder: Callable, + bounds: Union[Tuple[float], List[float], float] = 1.01, + num_chunks: int = 10000, + mc_level: float = 0.0, + octree_resolution: int = None, + min_resolution: int = 63, + enable_pbar: bool = True, + **kwargs, + ): + device = latents.device + dtype = latents.dtype + + resolutions = [] + if octree_resolution < min_resolution: + resolutions.append(octree_resolution) + while octree_resolution >= min_resolution: + resolutions.append(octree_resolution) + octree_resolution = octree_resolution // 2 + resolutions.reverse() + + # 1. generate query points + if isinstance(bounds, float): + bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] + bbox_min = np.array(bounds[0:3]) + bbox_max = np.array(bounds[3:6]) + bbox_size = bbox_max - bbox_min + + xyz_samples, grid_size, length = generate_dense_grid_points( + bbox_min=bbox_min, + bbox_max=bbox_max, + octree_resolution=resolutions[0], + indexing="ij", + ) + + dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype) + dilate.weight = torch.nn.Parameter( + torch.ones(dilate.weight.shape, dtype=dtype, device=device) + ) + + grid_size = np.array(grid_size) + xyz_samples = ( + torch.from_numpy(xyz_samples) + .to(device, dtype=dtype) + .contiguous() + .reshape(-1, 3) + ) + + # 2. latents to 3d volume + batch_logits = [] + batch_size = latents.shape[0] + for start in tqdm( + range(0, xyz_samples.shape[0], num_chunks), + desc=f"Hierarchical Volume Decoding [r{resolutions[0] + 1}]", + disable=not enable_pbar, + ): + queries = xyz_samples[start : start + num_chunks, :] + batch_queries = repeat(queries, "p c -> b p c", b=batch_size) + logits = geo_decoder(queries=batch_queries, latents=latents) + batch_logits.append(logits) + + grid_logits = torch.cat(batch_logits, dim=1).view( + (batch_size, grid_size[0], grid_size[1], grid_size[2]) + ) + + for octree_depth_now in resolutions[1:]: + grid_size = np.array([octree_depth_now + 1] * 3) + resolution = bbox_size / octree_depth_now + next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device) + next_logits = torch.full( + next_index.shape, -10000.0, dtype=dtype, device=device + ) + curr_points = extract_near_surface_volume_fn( + grid_logits.squeeze(0), mc_level + ) + curr_points += grid_logits.squeeze(0).abs() < 0.95 + + if octree_depth_now == resolutions[-1]: + expand_num = 0 + else: + expand_num = 1 + for i in range(expand_num): + curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0) + cidx_x, cidx_y, cidx_z = torch.where(curr_points > 0) + next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1 + for i in range(2 - expand_num): + next_index = dilate(next_index.unsqueeze(0)).squeeze(0) + nidx = torch.where(next_index > 0) + + next_points = torch.stack(nidx, dim=1) + next_points = next_points * torch.tensor( + resolution, dtype=next_points.dtype, device=device + ) + torch.tensor(bbox_min, dtype=next_points.dtype, device=device) + + # Check if next_points is empty + if next_points.shape[0] == 0: + logger.warning( + f"No valid surface points found at resolution {octree_depth_now}, " + f"skipping this level" + ) + continue + + batch_logits = [] + for start in tqdm( + range(0, next_points.shape[0], num_chunks), + desc=f"Hierarchical Volume Decoding [r{octree_depth_now + 1}]", + disable=not enable_pbar, + ): + queries = next_points[start : start + num_chunks, :] + batch_queries = repeat(queries, "p c -> b p c", b=batch_size) + logits = geo_decoder( + queries=batch_queries.to(latents.dtype), latents=latents + ) + batch_logits.append(logits) + grid_logits = torch.cat(batch_logits, dim=1) + next_logits[nidx] = grid_logits[0, ..., 0] + grid_logits = next_logits.unsqueeze(0) + grid_logits[grid_logits == -10000.0] = float("nan") + + return grid_logits + + +class FlashVDMVolumeDecoding: + """Flash VDM volume decoder with adaptive KV selection.""" + + def __init__(self, topk_mode="mean"): + if topk_mode not in ["mean", "merge"]: + raise ValueError(f"Unsupported topk_mode {topk_mode}") + + if topk_mode == "mean": + self.processor = FlashVDMCrossAttentionProcessor() + else: + self.processor = FlashVDMTopMCrossAttentionProcessor() + + @torch.no_grad() + def __call__( + self, + latents: torch.FloatTensor, + geo_decoder: CrossAttentionDecoder, + bounds: Union[Tuple[float], List[float], float] = 1.01, + num_chunks: int = 10000, + mc_level: float = 0.0, + octree_resolution: int = None, + min_resolution: int = 63, + mini_grid_num: int = 4, + enable_pbar: bool = True, + **kwargs, + ): + processor = self.processor + geo_decoder.set_cross_attention_processor(processor) + + device = latents.device + dtype = latents.dtype + + resolutions = [] + orig_resolution = octree_resolution + if octree_resolution < min_resolution: + resolutions.append(octree_resolution) + while octree_resolution >= min_resolution: + resolutions.append(octree_resolution) + octree_resolution = octree_resolution // 2 + resolutions.reverse() + resolutions[0] = round(resolutions[0] / mini_grid_num) * mini_grid_num - 1 + for i, resolution in enumerate(resolutions[1:]): + resolutions[i + 1] = resolutions[0] * 2 ** (i + 1) + + if isinstance(bounds, float): + bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] + bbox_min = np.array(bounds[0:3]) + bbox_max = np.array(bounds[3:6]) + bbox_size = bbox_max - bbox_min + + xyz_samples, grid_size, length = generate_dense_grid_points( + bbox_min=bbox_min, + bbox_max=bbox_max, + octree_resolution=resolutions[0], + indexing="ij", + ) + + logger.info(f"FlashVDMVolumeDecoding Resolution: {resolutions}") + + dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype) + dilate.weight = torch.nn.Parameter( + torch.ones(dilate.weight.shape, dtype=dtype, device=device) + ) + + grid_size = np.array(grid_size) + + xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype) + batch_size = latents.shape[0] + mini_grid_size = xyz_samples.shape[0] // mini_grid_num + xyz_samples = ( + xyz_samples.view( + mini_grid_num, + mini_grid_size, + mini_grid_num, + mini_grid_size, + mini_grid_num, + mini_grid_size, + 3, + ) + .permute(0, 2, 4, 1, 3, 5, 6) + .reshape(-1, mini_grid_size * mini_grid_size * mini_grid_size, 3) + ) + + batch_logits = [] + num_batchs = max(num_chunks // xyz_samples.shape[1], 1) + for start in tqdm( + range(0, xyz_samples.shape[0], num_batchs), + desc="FlashVDM Volume Decoding", + disable=not enable_pbar, + ): + queries = xyz_samples[start : start + num_batchs, :] + batch = queries.shape[0] + batch_latents = repeat(latents.squeeze(0), "p c -> b p c", b=batch) + processor.topk = True + logits = geo_decoder(queries=queries, latents=batch_latents) + batch_logits.append(logits) + + grid_logits = ( + torch.cat(batch_logits, dim=0) + .reshape( + mini_grid_num, + mini_grid_num, + mini_grid_num, + mini_grid_size, + mini_grid_size, + mini_grid_size, + ) + .permute(0, 3, 1, 4, 2, 5) + .contiguous() + .view((batch_size, grid_size[0], grid_size[1], grid_size[2])) + ) + + for octree_depth_now in resolutions[1:]: + grid_size = np.array([octree_depth_now + 1] * 3) + resolution = bbox_size / octree_depth_now + next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device) + next_logits = torch.full( + next_index.shape, -10000.0, dtype=dtype, device=device + ) + curr_points = extract_near_surface_volume_fn( + grid_logits.squeeze(0), mc_level + ) + curr_points += grid_logits.squeeze(0).abs() < 0.95 + + expand_num = 0 if octree_depth_now == resolutions[-1] else 1 + for _ in range(expand_num): + curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0) + + cidx_x, cidx_y, cidx_z = torch.where(curr_points > 0) + next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1 + for _ in range(2 - expand_num): + next_index = dilate(next_index.unsqueeze(0)).squeeze(0) + nidx = torch.where(next_index > 0) + + next_points = torch.stack(nidx, dim=1) + next_points = next_points * torch.tensor( + resolution, dtype=torch.float32, device=device + ) + torch.tensor(bbox_min, dtype=torch.float32, device=device) + + # Check if next_points is empty (no valid surface points found) + if next_points.shape[0] == 0: + # Skip this resolution level if no points found + # Use the previous grid_logits as fallback + logger.warning( + f"No valid surface points found at resolution {octree_depth_now}, " + f"skipping this level and using previous resolution grid_logits" + ) + continue + + query_grid_num = 6 + min_val = next_points.min(axis=0).values + max_val = next_points.max(axis=0).values + vol_queries_index = ( + (next_points - min_val) / (max_val - min_val) * (query_grid_num - 0.001) + ) + index = torch.floor(vol_queries_index).long() + index = ( + index[..., 0] * (query_grid_num**2) + + index[..., 1] * query_grid_num + + index[..., 2] + ) + index = index.sort() + next_points = next_points[index.indices].unsqueeze(0).contiguous() + unique_values = torch.unique(index.values, return_counts=True) + grid_logits_flat = torch.zeros( + (next_points.shape[1]), dtype=latents.dtype, device=latents.device + ) + input_grid = [[], []] + logits_grid_list = [] + start_num = 0 + sum_num = 0 + for grid_index, count in zip( + unique_values[0].cpu().tolist(), unique_values[1].cpu().tolist() + ): + if sum_num + count < num_chunks or sum_num == 0: + sum_num += count + input_grid[0].append(grid_index) + input_grid[1].append(count) + else: + processor.topk = input_grid + logits_grid = geo_decoder( + queries=next_points[:, start_num : start_num + sum_num], + latents=latents, + ) + start_num = start_num + sum_num + logits_grid_list.append(logits_grid) + input_grid = [[grid_index], [count]] + sum_num = count + if sum_num > 0: + processor.topk = input_grid + logits_grid = geo_decoder( + queries=next_points[:, start_num : start_num + sum_num], + latents=latents, + ) + logits_grid_list.append(logits_grid) + logits_grid = torch.cat(logits_grid_list, dim=1) + grid_logits_flat[index.indices] = logits_grid.squeeze(0).squeeze(-1) + next_logits[nidx] = grid_logits_flat + grid_logits = next_logits.unsqueeze(0) + + grid_logits[grid_logits == -10000.0] = float("nan") + return grid_logits + + +class Latent2MeshOutput: + """Container for mesh output from VAE decoder.""" + + def __init__(self, mesh_v=None, mesh_f=None): + self.mesh_v = mesh_v + self.mesh_f = mesh_f + + +def center_vertices(vertices): + """Translate vertices so bounding box is centered at zero.""" + vert_min = vertices.min(dim=0)[0] + vert_max = vertices.max(dim=0)[0] + vert_center = 0.5 * (vert_min + vert_max) + return vertices - vert_center + + +class SurfaceExtractor: + """Base class for surface extraction algorithms.""" + + def _compute_box_stat( + self, bounds: Union[Tuple[float], List[float], float], octree_resolution: int + ): + if isinstance(bounds, float): + bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] + + bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6]) + bbox_size = bbox_max - bbox_min + grid_size = [ + int(octree_resolution) + 1, + int(octree_resolution) + 1, + int(octree_resolution) + 1, + ] + return grid_size, bbox_min, bbox_size + + def run(self, *args, **kwargs): + raise NotImplementedError + + def __call__(self, grid_logits, **kwargs): + outputs = [] + for i in range(grid_logits.shape[0]): + try: + vertices, faces = self.run(grid_logits[i], **kwargs) + vertices = vertices.astype(np.float32) + faces = np.ascontiguousarray(faces) + outputs.append(Latent2MeshOutput(mesh_v=vertices, mesh_f=faces)) + except Exception: + import traceback + + traceback.print_exc() + outputs.append(None) + return outputs + + +class MCSurfaceExtractor(SurfaceExtractor): + """Marching Cubes surface extractor.""" + + def run(self, grid_logit, *, mc_level, bounds, octree_resolution, **kwargs): + from skimage import measure + + vertices, faces, normals, _ = measure.marching_cubes( + grid_logit.cpu().numpy(), mc_level, method="lewiner" + ) + grid_size, bbox_min, bbox_size = self._compute_box_stat( + bounds, octree_resolution + ) + vertices = vertices / grid_size * bbox_size + bbox_min + return vertices, faces + + +class DMCSurfaceExtractor(SurfaceExtractor): + """Differentiable Marching Cubes surface extractor.""" + + def run(self, grid_logit, *, octree_resolution, **kwargs): + device = grid_logit.device + if not hasattr(self, "dmc"): + try: + from diso import DiffDMC + + self.dmc = DiffDMC(dtype=torch.float32).to(device) + except ImportError: + raise ImportError( + "Please install diso via `pip install diso`, or set mc_algo to 'mc'" + ) + sdf = -grid_logit / octree_resolution + sdf = sdf.to(torch.float32).contiguous() + verts, faces = self.dmc(sdf, deform=None, return_quads=False, normalize=True) + verts = center_vertices(verts) + vertices = verts.detach().cpu().numpy() + faces = faces.detach().cpu().numpy()[:, ::-1] + return vertices, faces + + +SurfaceExtractors = { + "mc": MCSurfaceExtractor, + "dmc": DMCSurfaceExtractor, +} + + +class VectsetVAE(nn.Module): + """Base VAE class for vector set encoding.""" + + def __init__(self, volume_decoder=None, surface_extractor=None): + super().__init__() + if volume_decoder is None: + volume_decoder = VanillaVolumeDecoder() + if surface_extractor is None: + surface_extractor = MCSurfaceExtractor() + self.volume_decoder = volume_decoder + self.surface_extractor = surface_extractor + + def latents2mesh(self, latents: torch.FloatTensor, **kwargs): + """Convert latents to mesh.""" + grid_logits = self.volume_decoder(latents, self.geo_decoder, **kwargs) + outputs = self.surface_extractor(grid_logits, **kwargs) + return outputs + + def enable_flashvdm_decoder( + self, + enabled: bool = True, + adaptive_kv_selection=True, + topk_mode="mean", + mc_algo="dmc", + ): + """Enable or disable FlashVDM decoder for faster inference.""" + if enabled: + if adaptive_kv_selection: + self.volume_decoder = FlashVDMVolumeDecoding(topk_mode) + else: + self.volume_decoder = HierarchicalVolumeDecoding() + if mc_algo not in SurfaceExtractors: + raise ValueError( + f"Unsupported mc_algo {mc_algo}, available: {list(SurfaceExtractors.keys())}" + ) + self.surface_extractor = SurfaceExtractors[mc_algo]() + else: + self.volume_decoder = VanillaVolumeDecoder() + self.surface_extractor = MCSurfaceExtractor() + + +class ShapeVAE(VectsetVAE): + """Shape VAE for 3D mesh generation from latent codes.""" + + _aliases = ["hy3dgen.shapegen.models.ShapeVAE"] + + def __init__( + self, + *, + num_latents: int, + embed_dim: int, + width: int, + heads: int, + num_decoder_layers: int, + num_encoder_layers: int = 8, + pc_size: int = 5120, + pc_sharpedge_size: int = 5120, + point_feats: int = 3, + downsample_ratio: int = 20, + geo_decoder_downsample_ratio: int = 1, + geo_decoder_mlp_expand_ratio: int = 4, + geo_decoder_ln_post: bool = True, + num_freqs: int = 8, + include_pi: bool = True, + qkv_bias: bool = True, + qk_norm: bool = False, + label_type: str = "binary", + drop_path_rate: float = 0.0, + scale_factor: float = 1.0, + use_ln_post: bool = True, + ckpt_path=None, + ): + super().__init__() + self.geo_decoder_ln_post = geo_decoder_ln_post + self.downsample_ratio = downsample_ratio + + self.fourier_embedder = FourierEmbedder( + num_freqs=num_freqs, include_pi=include_pi + ) + + self.post_kl = nn.Linear(embed_dim, width) + + self.transformer = Transformer( + n_ctx=num_latents, + width=width, + layers=num_decoder_layers, + heads=heads, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + drop_path_rate=drop_path_rate, + ) + + self.geo_decoder = CrossAttentionDecoder( + fourier_embedder=self.fourier_embedder, + out_channels=1, + num_latents=num_latents, + mlp_expand_ratio=geo_decoder_mlp_expand_ratio, + downsample_ratio=geo_decoder_downsample_ratio, + enable_ln_post=self.geo_decoder_ln_post, + width=width // geo_decoder_downsample_ratio, + heads=heads // geo_decoder_downsample_ratio, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + label_type=label_type, + ) + + self.scale_factor = scale_factor + self.latent_shape = (num_latents, embed_dim) + + def forward(self, latents): + latents = self.post_kl(latents) + latents = self.transformer(latents) + return latents + + def decode(self, latents): + """Decode latents to features.""" + latents = self.post_kl(latents) + latents = self.transformer(latents) + return latents + + +# Entry class for model registry +EntryClass = ShapeVAE diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py new file mode 100644 index 000000000..c227ef89a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py @@ -0,0 +1,411 @@ +""" +Hunyuan3D image-to-mesh pipeline implementation. + +Shape pipeline: BeforeDenoising -> Denoising -> Export -> Save +Paint pipeline (optional): Preprocess -> TexGen -> Postprocess +""" + +from __future__ import annotations + +import glob +import importlib +import os +from itertools import chain +from typing import Any + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, +) +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.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + Hunyuan3DPaintPostprocessStage, + Hunyuan3DPaintPreprocessStage, + Hunyuan3DPaintTexGenStage, + Hunyuan3DShapeBeforeDenoisingStage, + Hunyuan3DShapeDenoisingStage, + Hunyuan3DShapeExportStage, + Hunyuan3DShapeSaveStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class Hunyuan3D2Pipeline(ComposedPipelineBase): + """Hunyuan3D 2.0 image-to-mesh pipeline. + + Shape pipeline: BeforeDenoising -> Denoising -> Export -> Save + Paint pipeline (optional): Preprocess -> TexGen -> Postprocess + """ + + pipeline_name = "Hunyuan3D2Pipeline" + _required_config_modules = [ + "hy3dshape_model", + "hy3dshape_vae", + "hy3dshape_scheduler", + "hy3dshape_conditioner", + "hy3dshape_image_processor", + ] + + def _load_config(self) -> dict[str, Any]: + return { + "_class_name": self.pipeline_name, + "_diffusers_version": "0.0.0", + "hy3dshape_model": ["diffusers", "Hunyuan3DShapeModel"], + "hy3dshape_vae": ["diffusers", "Hunyuan3DShapeVAE"], + "hy3dshape_scheduler": ["diffusers", "Hunyuan3DShapeScheduler"], + "hy3dshape_conditioner": ["diffusers", "Hunyuan3DShapeConditioner"], + "hy3dshape_image_processor": ["diffusers", "Hunyuan3DShapeImageProcessor"], + } + + # Class resolution + @staticmethod + def _resolve_class(target: str) -> Any: + """Resolve a YAML target string to a Python class.""" + from sglang.multimodal_gen.runtime.models.registry import ModelRegistry + + cls = ModelRegistry.resolve_by_alias(target) + if cls is not None: + return cls + + class_name = target.rsplit(".", 1)[-1] + try: + cls, _ = ModelRegistry.resolve_model_cls(class_name) + return cls + except Exception: + pass + + from sglang.multimodal_gen.runtime.utils.mesh3d_utils import ( + resolve_hunyuan3d_tool, + ) + + for name in (target, class_name): + tool_cls = resolve_hunyuan3d_tool(name) + if tool_cls is not None: + return tool_cls + + module, cls_name = target.rsplit(".", 1) + return getattr(importlib.import_module(module, package=None), cls_name) + + # Path / checkpoint resolution + @staticmethod + def _resolve_shape_dir( + model_path: str, + subfolder: str, + use_safetensors: bool, + variant: str | None, + ) -> tuple[str, str]: + """Locate (or download) the shape subfolder and return (config_path, ckpt_path).""" + local_path = os.path.join(model_path, subfolder) + if not os.path.exists(local_path): + local_path = os.path.expanduser(local_path) + + if not os.path.exists(local_path): + logger.info( + "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}/*"], + ) + local_path = os.path.join(downloaded, subfolder) + + config_path = os.path.join(local_path, "config.yaml") + if not os.path.exists(config_path): + for alt in ("config.yml", "model_config.yaml"): + alt_path = os.path.join(local_path, alt) + if os.path.exists(alt_path): + config_path = alt_path + break + + if use_safetensors: + ckpt_name = ( + f"model.{variant}.safetensors" if variant else "model.safetensors" + ) + else: + ckpt_name = f"model-{variant}.ckpt" if variant else "model.ckpt" + + ckpt_path = os.path.join(local_path, ckpt_name) + if not os.path.exists(ckpt_path): + pattern = "*.safetensors" if use_safetensors else "*.ckpt" + files = glob.glob(os.path.join(local_path, pattern)) + if files: + ckpt_path = files[0] + + logger.info("Config path: %s", config_path) + logger.info("Checkpoint path: %s", ckpt_path) + 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.""" + local_path = os.path.join(model_path, subfolder) + if not os.path.exists(local_path): + local_path = os.path.expanduser(local_path) + + if not os.path.exists(local_path): + logger.info( + "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}/*"], + ) + 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): + raise FileNotFoundError( + f"Paint model incomplete: {config_file} not found. " + "Download the model or check network connectivity." + ) + + logger.info("Resolved paint model directory: %s", local_path) + return local_path + + @staticmethod + def _load_and_split_checkpoint( + ckpt_path: str, use_safetensors: bool + ) -> 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") + ckpt: dict[str, dict[str, torch.Tensor]] = {} + for key, value in flat.items(): + component = key.split(".")[0] + sub_key = key[len(component) + 1 :] + ckpt.setdefault(component, {})[sub_key] = value + return ckpt + else: + return torch.load(ckpt_path, map_location="cpu", weights_only=True) + + # Component loading helpers + @classmethod + def _load_dit_model( + cls, + cfg: dict[str, Any], + weights: dict[str, torch.Tensor], + device: torch.device, + dtype: torch.dtype, + ) -> nn.Module: + """Load the DiT model using meta-device instantiation + standard weight loading.""" + if "target" not in cfg: + raise KeyError("Expected key 'target' in model config.") + target_cls = cls._resolve_class(cfg["target"]) + params = cfg.get("params", {}) + + if hasattr(target_cls, "build_config_from_params"): + dit_config = target_cls.build_config_from_params(params) + init_kwargs: dict[str, Any] = {"config": dit_config, "hf_config": {}} + else: + init_kwargs = params + + with set_default_torch_dtype(dtype), torch.device("meta"): + model = target_cls(**init_kwargs) + + weight_iterator = ((k, v) for k, v in weights.items()) + param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping) + + load_model_from_full_model_state_dict( + model, + weight_iterator, + device, + dtype, + strict=False, + param_names_mapping=param_names_mapping_fn, + ) + + for name, p in chain(model.named_parameters(), model.named_buffers()): + if p.is_meta: + raise RuntimeError(f"Unexpected param or buffer {name} on meta device.") + if isinstance(p, nn.Parameter): + p.requires_grad = False + + return model.eval() + + @classmethod + def _load_simple_component( + cls, + cfg: dict[str, Any], + weights: dict[str, torch.Tensor] | None, + device: torch.device, + dtype: torch.dtype, + ) -> nn.Module: + """Load a component (VAE / conditioner) with direct instantiation + state_dict.""" + if "target" not in cfg: + raise KeyError("Expected key 'target' in component config.") + target_cls = cls._resolve_class(cfg["target"]) + params = cfg.get("params", {}) + + with set_default_torch_dtype(dtype): + component = target_cls(**params) + + if weights is not None: + component.load_state_dict(weights, strict=False) + + component.to(device=device, dtype=dtype) + return component.eval() + + @classmethod + def _instantiate_component(cls, cfg: dict[str, Any]) -> Any: + """Instantiate a lightweight component (scheduler / image_processor) without weights.""" + if "target" not in cfg: + raise KeyError("Expected key 'target' in component config.") + target_cls = cls._resolve_class(cfg["target"]) + params = cfg.get("params", {}) + return target_cls(**params) + + # 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 + + config = server_args.pipeline_config + if not isinstance(config, Hunyuan3D2PipelineConfig): + raise TypeError(f"Expected Hunyuan3D2PipelineConfig, got {type(config)}") + + model_path = config.shape_model_path or server_args.model_path + + logger.info("Loading Hunyuan3D shape models from %s", model_path) + + config_path, ckpt_path = self._resolve_shape_dir( + model_path, + config.shape_subfolder, + config.shape_use_safetensors, + config.shape_variant, + ) + + with open(config_path, "r") as f: + model_config = yaml.safe_load(f) + + ckpt = self._load_and_split_checkpoint(ckpt_path, config.shape_use_safetensors) + + dtype = torch.float16 + if config.shape_variant and "bf16" in config.shape_variant: + dtype = torch.bfloat16 + device = get_local_torch_device() + + components: dict[str, Any] = {} + + components["hy3dshape_model"] = self._load_dit_model( + model_config["model"], ckpt["model"], device, dtype + ) + + components["hy3dshape_vae"] = self._load_simple_component( + model_config["vae"], ckpt.get("vae"), device, dtype + ) + + components["hy3dshape_conditioner"] = self._load_simple_component( + model_config["conditioner"], ckpt.get("conditioner"), device, dtype + ) + + components["hy3dshape_scheduler"] = self._instantiate_component( + model_config["scheduler"] + ) + components["hy3dshape_image_processor"] = self._instantiate_component( + 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) + + return components + + # Pipeline lifecycle + def initialize_pipeline(self, server_args: ServerArgs): + config = server_args.pipeline_config + if not isinstance(config, Hunyuan3D2PipelineConfig): + raise TypeError( + "Hunyuan3D2Pipeline requires Hunyuan3D2PipelineConfig, " + f"got {type(config)}" + ) + + def create_pipeline_stages(self, server_args: ServerArgs): + config = server_args.pipeline_config + assert isinstance(config, Hunyuan3D2PipelineConfig) + + # Shape: 4 stages + self.add_stage( + stage_name="shape_before_denoising", + stage=Hunyuan3DShapeBeforeDenoisingStage( + image_processor=self.get_module("hy3dshape_image_processor"), + conditioner=self.get_module("hy3dshape_conditioner"), + vae=self.get_module("hy3dshape_vae"), + model=self.get_module("hy3dshape_model"), + scheduler=self.get_module("hy3dshape_scheduler"), + config=config, + ), + ) + self.add_stage( + stage_name="shape_denoising", + stage=Hunyuan3DShapeDenoisingStage( + transformer=self.get_module("hy3dshape_model"), + scheduler=self.get_module("hy3dshape_scheduler"), + ), + ) + self.add_stage( + stage_name="shape_export", + stage=Hunyuan3DShapeExportStage( + vae=self.get_module("hy3dshape_vae"), + config=config, + ), + ) + self.add_stage( + stage_name="shape_save", + stage=Hunyuan3DShapeSaveStage(config=config), + ) + + # Paint: 3 stages (optional) + if config.paint_enable: + self.add_stage( + stage_name="paint_preprocess", + stage=Hunyuan3DPaintPreprocessStage(config=config), + ) + self.add_stage( + stage_name="paint_texgen", + stage=Hunyuan3DPaintTexGenStage( + config=config, + paint_dir=self.get_module("hy3dpaint_dir"), + ), + ) + self.add_stage( + stage_name="paint_postprocess", + stage=Hunyuan3DPaintPostprocessStage(config=config), + ) + + +EntryClass = Hunyuan3D2Pipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index bc43cbf61..6f8e14221 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -27,6 +27,21 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import ( DmdDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.encoding import EncodingStage + +# Hunyuan3D paint stages +from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_paint import ( + Hunyuan3DPaintPostprocessStage, + Hunyuan3DPaintPreprocessStage, + Hunyuan3DPaintTexGenStage, +) + +# Hunyuan3D shape stages +from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import ( + Hunyuan3DShapeBeforeDenoisingStage, + Hunyuan3DShapeDenoisingStage, + Hunyuan3DShapeExportStage, + Hunyuan3DShapeSaveStage, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import ( ImageEncodingStage, ImageVAEEncodingStage, @@ -68,4 +83,13 @@ __all__ = [ "ImageVAEEncodingStage", "TextEncodingStage", "LTX2TextConnectorStage", + # Hunyuan3D shape stages + "Hunyuan3DShapeBeforeDenoisingStage", + "Hunyuan3DShapeDenoisingStage", + "Hunyuan3DShapeExportStage", + "Hunyuan3DShapeSaveStage", + # Hunyuan3D paint stages + "Hunyuan3DPaintPreprocessStage", + "Hunyuan3DPaintTexGenStage", + "Hunyuan3DPaintPostprocessStage", ] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py new file mode 100644 index 000000000..15baaedd6 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py @@ -0,0 +1,1066 @@ +""" +Hunyuan3D paint/texture generation stages. + +Three-stage pipeline: Preprocess -> TexGen -> Postprocess. +""" + +from __future__ import annotations + +import os +from typing import Any + +import numpy as np +import torch +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import ( + retrieve_timesteps, +) +from einops import rearrange + +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, +) +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + StageParallelismType, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + StageValidators as V, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + VerificationResult, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# Utility functions +def guidance_scale_embedding( + w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Generate guidance scale embeddings.""" + assert len(w.shape) == 1 + w = w * 1000.0 + + half_dim = embedding_dim // 2 + emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) + emb = w.to(dtype)[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1)) + assert emb.shape == (w.shape[0], embedding_dim) + return emb + + +def extract_into_tensor( + a: torch.Tensor, t: torch.Tensor, x_shape: tuple, n_gen: int +) -> torch.Tensor: + """Extract values from tensor and reshape for multi-view generation.""" + out = a.gather(-1, t) + out = out.repeat(n_gen) + out = rearrange(out, "(b n) -> b n", n=n_gen) + b, c, *_ = out.shape + return out.reshape(b, c, *((1,) * (len(x_shape) - 2))) + + +def get_predicted_original_sample( + model_output: torch.Tensor, + timesteps: torch.Tensor, + sample: torch.Tensor, + prediction_type: str, + alphas: torch.Tensor, + sigmas: torch.Tensor, + n_gen: int, +) -> torch.Tensor: + """Get predicted original sample from model output.""" + alphas = extract_into_tensor(alphas, timesteps, sample.shape, n_gen) + sigmas = extract_into_tensor(sigmas, timesteps, sample.shape, n_gen) + model_output = rearrange(model_output, "(b n) c h w -> b n c h w", n=n_gen) + + if prediction_type == "epsilon": + pred_x_0 = (sample - sigmas * model_output) / alphas + elif prediction_type == "sample": + pred_x_0 = model_output + elif prediction_type == "v_prediction": + pred_x_0 = alphas * sample - sigmas * model_output + else: + raise ValueError( + f"Prediction type {prediction_type} is not supported; " + "currently, `epsilon`, `sample`, and `v_prediction` are supported." + ) + + return pred_x_0 + + +def get_predicted_noise( + model_output: torch.Tensor, + timesteps: torch.Tensor, + sample: torch.Tensor, + prediction_type: str, + alphas: torch.Tensor, + sigmas: torch.Tensor, + n_gen: int, +) -> torch.Tensor: + """Get predicted noise from model output.""" + alphas = extract_into_tensor(alphas, timesteps, sample.shape, n_gen) + sigmas = extract_into_tensor(sigmas, timesteps, sample.shape, n_gen) + model_output = rearrange(model_output, "(b n) c h w -> b n c h w", n=n_gen) + + if prediction_type == "epsilon": + pred_epsilon = model_output + elif prediction_type == "sample": + pred_epsilon = (sample - alphas * model_output) / sigmas + elif prediction_type == "v_prediction": + pred_epsilon = alphas * model_output + sigmas * sample + else: + raise ValueError( + f"Prediction type {prediction_type} is not supported; " + "currently, `epsilon`, `sample`, and `v_prediction` are supported." + ) + + return pred_epsilon + + +def to_rgb_image(maybe_rgba): + """Convert RGBA image to RGB.""" + from PIL import Image + + if maybe_rgba.mode == "RGB": + return maybe_rgba + if maybe_rgba.mode == "RGBA": + rgba = maybe_rgba + img = np.random.randint( + 127, 128, size=[rgba.size[1], rgba.size[0], 3], dtype=np.uint8 + ) + img = Image.fromarray(img, "RGB") + img.paste(rgba, mask=rgba.getchannel("A")) + return img + raise ValueError(f"Unsupported image type: {maybe_rgba.mode}") + + +class DDIMSolver: + """DDIM solver for fast sampling.""" + + def __init__( + self, + alpha_cumprods: np.ndarray, + timesteps: int = 1000, + ddim_timesteps: int = 50, + ): + step_ratio = timesteps // ddim_timesteps + self.ddim_timesteps = ( + np.arange(1, ddim_timesteps + 1) * step_ratio + ).round().astype(np.int64) - 1 + self.ddim_alpha_cumprods = alpha_cumprods[self.ddim_timesteps] + self.ddim_alpha_cumprods_prev = np.asarray( + [alpha_cumprods[0]] + alpha_cumprods[self.ddim_timesteps[:-1]].tolist() + ) + self.ddim_timesteps = torch.from_numpy(self.ddim_timesteps).long() + self.ddim_alpha_cumprods = torch.from_numpy(self.ddim_alpha_cumprods) + self.ddim_alpha_cumprods_prev = torch.from_numpy(self.ddim_alpha_cumprods_prev) + + def to(self, device: torch.device) -> "DDIMSolver": + self.ddim_timesteps = self.ddim_timesteps.to(device) + self.ddim_alpha_cumprods = self.ddim_alpha_cumprods.to(device) + self.ddim_alpha_cumprods_prev = self.ddim_alpha_cumprods_prev.to(device) + return self + + def ddim_step( + self, + pred_x0: torch.Tensor, + pred_noise: torch.Tensor, + timestep_index: torch.Tensor, + n_gen: int, + ) -> torch.Tensor: + alpha_cumprod_prev = extract_into_tensor( + self.ddim_alpha_cumprods_prev, timestep_index, pred_x0.shape, n_gen + ) + dir_xt = (1.0 - alpha_cumprod_prev).sqrt() * pred_noise + x_prev = alpha_cumprod_prev.sqrt() * pred_x0 + dir_xt + return x_prev + + +def _recorrect_rgb( + src_image: torch.Tensor, + target_image: torch.Tensor, + alpha_channel: torch.Tensor, + scale: float = 0.95, +) -> torch.Tensor: + """Correct RGB values to match target color distribution.""" + + def flat_and_mask(bgr, a): + mask = torch.where(a > 0.5, True, False) + bgr_flat = bgr.reshape(-1, bgr.shape[-1]) + mask_flat = mask.reshape(-1) + bgr_flat_masked = bgr_flat[mask_flat, :] + return bgr_flat_masked + + src_flat = flat_and_mask(src_image, alpha_channel) + target_flat = flat_and_mask(target_image, alpha_channel) + corrected_bgr = torch.zeros_like(src_image) + + for i in range(3): + src_mean, src_stddev = torch.mean(src_flat[:, i]), torch.std(src_flat[:, i]) + target_mean, target_stddev = torch.mean(target_flat[:, i]), torch.std( + target_flat[:, i] + ) + corrected_bgr[:, :, i] = torch.clamp( + (src_image[:, :, i] - scale * src_mean) * (target_stddev / src_stddev) + + scale * target_mean, + 0, + 1, + ) + + src_mse = torch.mean((src_image - target_image) ** 2) + modify_mse = torch.mean((corrected_bgr - target_image) ** 2) + if src_mse < modify_mse: + corrected_bgr = torch.cat([src_image, alpha_channel], dim=-1) + else: + corrected_bgr = torch.cat([corrected_bgr, alpha_channel], dim=-1) + + return corrected_bgr + + +# Stage 1: Preprocess (UV unwrap + delight + multi-view rendering) +class Hunyuan3DPaintPreprocessStage(PipelineStage): + """Preprocessing: UV unwrap + delight in parallel, then multi-view rendering.""" + + CAMERA_AZIMS = [0, 90, 180, 270, 0, 180] + CAMERA_ELEVS = [0, 0, 0, 0, 90, -90] + VIEW_WEIGHTS = [1, 0.1, 0.5, 0.1, 0.05, 0.05] + + @property + def parallelism_type(self) -> StageParallelismType: + return StageParallelismType.MAIN_RANK_ONLY + + def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: + super().__init__() + self.config = config + self._delight_pipeline = None + self._delight_loaded = False + self._renderer = None + self._renderer_loaded = False + + # --- UV unwrap --- + + def _do_uv_unwrap(self, batch: Req, server_args: ServerArgs) -> Req: + import time + + from sglang.multimodal_gen.runtime.utils.mesh3d_utils import mesh_uv_wrap + + mesh = batch.extra["shape_meshes"] + if isinstance(mesh, list): + mesh = mesh[0] + + try: + start_time = time.time() + mesh = mesh_uv_wrap(mesh) + elapsed = time.time() - start_time + logger.info(f"UV unwrapping completed in {elapsed:.2f}s") + except Exception as e: + logger.warning(f"UV unwrapping failed: {e}") + + batch.extra["paint_mesh"] = mesh + return batch + + # --- Delight --- + + def _load_delight_model(self, server_args: ServerArgs): + if self._delight_loaded: + return + + from diffusers import ( + EulerAncestralDiscreteScheduler, + StableDiffusionInstructPix2PixPipeline, + ) + from huggingface_hub import snapshot_download + + model_path = server_args.model_path + delight_subfolder = getattr( + self.config, "delight_subfolder", "hunyuan3d-delight-v2-0" + ) + + local_path = os.path.join(model_path, delight_subfolder) + if not os.path.exists(local_path): + local_path = os.path.expanduser(local_path) + + if not os.path.exists(local_path): + try: + downloaded = snapshot_download( + repo_id=model_path, + allow_patterns=[f"{delight_subfolder}/*"], + ) + local_path = os.path.join(downloaded, delight_subfolder) + except Exception as e: + logger.warning("Could not download delight model: %s", e) + local_path = None + + if local_path and os.path.exists(local_path): + pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained( + local_path, + torch_dtype=torch.float16, + safety_checker=None, + ) + pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( + pipeline.scheduler.config + ) + pipeline.set_progress_bar_config(disable=True) + self._delight_pipeline = pipeline.to(self.device, torch.float16) + logger.info("Delight model loaded successfully") + else: + logger.warning( + "Delight model not available, skipping delight preprocessing" + ) + + self._delight_loaded = True + + @torch.no_grad() + def _run_delight(self, image): + import cv2 + from PIL import Image as PILImage + + image = image.resize((512, 512)) + + if image.mode == "RGBA": + image_array = np.array(image) + alpha_channel = image_array[:, :, 3] + erosion_size = 3 + kernel = np.ones((erosion_size, erosion_size), np.uint8) + alpha_channel = cv2.erode(alpha_channel, kernel, iterations=1) + image_array[alpha_channel == 0, :3] = 255 + image_array[:, :, 3] = alpha_channel + image = PILImage.fromarray(image_array) + + image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device) + alpha = image_tensor[:, :, 3:] + rgb_target = image_tensor[:, :, :3] + else: + image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device) + alpha = torch.ones_like(image_tensor)[:, :, :1] + rgb_target = image_tensor[:, :, :3] + + image = image.convert("RGB") + + image = self._delight_pipeline( + prompt=self.config.delight_prompt, + image=image, + generator=torch.manual_seed(42), + height=512, + width=512, + num_inference_steps=self.config.delight_num_inference_steps, + image_guidance_scale=self.config.delight_cfg_image, + guidance_scale=self.config.delight_guidance_scale, + ).images[0] + + image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device) + rgb_src = image_tensor[:, :, :3] + image = _recorrect_rgb(rgb_src, rgb_target, alpha) + image = image[:, :, :3] * image[:, :, 3:] + torch.ones_like(image[:, :, :3]) * ( + 1.0 - image[:, :, 3:] + ) + image = PILImage.fromarray((image.cpu().numpy() * 255).astype(np.uint8)) + + return image + + def _do_delight(self, batch: Req, server_args: ServerArgs) -> Req: + from PIL import Image + + from sglang.multimodal_gen.runtime.utils.mesh3d_utils import recenter_image + + image = Image.open(batch.image_path) + image = recenter_image(image) + + if not self.config.delight_enable: + logger.info("Delight preprocessing disabled, using original image") + batch.extra["delighted_image"] = image + return batch + + self._load_delight_model(server_args) + if self._delight_pipeline is not None: + try: + image = self._run_delight(image) + logger.info("Image delight completed") + except Exception as e: + logger.warning(f"Image delight failed: {e}") + + batch.extra["delighted_image"] = image + return batch + + # --- Multi-view rendering --- + + def _init_renderer(self): + if self._renderer_loaded: + return + + from sglang.multimodal_gen.runtime.utils.mesh3d_utils import MeshRender + + self._renderer = MeshRender( + default_resolution=self.config.paint_render_size, + texture_size=self.config.paint_texture_size, + ) + self._renderer_loaded = True + logger.info("Mesh renderer initialized") + + def _render_multiview(self, mesh) -> tuple: + self._init_renderer() + self._renderer.load_mesh(mesh) + + normal_maps = self._renderer.render_normal_multiview( + self.CAMERA_ELEVS, self.CAMERA_AZIMS, use_abs_coor=True + ) + position_maps = self._renderer.render_position_multiview( + self.CAMERA_ELEVS, self.CAMERA_AZIMS + ) + + logger.info(f"Rendered {len(normal_maps)} views for texture generation") + return normal_maps, position_maps + + # --- Forward --- + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.extra.get("_mesh_failed"): + logger.warning("Mesh generation failed, skipping paint preprocessing") + batch.extra["paint_mesh"] = None + batch.extra["delighted_image"] = None + batch.extra["normal_maps"] = [] + batch.extra["position_maps"] = [] + batch.extra["camera_azims"] = self.CAMERA_AZIMS + batch.extra["camera_elevs"] = self.CAMERA_ELEVS + batch.extra["view_weights"] = self.VIEW_WEIGHTS + batch.extra["renderer"] = None + return batch + + import concurrent.futures + import copy + + # 1. UV unwrap + delight in parallel + batch_for_uv = batch + batch_for_delight = copy.copy(batch) + batch_for_delight.extra = batch.extra.copy() + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + uv_future = executor.submit(self._do_uv_unwrap, batch_for_uv, server_args) + delight_future = executor.submit( + self._do_delight, batch_for_delight, server_args + ) + uv_future.result() + delight_future.result() + + batch.extra["paint_mesh"] = batch_for_uv.extra.get("paint_mesh") + batch.extra["delighted_image"] = batch_for_delight.extra.get("delighted_image") + + # 2. Multi-view rendering + normal_maps, position_maps = self._render_multiview(batch.extra["paint_mesh"]) + batch.extra["normal_maps"] = normal_maps + batch.extra["position_maps"] = position_maps + batch.extra["camera_azims"] = self.CAMERA_AZIMS + batch.extra["camera_elevs"] = self.CAMERA_ELEVS + batch.extra["view_weights"] = self.VIEW_WEIGHTS + batch.extra["renderer"] = self._renderer + + return batch + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) + result.add_check("image_path", batch.image_path, V.not_none) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("paint_mesh", batch.extra.get("paint_mesh"), V.not_none) + result.add_check( + "delighted_image", batch.extra.get("delighted_image"), V.not_none + ) + result.add_check("normal_maps", batch.extra.get("normal_maps"), V.is_list) + result.add_check("position_maps", batch.extra.get("position_maps"), V.is_list) + result.add_check("renderer", batch.extra.get("renderer"), V.not_none) + return result + + +# Stage 2: TexGen (model loading + input prep + denoising + decode) +class Hunyuan3DPaintTexGenStage(PipelineStage): + def __init__( + self, + config: Hunyuan3D2PipelineConfig, + paint_dir: str | None = None, + transformer: Any = None, + scheduler: Any = None, + vae: Any = None, + vae_scale_factor: int = 8, + image_processor: Any = None, + solver: Any = None, + is_turbo: bool = False, + ) -> None: + super().__init__() + self.config = config + self.paint_dir = paint_dir + self.transformer = transformer + self.scheduler = scheduler + self.vae = vae + self.vae_scale_factor = vae_scale_factor + self.image_processor = image_processor + self.solver = solver + self.is_turbo = is_turbo + self._loaded = transformer is not None + + @property + def parallelism_type(self) -> StageParallelismType: + return StageParallelismType.MAIN_RANK_ONLY + + def _load_paint_models(self, server_args: ServerArgs) -> None: + """Load paint models from pre-resolved local path (no network).""" + if self._loaded: + return + if self.paint_dir is None: + logger.warning("No paint model directory resolved, skipping") + self._loaded = True + return + try: + self._do_load_paint(server_args) + logger.info("Paint pipeline loaded successfully") + except Exception as e: + logger.warning("Failed to load paint pipeline: %s", e) + self.vae = None + self.transformer = None + self.scheduler = None + self._loaded = True + + def _do_load_paint(self, server_args: ServerArgs) -> None: + import json + + from diffusers import AutoencoderKL + from diffusers.image_processor import VaeImageProcessor + + from sglang.multimodal_gen.runtime.models.dits.hunyuan3d import ( + UNet2p5DConditionModel, + ) + + local_path = self.paint_dir + logger.info("Loading paint model from %s", local_path) + vae_dir = os.path.join(local_path, "vae") + with open(os.path.join(vae_dir, "config.json"), "r") as f: + vae_config = json.load(f) + vae_config = {k: v for k, v in vae_config.items() if not k.startswith("_")} + self.vae = AutoencoderKL(**vae_config) + st_path = os.path.join(vae_dir, "diffusion_pytorch_model.safetensors") + bin_path = os.path.join(vae_dir, "diffusion_pytorch_model.bin") + if os.path.exists(st_path): + from safetensors.torch import load_file + + state_dict = load_file(st_path) + elif os.path.exists(bin_path): + state_dict = torch.load(bin_path, map_location="cpu", weights_only=True) + else: + raise FileNotFoundError(f"No VAE weights in {vae_dir}") + self.vae.load_state_dict(state_dict) + self.vae = self.vae.to(device=self.device, dtype=torch.float16).eval() + self.transformer = UNet2p5DConditionModel.from_pretrained( + os.path.join(local_path, "unet"), + torch_dtype=torch.float16, + ).to(self.device) + self.is_turbo = bool(getattr(self.config, "paint_turbo_mode", False)) + sched_path = os.path.join(local_path, "scheduler", "scheduler_config.json") + with open(sched_path, "r") as f: + sched_cfg = json.load(f) + if self.is_turbo: + from diffusers import LCMScheduler + + self.scheduler = LCMScheduler.from_config(sched_cfg) + else: + from diffusers import EulerAncestralDiscreteScheduler + + self.scheduler = EulerAncestralDiscreteScheduler.from_config( + sched_cfg, timestep_spacing="trailing" + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.solver = DDIMSolver( + self.scheduler.alphas_cumprod.cpu().numpy(), + timesteps=self.scheduler.config.num_train_timesteps, + ddim_timesteps=30, + ).to(self.device) + if server_args.enable_torch_compile: + compile_mode = os.environ.get( + "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs" + ) + logger.info("Compiling paint transformer with mode: %s", compile_mode) + self.transformer.compile(mode=compile_mode, fullgraph=False, dynamic=None) + + def _convert_pil_list_to_tensor( + self, images: list, device: torch.device + ) -> torch.Tensor: + bg_c = [1.0, 1.0, 1.0] + images_tensor = [] + for batch_imgs in images: + view_imgs = [] + for pil_img in batch_imgs: + if pil_img.mode == "L": + pil_img = pil_img.point( + lambda x: 255 if x > 1 else 0, mode="1" + ).convert("RGB") + img = np.asarray(pil_img, dtype=np.float32) / 255.0 + if img.shape[2] > 3: + alpha = img[:, :, 3:] + img = img[:, :, :3] * alpha + bg_c * (1 - alpha) + img = ( + torch.from_numpy(img) + .permute(2, 0, 1) + .unsqueeze(0) + .contiguous() + .to(device=device, dtype=self.vae.dtype) + ) + view_imgs.append(img) + view_imgs = torch.cat(view_imgs, dim=0) + images_tensor.append(view_imgs.unsqueeze(0)) + return torch.cat(images_tensor, dim=0) + + @torch.no_grad() + def _encode_images(self, images: torch.Tensor) -> torch.Tensor: + batch_size = images.shape[0] + images = rearrange(images, "b n c h w -> (b n) c h w") + dtype = next(self.vae.parameters()).dtype + images = (images - 0.5) * 2.0 + posterior = self.vae.encode(images.to(dtype)).latent_dist + latents = posterior.sample() * self.vae.config.scaling_factor + return rearrange(latents, "(b n) c h w -> b n c h w", b=batch_size) + + @staticmethod + def _compute_camera_index(azim: float, elev: float) -> int: + base_idx = int(((azim // 30) + 9) % 12) + if elev == 0: + base, divisor = 12, 1 + elif elev == 20: + base, divisor = 24, 1 + elif elev == -20: + base, divisor = 0, 1 + elif elev == 90: + base, divisor = 40, 3 + elif elev == -90: + base, divisor = 36, 3 + else: + base, divisor = 12, 1 + return base + (base_idx // divisor) + + def _prepare_denoising_inputs( + self, + batch: Req, + server_args: ServerArgs, + ) -> dict[str, Any]: + import random + + from diffusers.utils.torch_utils import randn_tensor + + device = self.device + normal_maps = batch.extra["normal_maps"] + position_maps = batch.extra["position_maps"] + camera_azims = batch.extra["camera_azims"] + camera_elevs = batch.extra["camera_elevs"] + + num_steps = self.config.paint_num_inference_steps + guidance_scale = self.config.paint_guidance_scale + render_size = self.config.paint_resolution + num_in_batch = len(normal_maps) + + seed = 0 + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + generator = torch.Generator(device=device).manual_seed(seed) + + image = batch.extra["delighted_image"] + if not isinstance(image, list): + image = [image] + image = [to_rgb_image(img) for img in image] + + image_vae = [ + torch.tensor(np.array(img, dtype=np.float32) / 255.0) for img in image + ] + image_vae = [ + iv.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(0) for iv in image_vae + ] + image_vae = torch.cat(image_vae, dim=1).to(device=device, dtype=self.vae.dtype) + ref_latents = self._encode_images(image_vae) + + target_size = render_size + if isinstance(normal_maps, list): + normal_maps = [ + ( + img.resize((target_size, target_size)) + if hasattr(img, "resize") + else img + ) + for img in normal_maps + ] + normal_maps = self._convert_pil_list_to_tensor([normal_maps], device) + if isinstance(position_maps, list): + position_maps = [ + ( + img.resize((target_size, target_size)) + if hasattr(img, "resize") + else img + ) + for img in position_maps + ] + position_maps = self._convert_pil_list_to_tensor([position_maps], device) + + normal_imgs = ( + self._encode_images(normal_maps) if normal_maps is not None else None + ) + position_imgs = ( + self._encode_images(position_maps) if position_maps is not None else None + ) + + camera_info = [ + self._compute_camera_index(azim, elev) + for azim, elev in zip(camera_azims, camera_elevs) + ] + camera_info_gen = torch.tensor([camera_info], device=device, dtype=torch.int64) + camera_info_ref = torch.tensor([[0]], device=device, dtype=torch.int64) + + do_cfg = guidance_scale > 1 and not self.is_turbo + + if self.is_turbo and position_maps is not None: + from sglang.multimodal_gen.runtime.models.dits.hunyuan3d import ( + compute_multi_resolution_discrete_voxel_indice, + compute_multi_resolution_mask, + ) + + position_attn_mask = compute_multi_resolution_mask(position_maps) + position_voxel_indices = compute_multi_resolution_discrete_voxel_indice( + position_maps + ) + else: + position_attn_mask = None + position_voxel_indices = None + + if do_cfg: + negative_ref_latents = torch.zeros_like(ref_latents) + ref_latents = torch.cat([negative_ref_latents, ref_latents]) + ref_scale = torch.as_tensor([0.0, 1.0]).to(ref_latents) + if normal_imgs is not None: + normal_imgs = torch.cat((normal_imgs, normal_imgs)) + if position_imgs is not None: + position_imgs = torch.cat((position_imgs, position_imgs)) + if position_maps is not None: + position_maps = torch.cat((position_maps, position_maps)) + camera_info_gen = torch.cat((camera_info_gen, camera_info_gen)) + camera_info_ref = torch.cat((camera_info_ref, camera_info_ref)) + else: + ref_scale = None + + model_kwargs = { + "ref_latents": ref_latents, + "num_in_batch": num_in_batch, + } + if ref_scale is not None: + model_kwargs["ref_scale"] = ref_scale + if normal_imgs is not None: + model_kwargs["normal_imgs"] = normal_imgs + if position_imgs is not None: + model_kwargs["position_imgs"] = position_imgs + if position_maps is not None: + model_kwargs["position_maps"] = position_maps + model_kwargs["camera_info_gen"] = camera_info_gen + model_kwargs["camera_info_ref"] = camera_info_ref + if position_attn_mask is not None: + model_kwargs["position_attn_mask"] = position_attn_mask + if position_voxel_indices is not None: + model_kwargs["position_voxel_indices"] = position_voxel_indices + + prompt_embeds = self.transformer.learned_text_clip_gen.repeat(1, 1, 1) + negative_prompt_embeds = torch.zeros_like(prompt_embeds) + + if self.is_turbo: + bsz = 3 + index = torch.arange(29, -1, -bsz, device=device).long() + timesteps = self.solver.ddim_timesteps[index] + self.scheduler.set_timesteps(timesteps=timesteps.cpu(), device=device) + timesteps = self.scheduler.timesteps + else: + timesteps, num_steps = retrieve_timesteps( + self.scheduler, num_steps, device, None, None + ) + + num_channels_latents = self.transformer.config.in_channels + latent_shape = ( + num_in_batch, + num_channels_latents, + render_size // self.vae_scale_factor, + render_size // self.vae_scale_factor, + ) + latents = randn_tensor( + latent_shape, generator=generator, device=device, dtype=prompt_embeds.dtype + ) + latents = latents * self.scheduler.init_noise_sigma + + return { + "timesteps": timesteps, + "latents": latents, + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "model_kwargs": model_kwargs, + "num_in_batch": num_in_batch, + "num_inference_steps": num_steps, + "guidance_scale": guidance_scale, + "do_cfg": do_cfg, + "generator": generator, + "num_channels_latents": num_channels_latents, + } + + @torch.no_grad() + def _denoise_loop( + self, + timesteps: torch.Tensor, + latents: torch.Tensor, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor, + model_kwargs: dict[str, Any], + num_in_batch: int, + guidance_scale: float, + do_cfg: bool, + generator: torch.Generator, + num_channels_latents: int, + ) -> torch.Tensor: + import inspect + + if do_cfg: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + + extra_step_kwargs = {} + if "eta" in inspect.signature(self.scheduler.step).parameters: + extra_step_kwargs["eta"] = 0.0 + if "generator" in inspect.signature(self.scheduler.step).parameters: + extra_step_kwargs["generator"] = generator + + for step_idx, t in enumerate(timesteps): + latents = rearrange(latents, "(b n) c h w -> b n c h w", n=num_in_batch) + latent_model_input = torch.cat([latents] * 2) if do_cfg else latents + latent_model_input = rearrange( + latent_model_input, "b n c h w -> (b n) c h w" + ) + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = rearrange( + latent_model_input, "(b n) c h w -> b n c h w", n=num_in_batch + ) + + with set_forward_context( + current_timestep=step_idx, + attn_metadata=None, + ): + noise_pred = self.transformer( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + timestep_cond=None, + cross_attention_kwargs=None, + added_cond_kwargs=None, + return_dict=False, + **model_kwargs, + )[0] + + latents = rearrange(latents, "b n c h w -> (b n) c h w") + + if do_cfg: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + + latents = self.scheduler.step( + noise_pred, + t, + latents[:, :num_channels_latents, :, :], + **extra_step_kwargs, + return_dict=False, + )[0] + + return latents + + @torch.no_grad() + def _decode_latents(self, latents: torch.Tensor) -> list: + image = self.vae.decode( + latents / self.vae.config.scaling_factor, return_dict=False + )[0] + return self.image_processor.postprocess(image, output_type="pil") + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.extra.get("_mesh_failed"): + logger.warning("Mesh generation failed, skipping paint texgen") + batch.extra["multiview_textures"] = [] + return batch + + self._load_paint_models(server_args) + + delighted_image = batch.extra["delighted_image"] + normal_maps = batch.extra["normal_maps"] + + if self.transformer is not None: + try: + prepared = self._prepare_denoising_inputs(batch, server_args) + + latents = self._denoise_loop( + timesteps=prepared["timesteps"], + latents=prepared["latents"], + prompt_embeds=prepared["prompt_embeds"], + negative_prompt_embeds=prepared["negative_prompt_embeds"], + model_kwargs=prepared["model_kwargs"], + num_in_batch=prepared["num_in_batch"], + guidance_scale=prepared["guidance_scale"], + do_cfg=prepared["do_cfg"], + generator=prepared["generator"], + num_channels_latents=prepared["num_channels_latents"], + ) + + multiview_textures = self._decode_latents(latents) + logger.info( + "Paint pipeline generated %d textures", len(multiview_textures) + ) + + except Exception as e: + logger.error(f"Paint pipeline execution failed: {e}") + import traceback + + traceback.print_exc() + render_size = self.config.paint_resolution + multiview_textures = [ + delighted_image.resize((render_size, render_size)) + for _ in range(len(normal_maps)) + ] + else: + logger.warning( + "Paint pipeline not available, using reference image for all views" + ) + render_size = self.config.paint_resolution + multiview_textures = [ + delighted_image.resize((render_size, render_size)) + for _ in range(len(normal_maps)) + ] + + batch.extra["multiview_textures"] = multiview_textures + logger.info(f"Generated {len(multiview_textures)} texture views") + return batch + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + if batch.extra.get("_mesh_failed"): + return VerificationResult() + result = VerificationResult() + result.add_check( + "delighted_image", batch.extra.get("delighted_image"), V.not_none + ) + result.add_check("normal_maps", batch.extra.get("normal_maps"), V.is_list) + result.add_check("position_maps", batch.extra.get("position_maps"), V.is_list) + result.add_check("camera_azims", batch.extra.get("camera_azims"), V.is_list) + result.add_check("camera_elevs", batch.extra.get("camera_elevs"), V.is_list) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check( + "multiview_textures", batch.extra.get("multiview_textures"), V.is_list + ) + return result + + +# Stage 3: Postprocess (texture baking + mesh export) +class Hunyuan3DPaintPostprocessStage(PipelineStage): + """Texture baking from multi-view images and final mesh export.""" + + @property + def parallelism_type(self) -> StageParallelismType: + return StageParallelismType.MAIN_RANK_ONLY + + def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: + super().__init__() + self.config = config + + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + if batch.extra.get("_mesh_failed"): + logger.warning("Mesh generation failed, skipping paint postprocess") + return OutputBatch(output_file_paths=[], metrics=batch.metrics) + + renderer = batch.extra["renderer"] + multiview_textures = batch.extra["multiview_textures"] + camera_elevs = batch.extra["camera_elevs"] + camera_azims = batch.extra["camera_azims"] + view_weights = batch.extra["view_weights"] + + render_size = getattr(self.config, "paint_render_size", 2048) + resized_textures = [] + for tex in multiview_textures: + if hasattr(tex, "resize"): + resized_textures.append(tex.resize((render_size, render_size))) + else: + resized_textures.append(tex) + + try: + texture, mask = renderer.bake_from_multiview( + resized_textures, + camera_elevs, + camera_azims, + view_weights, + method="fast", + ) + + mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype("uint8") + texture = renderer.texture_inpaint(texture, mask_np) + + renderer.set_texture(texture) + textured_mesh = renderer.save_mesh() + logger.info("Texture baking completed") + except Exception as e: + logger.error(f"Texture baking failed: {e}") + textured_mesh = batch.extra["paint_mesh"] + + obj_path = batch.extra["shape_obj_path"] + return_path = batch.extra["shape_return_path"] + + try: + textured_mesh.export(obj_path) + if self.config.paint_save_glb: + glb_path = obj_path[:-4] + ".glb" + textured_mesh.export(glb_path) + return_path = glb_path + self._cleanup_obj_artifacts(obj_path) + except Exception as e: + logger.error(f"Mesh export failed: {e}") + + return OutputBatch(output_file_paths=[return_path], metrics=batch.metrics) + + @staticmethod + def _cleanup_obj_artifacts(obj_path: str) -> None: + """Remove OBJ file and trimesh-generated material artifacts.""" + obj_dir = os.path.dirname(obj_path) or "." + targets = [obj_path] + for f in os.listdir(obj_dir): + if f.endswith(".mtl") or (f.startswith("material") and f.endswith(".png")): + targets.append(os.path.join(obj_dir, f)) + for path in targets: + try: + os.remove(path) + except OSError: + pass + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + if batch.extra.get("_mesh_failed"): + return VerificationResult() + result = VerificationResult() + result.add_check("renderer", batch.extra.get("renderer"), V.not_none) + result.add_check( + "multiview_textures", batch.extra.get("multiview_textures"), V.is_list + ) + result.add_check("camera_elevs", batch.extra.get("camera_elevs"), V.is_list) + result.add_check("camera_azims", batch.extra.get("camera_azims"), V.is_list) + result.add_check("view_weights", batch.extra.get("view_weights"), V.is_list) + return result + + +__all__ = [ + "Hunyuan3DPaintPreprocessStage", + "Hunyuan3DPaintTexGenStage", + "Hunyuan3DPaintPostprocessStage", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py new file mode 100644 index 000000000..6c3be73c1 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Hunyuan3D shape generation stages. + +Four-stage pipeline: BeforeDenoising -> Denoising -> Export -> Save. +""" + +from __future__ import annotations + +import os +from typing import Any + +import numpy as np +import torch + +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( + Hunyuan3D2PipelineConfig, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import ( + TransformerLoader, +) +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + StageValidators as V, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + VerificationResult, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.mesh3d_utils import export_to_trimesh + +logger = init_logger(__name__) + + +def retrieve_timesteps( + scheduler, + num_inference_steps=None, + device=None, + timesteps=None, + sigmas=None, + **kwargs, +): + """Retrieve timesteps from scheduler.""" + import inspect + + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of timesteps or sigmas can be passed.") + + if timesteps is not None: + accepts_timesteps = "timesteps" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys() + ) + if not accepts_timesteps: + raise ValueError( + f"Scheduler {scheduler.__class__} doesn't support custom timesteps." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + + elif sigmas is not None: + accepts_sigmas = "sigmas" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys() + ) + if not accepts_sigmas: + raise ValueError( + f"Scheduler {scheduler.__class__} doesn't support custom sigmas." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + + return timesteps, num_inference_steps + + +def _prepare_shape_image(image_processor, image, mask=None) -> dict: + """Prepare shape image for conditioning.""" + if isinstance(image, torch.Tensor) and isinstance(mask, torch.Tensor): + return {"image": image, "mask": mask} + + if isinstance(image, str) and not os.path.exists(image): + raise FileNotFoundError(f"Couldn't find image at path {image}") + + if not isinstance(image, list): + image = [image] + + outputs = [image_processor(img) for img in image] + cond_input = {k: [] for k in outputs[0].keys()} + for output in outputs: + for key, value in output.items(): + cond_input[key].append(value) + for key, value in cond_input.items(): + if isinstance(value[0], torch.Tensor): + cond_input[key] = torch.cat(value, dim=0) + return cond_input + + +def _move_to_device(payload, device, dtype): + """Recursively move tensors in payload to specified device and dtype.""" + if isinstance(payload, torch.Tensor): + return payload.to(device=device, dtype=dtype) + if isinstance(payload, dict): + return {k: _move_to_device(v, device, dtype) for k, v in payload.items()} + if isinstance(payload, list): + return [_move_to_device(v, device, dtype) for v in payload] + return payload + + +class Hunyuan3DShapeBeforeDenoisingStage(PipelineStage): + """Monolithic pre-processing stage for Hunyuan3D shape generation. + + Consolidates input validation, image preprocessing, conditioning, and + latent/timestep preparation into a single stage. + """ + + def __init__( + self, + image_processor: Any, + conditioner: Any, + vae: Any, + model: Any, + scheduler: Any, + config: Hunyuan3D2PipelineConfig, + ) -> None: + super().__init__() + self.image_processor = image_processor + self.conditioner = conditioner + self.vae = vae + self.model = model + self.scheduler = scheduler + self.config = config + + def _validate_input(self, batch: Req, server_args: ServerArgs) -> None: + if batch.image_path is None: + raise ValueError("Hunyuan3D requires 'image_path' input.") + if isinstance(batch.image_path, list): + if len(batch.image_path) != 1: + raise ValueError("Hunyuan3D only supports a single image input.") + batch.image_path = batch.image_path[0] + if not isinstance(batch.image_path, str): + raise ValueError( + f"Hunyuan3D expects image_path as str, got {type(batch.image_path)}" + ) + if not os.path.exists(batch.image_path): + raise FileNotFoundError(f"Image path not found: {batch.image_path}") + if batch.num_outputs_per_prompt != 1: + raise ValueError("Hunyuan3D only supports num_outputs_per_prompt=1.") + + def _prepare_latents(self, batch_size, dtype, device, generator): + from diffusers.utils.torch_utils import randn_tensor + + shape = (batch_size, *self.vae.latent_shape) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents * getattr(self.scheduler, "init_noise_sigma", 1.0) + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + # 1. Input validation + self._validate_input(batch, server_args) + + # 2. Image preprocessing + cond_inputs = _prepare_shape_image(self.image_processor, batch.image_path) + image = cond_inputs.pop("image") + + device = self.device + dtype = next(self.model.parameters()).dtype + image = _move_to_device(image, device, dtype) + cond_inputs = _move_to_device(cond_inputs, device, dtype) + + # 3. Conditioning with CFG + do_cfg = batch.guidance_scale >= 0 and not ( + hasattr(self.model, "guidance_embed") and self.model.guidance_embed is True + ) + + cond = self.conditioner(image=image, **cond_inputs) + if do_cfg: + un_cond = self.conditioner.unconditional_embedding( + image.shape[0], **cond_inputs + ) + + def cat_recursive(a, b): + if isinstance(a, torch.Tensor): + return torch.cat([a, b], dim=0).to(dtype) + out = {} + for key in a.keys(): + out[key] = cat_recursive(a[key], b[key]) + return out + + cond = cat_recursive(cond, un_cond) + + # 4. Latent and timestep preparation + batch_size = image.shape[0] + sigmas = np.linspace(0, 1, batch.num_inference_steps) + timesteps, _ = retrieve_timesteps( + self.scheduler, + batch.num_inference_steps, + device, + sigmas=sigmas, + ) + + generator = batch.generator + if generator is None and batch.seed is not None: + generator = torch.Generator(device=device).manual_seed(batch.seed) + + latents = self._prepare_latents(batch_size, dtype, device, generator) + + guidance = None + if hasattr(self.model, "guidance_embed") and self.model.guidance_embed is True: + guidance = torch.tensor( + [batch.guidance_scale] * batch_size, device=device, dtype=dtype + ) + + # 5. Populate batch + batch.prompt_embeds = [cond] + batch.do_classifier_free_guidance = do_cfg + batch.timesteps = timesteps + batch.latents = latents + batch.extra["shape_guidance"] = guidance + batch.extra["shape_image"] = image + return batch + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("image_path", batch.image_path, V.not_none) + result.add_check( + "num_inference_steps", batch.num_inference_steps, V.positive_int + ) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) + result.add_check("latents", batch.latents, V.is_tensor) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + return result + + +class Hunyuan3DShapeDenoisingStage(DenoisingStage): + """Denoising stage for Hunyuan3D shape generation.""" + + def __init__(self, transformer: Any, scheduler: Any, **kwargs) -> None: + super().__init__(transformer=transformer, scheduler=scheduler, **kwargs) + + def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): + """Prepare Hunyuan3D-specific variables for the base denoising loop.""" + assert self.transformer is not None + pipeline = self.pipeline() if self.pipeline else None + cache_dit_num_inference_steps = batch.extra.get( + "cache_dit_num_inference_steps", batch.num_inference_steps + ) + if not server_args.model_loaded["transformer"]: + loader = TransformerLoader() + self.transformer = loader.load( + server_args.model_paths["transformer"], server_args, "transformer" + ) + self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) + self._maybe_enable_torch_compile(self.transformer) + if pipeline: + pipeline.add_module("transformer", self.transformer) + server_args.model_loaded["transformer"] = True + else: + self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) + + timesteps = batch.timesteps + if timesteps is None: + raise ValueError("Timesteps must be provided") + + latents = batch.latents + if latents is None: + raise ValueError("Latents must be provided") + + cond = batch.prompt_embeds[0] if batch.prompt_embeds else None + if cond is None: + raise ValueError("Conditioning (prompt_embeds) must be provided") + + if batch.raw_latent_shape is None: + batch.raw_latent_shape = latents.shape + + guidance = batch.extra.get("shape_guidance") + num_inference_steps = batch.num_inference_steps + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + {"generator": batch.generator, "eta": batch.eta}, + ) + + target_dtype = next(self.transformer.parameters()).dtype + autocast_enabled = False + + pos_cond_kwargs = {"encoder_hidden_states": cond} + neg_cond_kwargs = {} + + return { + "extra_step_kwargs": extra_step_kwargs, + "target_dtype": target_dtype, + "autocast_enabled": autocast_enabled, + "timesteps": timesteps, + "num_inference_steps": num_inference_steps, + "num_warmup_steps": num_warmup_steps, + "image_kwargs": {}, + "pos_cond_kwargs": pos_cond_kwargs, + "neg_cond_kwargs": neg_cond_kwargs, + "latents": latents, + "prompt_embeds": batch.prompt_embeds, + "neg_prompt_embeds": None, + "boundary_timestep": None, + "z": None, + "reserved_frames_mask": None, + "seq_len": None, + "guidance": guidance, + } + + def _predict_noise( + self, + current_model, + latent_model_input, + timestep, + target_dtype, + guidance: torch.Tensor, + **kwargs, + ): + """Hunyuan3D-specific noise prediction with normalized timestep.""" + cond = kwargs.get("encoder_hidden_states") + timestep_norm = timestep / self.scheduler.config.num_train_timesteps + return current_model(latent_model_input, timestep_norm, cond, guidance=guidance) + + def _predict_noise_with_cfg( + self, + current_model, + latent_model_input: torch.Tensor, + timestep, + batch: Req, + timestep_index: int, + attn_metadata, + target_dtype, + current_guidance_scale, + image_kwargs: dict[str, Any], + pos_cond_kwargs: dict[str, Any], + neg_cond_kwargs: dict[str, Any], + server_args, + guidance, + latents, + ): + """Hunyuan3D-specific CFG: concat latents, single forward, then split.""" + cond = pos_cond_kwargs.get("encoder_hidden_states") + do_cfg = batch.do_classifier_free_guidance + + if do_cfg: + latent_input = torch.cat([latent_model_input] * 2) + else: + latent_input = latent_model_input + + timestep_expanded = timestep.expand(latent_input.shape[0]).to(latents.dtype) + + with set_forward_context( + current_timestep=timestep_index, + attn_metadata=attn_metadata, + forward_batch=batch, + ): + noise_pred = self._predict_noise( + current_model=current_model, + latent_model_input=latent_input, + timestep=timestep_expanded, + target_dtype=target_dtype, + guidance=guidance, + encoder_hidden_states=cond, + ) + + if do_cfg: + noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + current_guidance_scale * ( + noise_pred_cond - noise_pred_uncond + ) + + return noise_pred + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) + result.add_check("latents", batch.latents, V.is_tensor) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check( + "num_inference_steps", batch.num_inference_steps, V.positive_int + ) + result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, V.is_tensor) + return result + + +class Hunyuan3DShapeExportStage(PipelineStage): + """VAE decoding and mesh extraction stage.""" + + def __init__(self, vae: Any, config: Hunyuan3D2PipelineConfig) -> None: + super().__init__() + self.vae = vae + self.config = config + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if self.config.shape_mc_algo is not None: + try: + from sglang.multimodal_gen.runtime.models.vaes.hunyuan3d_vae import ( + SurfaceExtractors, + ) + + self.vae.surface_extractor = SurfaceExtractors[ + self.config.shape_mc_algo + ]() + except ImportError: + logger.warning( + f"Could not load SurfaceExtractors for mc_algo={self.config.shape_mc_algo}" + ) + + latents = batch.latents + + if self.config.shape_output_type != "latent": + latents = 1.0 / self.vae.scale_factor * latents + latents = self.vae(latents) + + outputs = self.vae.latents2mesh( + latents, + bounds=self.config.shape_box_v, + mc_level=self.config.shape_mc_level, + num_chunks=self.config.shape_num_chunks, + octree_resolution=self.config.shape_octree_resolution, + mc_algo=self.config.shape_mc_algo, + enable_pbar=False, + ) + else: + outputs = latents + + if self.config.shape_output_type == "trimesh": + outputs = export_to_trimesh(outputs) + + batch.extra["shape_meshes"] = outputs + return batch + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, V.is_tensor) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) + return result + + +class Hunyuan3DShapeSaveStage(PipelineStage): + """Mesh file export and output decision stage.""" + + def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: + super().__init__() + self.config = config + + def _get_output_paths(self, batch: Req) -> tuple[str, str]: + output_path = batch.output_file_path() or os.path.join( + batch.output_path, "output.obj" + ) + if output_path.endswith(".glb"): + obj_path = output_path[:-4] + ".obj" + return obj_path, output_path + if output_path.endswith(".obj"): + return output_path, output_path + return output_path + ".obj", output_path + ".obj" + + def forward(self, batch: Req, server_args: ServerArgs) -> Req | OutputBatch: + mesh_outputs = batch.extra["shape_meshes"] + mesh = mesh_outputs[0] if isinstance(mesh_outputs, list) else mesh_outputs + if isinstance(mesh, list): + mesh = mesh[0] + + if mesh is None: + if batch.is_warmup: + logger.info( + "Skipping mesh export during warmup " + "(surface extraction returned None)" + ) + batch.extra["_mesh_failed"] = True + if self.config.paint_enable: + return batch + return OutputBatch(output_file_paths=[], metrics=batch.metrics) + raise RuntimeError( + "Mesh generation failed: surface extraction returned None. " + "The surface level may be outside the volume data range." + ) + + obj_path, return_path = self._get_output_paths(batch) + output_dir = os.path.dirname(obj_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + mesh.export(obj_path) + + batch.extra["shape_obj_path"] = obj_path + batch.extra["shape_return_path"] = return_path + + if self.config.paint_enable: + return batch + + if return_path.endswith(".glb"): + return_path = obj_path + return OutputBatch(output_file_paths=[return_path], timings=batch.timings) + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) + return result + + +__all__ = [ + "retrieve_timesteps", + "Hunyuan3DShapeBeforeDenoisingStage", + "Hunyuan3DShapeDenoisingStage", + "Hunyuan3DShapeExportStage", + "Hunyuan3DShapeSaveStage", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 971a9b0a0..aa0b686b8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -241,8 +241,21 @@ class InputValidationStage(PipelineStage): self._generate_seeds(batch, server_args) - # Ensure prompt is properly formatted - if batch.prompt is None and batch.prompt_embeds is None: + if ( + server_args.pipeline_config.task_type == ModelTaskType.I2M + and batch.num_inference_steps is None + and hasattr(server_args.pipeline_config, "shape_num_inference_steps") + ): + batch.num_inference_steps = ( + server_args.pipeline_config.shape_num_inference_steps + ) + + # Ensure prompt is properly formatted (I2M can be image-only) + if ( + server_args.pipeline_config.task_type != ModelTaskType.I2M + and batch.prompt is None + and batch.prompt_embeds is None + ): raise ValueError("Either `prompt` or `prompt_embeds` must be provided") # Ensure negative prompt is properly formatted if using classifier-free guidance @@ -299,9 +312,10 @@ class InputValidationStage(PipelineStage): ) batch.original_condition_image_size = image.size - self.preprocess_condition_image( - batch, server_args, condition_image_width, condition_image_height - ) + if server_args.pipeline_config.task_type != ModelTaskType.I2M: + self.preprocess_condition_image( + batch, server_args, condition_image_width, condition_image_height + ) # if height or width is not specified at this point, set default to 720p default_height = 720 @@ -323,16 +337,24 @@ class InputValidationStage(PipelineStage): result.add_check( "num_videos_per_prompt", batch.num_outputs_per_prompt, V.positive_int ) - result.add_check( - "prompt_or_embeds", - None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), - ) + if server_args.pipeline_config.task_type != ModelTaskType.I2M: + result.add_check( + "prompt_or_embeds", + None, + lambda _: V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds), + ) - result.add_check( - "num_inference_steps", batch.num_inference_steps, V.positive_int - ) + if server_args.pipeline_config.task_type != ModelTaskType.I2M: + result.add_check( + "num_inference_steps", batch.num_inference_steps, V.positive_int + ) + else: + result.add_check( + "num_inference_steps", + batch.num_inference_steps, + lambda x: x is None or V.positive_int(x), + ) result.add_check( "guidance_scale", batch.guidance_scale, diff --git a/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py b/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py new file mode 100644 index 000000000..8d17eaf9d --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py @@ -0,0 +1,1114 @@ +"""Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2""" + +from __future__ import annotations + +import math +from typing import Any, List, Optional, Tuple, Union + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +import trimesh +from einops import rearrange, repeat +from PIL import Image + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Import C++ mesh processor extension +from sglang.multimodal_gen.csrc.render.mesh_processor import meshVerticeInpaint + + +def transform_pos( + mtx: Union[np.ndarray, torch.Tensor], + pos: torch.Tensor, + keepdim: bool = False, +) -> torch.Tensor: + """Transform positions by a matrix.""" + t_mtx = torch.from_numpy(mtx).to(pos.device) if isinstance(mtx, np.ndarray) else mtx + + if pos.shape[-1] == 3: + posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).to(pos.device)], axis=1) + else: + posw = pos + + if keepdim: + return torch.matmul(posw, t_mtx.t())[...] + else: + return torch.matmul(posw, t_mtx.t())[None, ...] + + +def get_mv_matrix( + elev: float, + azim: float, + camera_distance: float, + center: Optional[np.ndarray] = None, +) -> np.ndarray: + """Compute model-view matrix from camera parameters.""" + elev = -elev + azim += 90 + + elev_rad = math.radians(elev) + azim_rad = math.radians(azim) + + camera_position = np.array( + [ + camera_distance * math.cos(elev_rad) * math.cos(azim_rad), + camera_distance * math.cos(elev_rad) * math.sin(azim_rad), + camera_distance * math.sin(elev_rad), + ] + ) + + if center is None: + center = np.array([0, 0, 0]) + else: + center = np.array(center) + + lookat = center - camera_position + lookat = lookat / np.linalg.norm(lookat) + + up = np.array([0, 0, 1.0]) + right = np.cross(lookat, up) + right = right / np.linalg.norm(right) + up = np.cross(right, lookat) + up = up / np.linalg.norm(up) + + c2w = np.concatenate( + [np.stack([right, up, -lookat], axis=-1), camera_position[:, None]], axis=-1 + ) + + w2c = np.zeros((4, 4)) + w2c[:3, :3] = np.transpose(c2w[:3, :3], (1, 0)) + w2c[:3, 3:] = -np.matmul(np.transpose(c2w[:3, :3], (1, 0)), c2w[:3, 3:]) + w2c[3, 3] = 1.0 + + return w2c.astype(np.float32) + + +def get_orthographic_projection_matrix( + left: float = -1, + right: float = 1, + bottom: float = -1, + top: float = 1, + near: float = 0, + far: float = 2, +) -> np.ndarray: + """Compute orthographic projection matrix.""" + ortho_matrix = np.eye(4, dtype=np.float32) + ortho_matrix[0, 0] = 2 / (right - left) + ortho_matrix[1, 1] = 2 / (top - bottom) + ortho_matrix[2, 2] = -2 / (far - near) + ortho_matrix[0, 3] = -(right + left) / (right - left) + ortho_matrix[1, 3] = -(top + bottom) / (top - bottom) + ortho_matrix[2, 3] = -(far + near) / (far - near) + return ortho_matrix + + +def get_perspective_projection_matrix( + fovy: float, + aspect_wh: float, + near: float, + far: float, +) -> np.ndarray: + """Compute perspective projection matrix.""" + fovy_rad = math.radians(fovy) + return np.array( + [ + [1.0 / (math.tan(fovy_rad / 2.0) * aspect_wh), 0, 0, 0], + [0, 1.0 / math.tan(fovy_rad / 2.0), 0, 0], + [0, 0, -(far + near) / (far - near), -2.0 * far * near / (far - near)], + [0, 0, -1, 0], + ] + ).astype(np.float32) + + +def export_to_trimesh(mesh_output: Any) -> Any: + """Convert mesh output to trimesh format.""" + if isinstance(mesh_output, list): + outputs = [] + for mesh in mesh_output: + if mesh is None: + outputs.append(None) + else: + # Reverse face winding + mesh.mesh_f = mesh.mesh_f[:, ::-1] + mesh_obj = trimesh.Trimesh(mesh.mesh_v, mesh.mesh_f) + outputs.append(mesh_obj) + return outputs + else: + mesh_output.mesh_f = mesh_output.mesh_f[:, ::-1] + return trimesh.Trimesh(mesh_output.mesh_v, mesh_output.mesh_f) + + +def mesh_uv_wrap(mesh: Any) -> Any: + """Apply UV unwrapping to mesh. In-place like native Hunyuan3D-2 for same layout.""" + try: + import xatlas + except ImportError: + logger.warning("xatlas not available, skipping UV unwrap") + return mesh + + if isinstance(mesh, trimesh.Scene): + mesh = mesh.dump(concatenate=True) + + if len(mesh.faces) > 500000000: + raise ValueError( + "The mesh has more than 500,000,000 faces, which is not supported." + ) + + vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces) + + mesh.vertices = mesh.vertices[vmapping] + mesh.faces = indices + if not hasattr(mesh.visual, "uv"): + mesh.visual = trimesh.visual.TextureVisuals( + uv=uvs, material=trimesh.visual.material.SimpleMaterial() + ) + else: + mesh.visual.uv = uvs + + return mesh + + +def stride_from_shape(shape: Tuple[int, ...]) -> List[int]: + """Compute stride from shape for scatter operations.""" + stride = [1] + for x in reversed(shape[1:]): + stride.append(stride[-1] * x) + return list(reversed(stride)) + + +def scatter_add_nd_with_count( + input: torch.Tensor, + count: torch.Tensor, + indices: torch.Tensor, + values: torch.Tensor, + weights: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Scatter add with counting for texture baking.""" + D = indices.shape[-1] + C = input.shape[-1] + size = input.shape[:-1] + stride = stride_from_shape(size) + + assert len(size) == D + + input = input.view(-1, C) + count = count.view(-1, 1) + + flatten_indices = ( + indices * torch.tensor(stride, dtype=torch.long, device=indices.device) + ).sum(-1) + + if weights is None: + weights = torch.ones_like(values[..., :1]) + + input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values) + count.scatter_add_(0, flatten_indices.unsqueeze(1), weights) + + return input.view(*size, C), count.view(*size, 1) + + +def linear_grid_put_2d( + H: int, + W: int, + coords: torch.Tensor, + values: torch.Tensor, + return_count: bool = False, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Put values into a 2D grid using linear interpolation.""" + C = values.shape[-1] + + indices = coords * torch.tensor( + [H - 1, W - 1], dtype=torch.float32, device=coords.device + ) + indices_00 = indices.floor().long() + indices_00[:, 0].clamp_(0, H - 2) + indices_00[:, 1].clamp_(0, W - 2) + indices_01 = indices_00 + torch.tensor( + [0, 1], dtype=torch.long, device=indices.device + ) + indices_10 = indices_00 + torch.tensor( + [1, 0], dtype=torch.long, device=indices.device + ) + indices_11 = indices_00 + torch.tensor( + [1, 1], dtype=torch.long, device=indices.device + ) + + h = indices[..., 0] - indices_00[..., 0].float() + w = indices[..., 1] - indices_00[..., 1].float() + w_00 = (1 - h) * (1 - w) + w_01 = (1 - h) * w + w_10 = h * (1 - w) + w_11 = h * w + + result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype) + count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype) + weights = torch.ones_like(values[..., :1]) + + result, count = scatter_add_nd_with_count( + result, + count, + indices_00, + values * w_00.unsqueeze(1), + weights * w_00.unsqueeze(1), + ) + result, count = scatter_add_nd_with_count( + result, + count, + indices_01, + values * w_01.unsqueeze(1), + weights * w_01.unsqueeze(1), + ) + result, count = scatter_add_nd_with_count( + result, + count, + indices_10, + values * w_10.unsqueeze(1), + weights * w_10.unsqueeze(1), + ) + result, count = scatter_add_nd_with_count( + result, + count, + indices_11, + values * w_11.unsqueeze(1), + weights * w_11.unsqueeze(1), + ) + + if return_count: + return result, count + + mask = count.squeeze(-1) > 0 + result[mask] = result[mask] / count[mask].repeat(1, C) + + return result + + +class MeshRender: + """Mesh renderer using CUDA rasterization for texture generation.""" + + def __init__( + self, + camera_distance: float = 1.45, + camera_type: str = "orth", + default_resolution: int = 1024, + texture_size: int = 1024, + bake_mode: str = "linear", + device: str = "cuda", + ): + """Initialize the mesh renderer.""" + self.device = device + + self.set_default_render_resolution(default_resolution) + self.set_default_texture_resolution(texture_size) + + self.camera_distance = camera_distance + self.camera_type = camera_type + self.bake_angle_thres = 75 + self.bake_unreliable_kernel_size = int( + (2 / 512) * max(self.default_resolution[0], self.default_resolution[1]) + ) + self.bake_mode = bake_mode + + # Set up camera projection matrix + if camera_type == "orth": + self.ortho_scale = 1.2 + self.camera_proj_mat = get_orthographic_projection_matrix( + left=-self.ortho_scale * 0.5, + right=self.ortho_scale * 0.5, + bottom=-self.ortho_scale * 0.5, + top=self.ortho_scale * 0.5, + near=0.1, + far=100, + ) + elif camera_type == "perspective": + self.camera_proj_mat = get_perspective_projection_matrix( + 49.13, + self.default_resolution[1] / self.default_resolution[0], + 0.01, + 100.0, + ) + else: + raise ValueError(f"Unknown camera type: {camera_type}") + + # Mesh data + self.vtx_pos = None + self.pos_idx = None + self.vtx_uv = None + self.uv_idx = None + self.tex = None + self.mesh_copy = None + self.scale_factor = 1.0 + + def set_default_render_resolution( + self, default_resolution: Union[int, Tuple[int, int]] + ): + """Set default rendering resolution.""" + if isinstance(default_resolution, int): + default_resolution = (default_resolution, default_resolution) + self.default_resolution = default_resolution + + def set_default_texture_resolution(self, texture_size: Union[int, Tuple[int, int]]): + """Set default texture resolution.""" + if isinstance(texture_size, int): + texture_size = (texture_size, texture_size) + self.texture_size = texture_size + + def _rasterize( + self, + pos_clip: torch.Tensor, + tri: torch.Tensor, + resolution: Tuple[int, int], + ) -> torch.Tensor: + """Rasterize using CUDA rasterizer.""" + from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import rasterize + + if pos_clip.dim() == 2: + pos_clip = pos_clip.unsqueeze(0) + + findices, barycentric = rasterize(pos_clip, tri, resolution) + rast_out = torch.cat((barycentric, findices.unsqueeze(-1).float()), dim=-1) + rast_out = rast_out.unsqueeze(0) + return rast_out + + def _interpolate( + self, + attr: torch.Tensor, + rast_out: torch.Tensor, + tri: torch.Tensor, + ) -> torch.Tensor: + """Interpolate vertex attributes.""" + from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import interpolate + + barycentric = rast_out[0, ..., :-1] + findices = rast_out[0, ..., -1].int() + + if attr.dim() == 2: + attr = attr.unsqueeze(0) + + result = interpolate(attr, findices, barycentric, tri) + return result + + def load_mesh( + self, + mesh: Union[trimesh.Trimesh, trimesh.Scene], + scale_factor: float = 1.15, + auto_center: bool = True, + ): + """Load a mesh for rendering.""" + if isinstance(mesh, trimesh.Scene): + mesh = mesh.dump(concatenate=True) + + self.mesh_copy = mesh.copy() + + vtx_pos = mesh.vertices.astype(np.float32) + pos_idx = mesh.faces.astype(np.int32) + + # Get UV coordinates if available + if hasattr(mesh.visual, "uv") and mesh.visual.uv is not None: + vtx_uv = mesh.visual.uv.astype(np.float32) + uv_idx = pos_idx.copy() + else: + vtx_uv = None + uv_idx = None + + self.vtx_pos = torch.from_numpy(vtx_pos).to(self.device).float() + self.pos_idx = torch.from_numpy(pos_idx).to(self.device).to(torch.int32) + + if vtx_uv is not None and uv_idx is not None: + self.vtx_uv = torch.from_numpy(vtx_uv).to(self.device).float() + self.uv_idx = torch.from_numpy(uv_idx).to(self.device).to(torch.int32) + else: + self.vtx_uv = None + self.uv_idx = None + + # Coordinate transformation (Y-up to Z-up) + self.vtx_pos[:, [0, 1]] = -self.vtx_pos[:, [0, 1]] + self.vtx_pos[:, [1, 2]] = self.vtx_pos[:, [2, 1]] + if self.vtx_uv is not None: + self.vtx_uv[:, 1] = 1.0 - self.vtx_uv[:, 1] + + if auto_center: + max_bb = (self.vtx_pos - 0).max(0)[0] + min_bb = (self.vtx_pos - 0).min(0)[0] + center = (max_bb + min_bb) / 2 + scale = torch.norm(self.vtx_pos - center, dim=1).max() * 2.0 + self.vtx_pos = (self.vtx_pos - center) * (scale_factor / float(scale)) + self.scale_factor = scale_factor + + def save_mesh(self) -> trimesh.Trimesh: + """Save mesh with current texture, reusing the original mesh object.""" + texture_data = self.get_texture() + texture_img = Image.fromarray((texture_data * 255).astype(np.uint8)) + + material = trimesh.visual.material.SimpleMaterial( + image=texture_img, diffuse=(255, 255, 255) + ) + self.mesh_copy.visual = trimesh.visual.TextureVisuals( + uv=self.mesh_copy.visual.uv, image=texture_img, material=material + ) + return self.mesh_copy + + def get_mesh(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Get mesh data with inverse coordinate transformation.""" + vtx_pos = self.vtx_pos.cpu().numpy().copy() + pos_idx = self.pos_idx.cpu().numpy() + + # Inverse coordinate transformation + vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]] + vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]] + + if self.vtx_uv is not None: + vtx_uv = self.vtx_uv.cpu().numpy().copy() + vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1] + uv_idx = self.uv_idx.cpu().numpy() + else: + vtx_uv = None + uv_idx = None + + return vtx_pos, pos_idx, vtx_uv, uv_idx + + def set_texture(self, tex: Union[np.ndarray, torch.Tensor, Image.Image]): + """Set texture for the mesh.""" + if isinstance(tex, np.ndarray): + if tex.max() <= 1.0: + tex = (tex * 255).astype(np.uint8) + tex = Image.fromarray(tex.astype(np.uint8)) + elif isinstance(tex, torch.Tensor): + tex_np = tex.cpu().numpy() + if tex_np.max() <= 1.0: + tex_np = (tex_np * 255).astype(np.uint8) + tex = Image.fromarray(tex_np.astype(np.uint8)) + + tex = tex.resize(self.texture_size).convert("RGB") + tex = np.array(tex) / 255.0 + self.tex = torch.from_numpy(tex).to(self.device).float() + + def get_texture(self) -> np.ndarray: + """Get current texture as numpy array.""" + if self.tex is None: + return np.ones((*self.texture_size, 3), dtype=np.float32) + return self.tex.cpu().numpy() + + def _get_pos_from_mvp( + self, + elev: float, + azim: float, + camera_distance: Optional[float] = None, + center: Optional[np.ndarray] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Get camera-space and clip-space positions.""" + proj = self.camera_proj_mat + r_mv = get_mv_matrix( + elev=elev, + azim=azim, + camera_distance=( + self.camera_distance if camera_distance is None else camera_distance + ), + center=center, + ) + + pos_camera = transform_pos(r_mv, self.vtx_pos, keepdim=True) + pos_clip = transform_pos(proj, pos_camera) + + return pos_camera, pos_clip + + def render_normal( + self, + elev: float, + azim: float, + camera_distance: Optional[float] = None, + center: Optional[np.ndarray] = None, + resolution: Optional[Tuple[int, int]] = None, + bg_color: List[float] = [1, 1, 1], + use_abs_coor: bool = False, + normalize_rgb: bool = True, + return_type: str = "th", + ) -> Union[torch.Tensor, np.ndarray, Image.Image]: + """Render normal map from a viewpoint.""" + pos_camera, pos_clip = self._get_pos_from_mvp( + elev, azim, camera_distance, center + ) + + if resolution is None: + resolution = self.default_resolution + if isinstance(resolution, (int, float)): + resolution = (int(resolution), int(resolution)) + + rast_out = self._rasterize(pos_clip, self.pos_idx, resolution) + + # Compute face normals + if use_abs_coor: + mesh_triangles = self.vtx_pos[self.pos_idx[:, :3].long(), :] + else: + pos_camera_3d = pos_camera[:, :3] / pos_camera[:, 3:4] + mesh_triangles = pos_camera_3d[self.pos_idx[:, :3].long(), :] + + face_normals = F.normalize( + torch.cross( + mesh_triangles[:, 1, :] - mesh_triangles[:, 0, :], + mesh_triangles[:, 2, :] - mesh_triangles[:, 0, :], + dim=-1, + ), + dim=-1, + ) + + # Compute vertex normals + vertex_normals = trimesh.geometry.mean_vertex_normals( + vertex_count=self.vtx_pos.shape[0], + faces=self.pos_idx.cpu().numpy(), + face_normals=face_normals.cpu().numpy(), + ) + vertex_normals = ( + torch.from_numpy(vertex_normals).float().to(self.device).contiguous() + ) + + # Interpolate normals + normal = self._interpolate(vertex_normals[None, ...], rast_out, self.pos_idx) + + # Apply visibility mask + visible_mask = torch.clamp(rast_out[..., -1:], 0, 1) + bg_tensor = torch.tensor(bg_color, dtype=torch.float32, device=self.device) + normal = normal * visible_mask + bg_tensor * (1 - visible_mask) + + if normalize_rgb: + normal = (normal + 1) * 0.5 + + image = normal[0, ...] + + if return_type == "np": + image = image.cpu().numpy() + elif return_type == "pl": + image = image.cpu().numpy() * 255 + image = Image.fromarray(image.astype(np.uint8)) + + return image + + def render_position( + self, + elev: float, + azim: float, + camera_distance: Optional[float] = None, + center: Optional[np.ndarray] = None, + resolution: Optional[Tuple[int, int]] = None, + bg_color: List[float] = [1, 1, 1], + return_type: str = "th", + ) -> Union[torch.Tensor, np.ndarray, Image.Image]: + """Render position map from a viewpoint.""" + pos_camera, pos_clip = self._get_pos_from_mvp( + elev, azim, camera_distance, center + ) + + if resolution is None: + resolution = self.default_resolution + if isinstance(resolution, (int, float)): + resolution = (int(resolution), int(resolution)) + + rast_out = self._rasterize(pos_clip, self.pos_idx, resolution) + + # Position colors (normalized vertex positions) + tex_position = 0.5 - self.vtx_pos[:, :3] / self.scale_factor + tex_position = tex_position.contiguous() + + # Interpolate positions + position = self._interpolate(tex_position[None, ...], rast_out, self.pos_idx) + + # Apply visibility mask + visible_mask = torch.clamp(rast_out[..., -1:], 0, 1) + bg_tensor = torch.tensor(bg_color, dtype=torch.float32, device=self.device) + position = position * visible_mask + bg_tensor * (1 - visible_mask) + + image = position[0, ...] + + if return_type == "np": + image = image.cpu().numpy() + elif return_type == "pl": + image = image.cpu().numpy() * 255 + image = Image.fromarray(image.astype(np.uint8)) + + return image + + def render_normal_multiview( + self, + camera_elevs: List[float], + camera_azims: List[float], + use_abs_coor: bool = True, + ) -> List[Image.Image]: + """Render normal maps from multiple viewpoints.""" + normal_maps = [] + for elev, azim in zip(camera_elevs, camera_azims): + normal_map = self.render_normal( + elev, azim, use_abs_coor=use_abs_coor, return_type="pl" + ) + normal_maps.append(normal_map) + return normal_maps + + def render_position_multiview( + self, + camera_elevs: List[float], + camera_azims: List[float], + ) -> List[Image.Image]: + """Render position maps from multiple viewpoints.""" + position_maps = [] + for elev, azim in zip(camera_elevs, camera_azims): + position_map = self.render_position(elev, azim, return_type="pl") + position_maps.append(position_map) + return position_maps + + def _render_sketch_from_depth(self, depth_image: torch.Tensor) -> torch.Tensor: + """Render sketch from depth using edge detection.""" + depth_image_np = depth_image.cpu().numpy() + depth_image_np = (depth_image_np * 255).astype(np.uint8) + depth_edges = cv2.Canny(depth_image_np, 30, 80) + sketch_image = ( + torch.from_numpy(depth_edges).to(depth_image.device).float() / 255.0 + ) + sketch_image = sketch_image.unsqueeze(-1) + return sketch_image + + def back_project( + self, + image: Union[Image.Image, np.ndarray, torch.Tensor], + elev: float, + azim: float, + camera_distance: Optional[float] = None, + center: Optional[np.ndarray] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Back-project an image onto mesh UV space.""" + if isinstance(image, Image.Image): + image = torch.tensor(np.array(image) / 255.0) + elif isinstance(image, np.ndarray): + image = torch.tensor(image) + if image.dim() == 2: + image = image.unsqueeze(-1) + image = image.float().to(self.device) + resolution = image.shape[:2] + channel = image.shape[-1] + + pos_camera, pos_clip = self._get_pos_from_mvp( + elev, azim, camera_distance, center + ) + + rast_out = self._rasterize(pos_clip, self.pos_idx, resolution) + visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0, ...] + + # Compute vertex normals for angle-based weighting + pos_camera_3d = pos_camera[:, :3] / pos_camera[:, 3:4] + v0 = pos_camera_3d[self.pos_idx[:, 0].long(), :] + v1 = pos_camera_3d[self.pos_idx[:, 1].long(), :] + v2 = pos_camera_3d[self.pos_idx[:, 2].long(), :] + face_normals = F.normalize(torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1) + + vertex_normals = trimesh.geometry.mean_vertex_normals( + vertex_count=self.vtx_pos.shape[0], + faces=self.pos_idx.cpu().numpy(), + face_normals=face_normals.cpu().numpy(), + ) + vertex_normals = ( + torch.from_numpy(vertex_normals).float().to(self.device).contiguous() + ) + + # Interpolate normals and UVs + normal = self._interpolate(vertex_normals[None, ...], rast_out, self.pos_idx) + normal = normal[0, ...] + + if self.vtx_uv is not None: + uv = self._interpolate(self.vtx_uv[None, ...], rast_out, self.uv_idx) + else: + # No UV coordinates + texture = torch.zeros( + self.texture_size[1], self.texture_size[0], channel, device=self.device + ) + cos_map = torch.zeros( + self.texture_size[1], self.texture_size[0], 1, device=self.device + ) + boundary_map = torch.zeros_like(cos_map) + return texture, cos_map, boundary_map + + # Compute depth for sketch + tex_depth = pos_camera_3d[:, 2].reshape(1, -1, 1).contiguous() + depth = self._interpolate(tex_depth, rast_out, self.pos_idx)[0, ...] + depth_masked = depth[visible_mask > 0] + if depth_masked.numel() > 0: + depth_max, depth_min = depth_masked.max(), depth_masked.min() + depth_normalized = (depth - depth_min) / (depth_max - depth_min + 1e-8) + else: + depth_normalized = depth + depth_image = depth_normalized * visible_mask + + sketch_image = self._render_sketch_from_depth(depth_image) + + # Cosine weighting + lookat = torch.tensor([[0, 0, -1]], device=self.device) + cos_image = torch.nn.functional.cosine_similarity(lookat, normal.view(-1, 3)) + cos_image = cos_image.view(normal.shape[0], normal.shape[1], 1) + + cos_thres = np.cos(self.bake_angle_thres / 180 * np.pi) + cos_image[cos_image < cos_thres] = 0 + + # Shrink visible mask + kernel_size = self.bake_unreliable_kernel_size * 2 + 1 + kernel = torch.ones((1, 1, kernel_size, kernel_size), dtype=torch.float32).to( + sketch_image.device + ) + + visible_mask_proc = visible_mask.permute(2, 0, 1).unsqueeze(0).float() + visible_mask_proc = F.conv2d( + 1.0 - visible_mask_proc, kernel, padding=kernel_size // 2 + ) + visible_mask_proc = 1.0 - (visible_mask_proc > 0).float() + visible_mask_proc = visible_mask_proc.squeeze(0).permute(1, 2, 0) + + sketch_proc = sketch_image.permute(2, 0, 1).unsqueeze(0) + sketch_proc = F.conv2d(sketch_proc, kernel, padding=kernel_size // 2) + sketch_proc = (sketch_proc > 0).float() + sketch_proc = sketch_proc.squeeze(0).permute(1, 2, 0) + visible_mask_proc = visible_mask_proc * (sketch_proc < 0.5) + + cos_image[visible_mask_proc == 0] = 0 + + # Linear baking + proj_mask = (visible_mask_proc != 0).view(-1) + uv_flat = uv.squeeze(0).contiguous().view(-1, 2)[proj_mask] + image_flat = image.squeeze(0).contiguous().view(-1, channel)[proj_mask] + cos_flat = cos_image.contiguous().view(-1, 1)[proj_mask] + sketch_flat = sketch_image.contiguous().view(-1, 1)[proj_mask] + + texture = linear_grid_put_2d( + self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], image_flat + ) + cos_map = linear_grid_put_2d( + self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], cos_flat + ) + boundary_map = linear_grid_put_2d( + self.texture_size[1], + self.texture_size[0], + uv_flat[..., [1, 0]], + sketch_flat, + ) + + return texture, cos_map, boundary_map + + def bake_from_multiview( + self, + views: List[Image.Image], + camera_elevs: List[float], + camera_azims: List[float], + view_weights: List[float], + method: str = "fast", + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Bake texture from multiple views.""" + project_textures, project_weighted_cos_maps = [], [] + bake_exp = 4 + + for view, camera_elev, camera_azim, weight in zip( + views, camera_elevs, camera_azims, view_weights + ): + project_texture, project_cos_map, _ = self.back_project( + view, camera_elev, camera_azim + ) + project_cos_map = weight * (project_cos_map**bake_exp) + project_textures.append(project_texture) + project_weighted_cos_maps.append(project_cos_map) + + if method == "fast": + texture, ori_trust_map = self.fast_bake_texture( + project_textures, project_weighted_cos_maps + ) + else: + raise ValueError(f"Unknown bake method: {method}") + + return texture, ori_trust_map > 1e-8 + + @torch.no_grad() + def fast_bake_texture( + self, + textures: List[torch.Tensor], + cos_maps: List[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Fast texture baking by weighted averaging.""" + channel = textures[0].shape[-1] + texture_merge = torch.zeros(self.texture_size + (channel,)).to(self.device) + trust_map_merge = torch.zeros(self.texture_size + (1,)).to(self.device) + + for texture, cos_map in zip(textures, cos_maps): + view_sum = (cos_map > 0).sum() + painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum() + if view_sum > 0 and painted_sum / view_sum > 0.99: + continue + texture_merge += texture * cos_map + trust_map_merge += cos_map + + texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8) + texture_merge = texture_merge.clamp(0.0, 1.0) + + return texture_merge, trust_map_merge > 1e-8 + + def texture_inpaint( + self, + texture: torch.Tensor, + mask: Union[torch.Tensor, np.ndarray], + ) -> torch.Tensor: + """Inpaint missing regions in UV texture using mesh-aware method.""" + if isinstance(texture, torch.Tensor): + texture_np = texture.cpu().numpy() + else: + texture_np = texture + + if isinstance(mask, torch.Tensor): + mask_np = mask.cpu().numpy() + else: + mask_np = mask + + # Ensure proper format + if texture_np.max() <= 1.0: + texture_np = texture_np.astype(np.float32) + else: + texture_np = (texture_np / 255.0).astype(np.float32) + + if mask_np.ndim == 3: + mask_np = mask_np.squeeze(-1) + if mask_np.dtype == np.uint8: + mask_uint8 = mask_np + else: + mask_uint8 = ((mask_np > 0) * 255).astype(np.uint8) + + # Get mesh data for mesh-aware inpainting + vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh() + + if vtx_uv is not None and uv_idx is not None: + texture_np, mask_uint8 = meshVerticeInpaint( + texture_np, mask_uint8, vtx_pos, vtx_uv, pos_idx, uv_idx + ) + + # Final OpenCV inpainting for remaining holes + texture_uint8 = (texture_np * 255).astype(np.uint8) + inpaint_mask = 255 - mask_uint8 + texture_inpainted = cv2.inpaint(texture_uint8, inpaint_mask, 3, cv2.INPAINT_NS) + + return torch.from_numpy(texture_inpainted / 255.0).float().to(self.device) + + # Alias for compatibility + uv_inpaint = texture_inpaint + + +def array_to_tensor(np_array): + """Convert numpy array to normalized tensor.""" + image_pt = torch.tensor(np_array).float() + image_pt = image_pt / 255 * 2 - 1 + image_pt = rearrange(image_pt, "h w c -> c h w") + image_pts = repeat(image_pt, "c h w -> b c h w", b=1) + return image_pts + + +def recenter_image(image, border_ratio=0.2): + """Recenter a PIL image, cropping to non-transparent content with a border.""" + from PIL import Image as PILImage + + if image.mode == "RGB": + return image + elif image.mode == "L": + return image.convert("RGB") + if image.mode != "RGBA": + image = image.convert("RGBA") + + alpha_channel = np.array(image)[:, :, 3] + non_zero_indices = np.argwhere(alpha_channel > 0) + if non_zero_indices.size == 0: + raise ValueError("Image is fully transparent") + + min_row, min_col = non_zero_indices.min(axis=0) + max_row, max_col = non_zero_indices.max(axis=0) + + cropped_image = image.crop((min_col, min_row, max_col + 1, max_row + 1)) + + width, height = cropped_image.size + border_width = int(width * border_ratio) + border_height = int(height * border_ratio) + + new_width = width + 2 * border_width + new_height = height + 2 * border_height + square_size = max(new_width, new_height) + + new_image = PILImage.new("RGBA", (square_size, square_size), (255, 255, 255, 0)) + + paste_x = (square_size - new_width) // 2 + border_width + paste_y = (square_size - new_height) // 2 + border_height + new_image.paste(cropped_image, (paste_x, paste_y)) + return new_image + + +class ImageProcessorV2: + """Image processor for Hunyuan3D single-view input.""" + + # External module path aliases for compatibility with Hunyuan3D configs + _aliases = [ + "hy3dshape.preprocessors.ImageProcessorV2", + "hy3dgen.shapegen.preprocessors.ImageProcessorV2", + ] + + def __init__(self, size=512, border_ratio=None): + self.size = size + self.border_ratio = border_ratio + + @staticmethod + def recenter(image, border_ratio: float = 0.2): + """recenter an image to leave some empty space at the image border.""" + + if image.shape[-1] == 4: + mask = image[..., 3] + else: + mask = np.ones_like(image[..., 0:1]) * 255 + image = np.concatenate([image, mask], axis=-1) + mask = mask[..., 0] + + height, width, channels = image.shape + + size = max(height, width) + result = np.zeros((size, size, channels), dtype=np.uint8) + + coords = np.nonzero(mask) + x_min, x_max = coords[0].min(), coords[0].max() + y_min, y_max = coords[1].min(), coords[1].max() + crop_h = x_max - x_min + crop_w = y_max - y_min + if crop_h == 0 or crop_w == 0: + raise ValueError("input image is empty") + desired_size = int(size * (1 - border_ratio)) + scale = desired_size / max(crop_h, crop_w) + scaled_h = int(crop_h * scale) + scaled_w = int(crop_w * scale) + x2_min = (size - scaled_h) // 2 + x2_max = x2_min + scaled_h + + y2_min = (size - scaled_w) // 2 + y2_max = y2_min + scaled_w + + result[x2_min:x2_max, y2_min:y2_max] = cv2.resize( + image[x_min:x_max, y_min:y_max], + (scaled_w, scaled_h), + interpolation=cv2.INTER_AREA, + ) + + bg = np.ones((result.shape[0], result.shape[1], 3), dtype=np.uint8) * 255 + + mask = result[..., 3:].astype(np.float32) / 255 + result = result[..., :3] * mask + bg * (1 - mask) + + mask = mask * 255 + result = result.clip(0, 255).astype(np.uint8) + mask = mask.clip(0, 255).astype(np.uint8) + return result, mask + + def load_image(self, image, border_ratio=0.15, to_tensor=True): + if isinstance(image, str): + image = cv2.imread(image, cv2.IMREAD_UNCHANGED) + image, mask = self.recenter(image, border_ratio=border_ratio) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + elif isinstance(image, Image.Image): + image = image.convert("RGBA") + image = np.asarray(image) + image, mask = self.recenter(image, border_ratio=border_ratio) + + image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC) + mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST) + mask = mask[..., np.newaxis] + + if to_tensor: + image = array_to_tensor(image) + mask = array_to_tensor(mask) + return image, mask + + def __call__(self, image, border_ratio=0.15, to_tensor=True, **kwargs): + if self.border_ratio is not None: + border_ratio = self.border_ratio + image, mask = self.load_image( + image, border_ratio=border_ratio, to_tensor=to_tensor + ) + outputs = {"image": image, "mask": mask} + return outputs + + +class MVImageProcessorV2(ImageProcessorV2): + """Multi-view image processor for Hunyuan3D.""" + + # External module path aliases for compatibility with Hunyuan3D configs + _aliases = [ + "hy3dshape.preprocessors.MVImageProcessorV2", + ] + + return_view_idx = True + + def __init__(self, size=512, border_ratio=None): + super().__init__(size, border_ratio) + self.view2idx = {"front": 0, "left": 1, "back": 2, "right": 3} + + def __call__(self, image_dict, border_ratio=0.15, to_tensor=True, **kwargs): + if self.border_ratio is not None: + border_ratio = self.border_ratio + + images = [] + masks = [] + view_idxs = [] + for view_tag, image in image_dict.items(): + view_idxs.append(self.view2idx[view_tag]) + image, mask = self.load_image( + image, border_ratio=border_ratio, to_tensor=to_tensor + ) + images.append(image) + masks.append(mask) + + zipped_lists = zip(view_idxs, images, masks) + sorted_zipped_lists = sorted(zipped_lists) + view_idxs, images, masks = zip(*sorted_zipped_lists) + + image = torch.cat(images, 0).unsqueeze(0) + mask = torch.cat(masks, 0).unsqueeze(0) + outputs = {"image": image, "mask": mask, "view_idxs": view_idxs} + return outputs + + +# All tool classes available in this module for resolution +TOOL_CLASSES = ( + ImageProcessorV2, + MVImageProcessorV2, +) + + +def resolve_hunyuan3d_tool(target: str): + """Resolve a Hunyuan3D tool class by target string.""" + # First, try to match against _aliases + for cls in TOOL_CLASSES: + aliases = getattr(cls, "_aliases", []) + if target in aliases: + return cls + + # Then, try to match against class names + for cls in TOOL_CLASSES: + if cls.__name__ == target: + return cls + + return None + + +__all__ = [ + "transform_pos", + "get_mv_matrix", + "get_orthographic_projection_matrix", + "get_perspective_projection_matrix", + "export_to_trimesh", + "mesh_uv_wrap", + "meshVerticeInpaint", + "stride_from_shape", + "scatter_add_nd_with_count", + "linear_grid_put_2d", + "MeshRender", + "recenter_image", + "array_to_tensor", + "ImageProcessorV2", + "MVImageProcessorV2", + "TOOL_CLASSES", + "resolve_hunyuan3d_tool", +] diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 08abfd296..cdd9f6700 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -1920,6 +1920,72 @@ "expected_avg_denoise_ms": 173.83, "expected_median_denoise_ms": 178.08 }, + "hunyuan3d_shape_gen": { + "stages_ms": { + "Hunyuan3DShapeBeforeDenoisingStage": 31.42, + "Hunyuan3DShapeDenoisingStage": 3259.83, + "Hunyuan3DShapeExportStage": 8735.55, + "Hunyuan3DShapeSaveStage": 981.64, + "Hunyuan3DPaintPreprocessStage": 226071.67, + "Hunyuan3DPaintTexGenStage": 11083.05, + "Hunyuan3DPaintPostprocessStage": 7469.29 + }, + "denoise_step_ms": { + "0": 32.26, + "1": 63.34, + "2": 65.44, + "3": 65.44, + "4": 65.6, + "5": 65.81, + "6": 65.82, + "7": 65.48, + "8": 65.9, + "9": 65.77, + "10": 65.54, + "11": 65.68, + "12": 65.85, + "13": 65.77, + "14": 65.7, + "15": 65.78, + "16": 66.0, + "17": 66.15, + "18": 65.91, + "19": 66.5, + "20": 65.76, + "21": 66.08, + "22": 66.06, + "23": 66.23, + "24": 65.79, + "25": 65.58, + "26": 65.88, + "27": 65.67, + "28": 65.87, + "29": 66.09, + "30": 65.81, + "31": 65.91, + "32": 66.18, + "33": 65.93, + "34": 66.26, + "35": 66.26, + "36": 66.27, + "37": 65.57, + "38": 66.02, + "39": 66.19, + "40": 65.23, + "41": 66.11, + "42": 66.18, + "43": 65.86, + "44": 65.86, + "45": 65.92, + "46": 65.65, + "47": 65.78, + "48": 66.01, + "49": 66.08 + }, + "expected_e2e_ms": 257696.97, + "expected_avg_denoise_ms": 65.16, + "expected_median_denoise_ms": 65.86 + }, "wan2_1_t2v_1.3b_frame_interp_2x": { "stages_ms": { "InputValidationStage": 0.03, diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 3491a3c3b..690c5ba62 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -734,6 +734,7 @@ Consider updating perf_baselines.json with the snippets below: modality_to_valid_task_types = { "image": {"T2I", "I2I", "TI2I"}, "video": {"T2V", "I2V", "TI2V"}, + "3d": {"I2M"}, } valid_task_types = modality_to_valid_task_types.get( case.server_args.modality, set() @@ -858,6 +859,17 @@ Consider updating perf_baselines.json with the snippets below: # Validation 1: Performance self._validate_and_record(case, perf_record) + # Mesh correctness check (Chamfer Distance) for 3D models + if case.server_args.custom_validator == "mesh": + from sglang.multimodal_gen.test.server.test_server_utils import ( + MESH_OUTPUT_PATHS, + validate_mesh_correctness, + ) + + mesh_path = MESH_OUTPUT_PATHS.pop(case.id, None) + if mesh_path: + validate_mesh_correctness(mesh_path) + # Test /v1/models endpoint for router compatibility self._test_v1_models_endpoint(diffusion_server, case) self._test_t2v_rejects_input_reference(diffusion_server, case) diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 094487b0f..3ba45f0db 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -50,6 +50,10 @@ logger = init_logger(__name__) globally_suppress_loggers() +# Tracks mesh output file paths from generate_mesh for later correctness validation. +# Keyed by case_id, cleaned up after use. +MESH_OUTPUT_PATHS: dict[str, str] = {} + def download_image_from_url(url: str) -> Path: """Download an image from a URL to a temporary file. @@ -637,10 +641,99 @@ class VideoPerformanceValidator(PerformanceValidator): ) +class MeshValidator(PerformanceValidator): + """Validator for 3D mesh generation. Inherits perf validation from PerformanceValidator.""" + + pass + + +HUNYUAN3D_REFERENCE_URL = ( + "https://raw.githubusercontent.com/sgl-project/sgl-test-files/" + "main/diffusion-ci/consistency_gt/1-gpu/hunyuan3d_2_0/hunyuan3d.glb" +) + + +def _download_reference_mesh(url: str) -> Path: + """Download a reference mesh from URL, caching in temp dir.""" + import hashlib + + cache_name = f"ref_mesh_{hashlib.md5(url.encode()).hexdigest()}.glb" + cache_path = Path(tempfile.gettempdir()) / cache_name + if cache_path.exists(): + logger.info(f"Using cached reference mesh: {cache_path}") + return cache_path + + logger.info(f"Downloading reference mesh from: {url}") + with urlopen(url, timeout=60) as resp: + cache_path.write_bytes(resp.read()) + logger.info(f"Reference mesh cached at: {cache_path}") + return cache_path + + +def validate_mesh_correctness( + generated_mesh_path: str, + reference_url: str = HUNYUAN3D_REFERENCE_URL, + num_sample_points: int = 4096, + cd_threshold_ratio: float = 0.01, + random_seed: int = 42, +): + """Validate mesh geometric similarity against a reference via Chamfer Distance. + + Downloads the reference mesh from a URL (cached), samples point clouds from + both meshes, and asserts Chamfer Distance is within threshold. + """ + import numpy as np + + try: + import trimesh + except ImportError: + pytest.fail("trimesh is required for mesh validation: pip install trimesh") + + from scipy.spatial import cKDTree + + # Load generated mesh + generated_mesh = trimesh.load(generated_mesh_path) + if isinstance(generated_mesh, trimesh.Scene): + generated_mesh = generated_mesh.dump(concatenate=True) + + # Download and load reference mesh + ref_path = _download_reference_mesh(reference_url) + reference_mesh = trimesh.load(str(ref_path)) + if isinstance(reference_mesh, trimesh.Scene): + reference_mesh = reference_mesh.dump(concatenate=True) + + # Bounding box diagonal for threshold normalization + ref_bbox = reference_mesh.bounding_box.bounds + bbox_diagonal = float(np.linalg.norm(ref_bbox[1] - ref_bbox[0])) + cd_threshold = cd_threshold_ratio * bbox_diagonal + + # Sample point clouds + np.random.seed(random_seed) + gen_points = np.array( + generated_mesh.sample(num_sample_points, return_index=True)[0] + ) + ref_points = np.array( + reference_mesh.sample(num_sample_points, return_index=True)[0] + ) + + # Bidirectional Chamfer Distance + tree1 = cKDTree(gen_points) + tree2 = cKDTree(ref_points) + forward_cd = float(np.mean(tree2.query(gen_points)[0] ** 2)) + backward_cd = float(np.mean(tree1.query(ref_points)[0] ** 2)) + total_cd = forward_cd + backward_cd + + assert total_cd <= cd_threshold, ( + f"Chamfer Distance check failed: total_cd={total_cd:.6f}, " + f"threshold={cd_threshold:.6f} ({cd_threshold_ratio * 100:.2f}% of bbox diagonal {bbox_diagonal:.4f})" + ) + + # Registry of validators by name VALIDATOR_REGISTRY = { "default": PerformanceValidator, "video": VideoPerformanceValidator, + "mesh": MeshValidator, } @@ -1095,7 +1188,103 @@ def get_generate_fn( }, ) - if modality == "video": + def generate_mesh(case_id, client) -> tuple[str, bytes]: + """I2M: Image to Mesh generation using async /v1/meshes API.""" + import requests as http_requests + + if not sampling_params.image_path: + pytest.skip(f"{case_id}: no input image configured for mesh generation") + + image_path = sampling_params.image_path + if isinstance(image_path, str) and is_image_url(image_path): + image_path = download_image_from_url(image_path) + elif isinstance(image_path, Path): + if not image_path.exists(): + pytest.skip(f"{case_id}: image file missing: {image_path}") + else: + image_path = Path(str(image_path)) + if not image_path.exists(): + pytest.skip(f"{case_id}: image file missing: {image_path}") + + base_url = str(client.base_url).rstrip("/") + if base_url.endswith("/v1"): + base_url = base_url[:-3] + + create_url = f"{base_url}/v1/meshes" + + with open(str(image_path), "rb") as img_file: + files = {"image": (Path(str(image_path)).name, img_file, "image/png")} + data = { + "prompt": "generate 3d mesh", + "model": model_path, + "seed": "0", + "guidance_scale": "5.0", + "num_inference_steps": "50", + } + + logger.info(f"[Mesh Gen] Sending request to {create_url}") + + try: + response = http_requests.post( + create_url, files=files, data=data, timeout=60 + ) + except Exception as e: + pytest.fail(f"{case_id}: mesh creation request failed: {e}") + + if response.status_code != 200: + pytest.fail(f"{case_id}: mesh creation failed: {response.text}") + + job = response.json() + mesh_id = job.get("id") + if not mesh_id: + pytest.fail(f"{case_id}: no mesh id in response: {job}") + + poll_url = f"{base_url}/v1/meshes/{mesh_id}" + poll_interval = 5 + max_wait = 1200 + elapsed = 0 + + while elapsed < max_wait: + time.sleep(poll_interval) + elapsed += poll_interval + + try: + poll_resp = http_requests.get(poll_url, timeout=30) + except Exception as e: + logger.warning(f"[Mesh Gen] Poll failed: {e}") + continue + + if poll_resp.status_code != 200: + continue + + status_data = poll_resp.json() + status = status_data.get("status", "") + + if status == "completed": + content_url = f"{base_url}/v1/meshes/{mesh_id}/content" + try: + content_resp = http_requests.get(content_url, timeout=60) + except Exception as e: + pytest.fail(f"{case_id}: mesh download failed: {e}") + + if content_resp.status_code != 200: + pytest.fail(f"{case_id}: mesh download failed: {content_resp.text}") + + temp_path = Path(tempfile.gettempdir()) / f"mesh_test_{mesh_id}.glb" + temp_path.write_bytes(content_resp.content) + MESH_OUTPUT_PATHS[case_id] = str(temp_path) + + logger.info(f"[Mesh Gen] Mesh downloaded to {temp_path}") + return (mesh_id, b"") + elif status == "failed": + error = status_data.get("error", {}) + pytest.fail(f"{case_id}: mesh generation failed: {error}") + + pytest.fail(f"{case_id}: mesh generation timed out after {max_wait}s") + + if modality == "3d": + fn = generate_mesh + elif modality == "video": if sampling_params.image_path and sampling_params.prompt: if getattr(sampling_params, "direct_url_test", False): fn = generate_text_url_image_to_video diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index f8c980b2f..1a4e32d59 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -190,6 +190,8 @@ class DiffusionServerArgs: self.custom_validator = "image" elif self.modality == "video": self.custom_validator = "video" + elif self.modality == "3d": + self.custom_validator = "mesh" @dataclass(frozen=True) @@ -460,6 +462,11 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ ), ] +HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams( + prompt="", + image_path="https://raw.githubusercontent.com/sgl-project/sgl-test-files/main/diffusion-ci/consistency_gt/1-gpu/hunyuan3d_2_0/hunyuan3d.png", +) + ONE_GPU_CASES_B: list[DiffusionTestCase] = [ # === Text to Video (T2V) === DiffusionTestCase( @@ -591,7 +598,19 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ ), ] -# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536. +# Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm. +if not current_platform.is_hip(): + ONE_GPU_CASES_B.append( + DiffusionTestCase( + "hunyuan3d_shape_gen", + DiffusionServerArgs( + model_path="tencent/Hunyuan3D-2", + modality="3d", + ), + HUNYUAN3D_SHAPE_sampling_params, + ), + ) +# Skip turbowan on AMD: Triton requires 81920 shared memory, but AMD only has 65536. if not current_platform.is_hip(): ONE_GPU_CASES_B.append( DiffusionTestCase(