[diffusion] model: support Hunyuan3D-2 (#18170)

Co-authored-by: yingluosanqian <yingluosanqian@gmail.com>
Co-authored-by: daiweitao <dwti614707404@163.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Prozac614
2026-03-02 12:28:05 +08:00
committed by GitHub
parent f6ee6dc8c3
commit 57c5c343d7
43 changed files with 8082 additions and 65 deletions

View File

@@ -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 = [

View File

@@ -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 = [

View File

@@ -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 = [

View File

@@ -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():

View File

@@ -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",
]

View File

@@ -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"

View File

@@ -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",
]

View File

@@ -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

View File

@@ -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",

View File

@@ -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:

View File

@@ -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

View File

@@ -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__()

View File

@@ -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:

View File

@@ -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"]

View File

@@ -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<torch::Tensor> 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<float>(), F.data_ptr<int>(), 0,
(INT64*)z_min.data_ptr<int64_t>(), occlusion_truncation, width, height, num_vertices, num_faces, i);
}
} else {
for (int i = 0; i < num_faces; ++i)
rasterizeImagecoordsKernelCPU(V.data_ptr<float>(), F.data_ptr<int>(), D.data_ptr<float>(),
(INT64*)z_min.data_ptr<int64_t>(), 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<float>(), F.data_ptr<int>(),
findices.data_ptr<int>(), (INT64*)z_min.data_ptr<int64_t>(), width, height, num_vertices, num_faces, barycentric.data_ptr<float>(), i);
return {findices, barycentric};
}
std::vector<torch::Tensor> 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");
}

View File

@@ -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 <torch/extension.h>
#include <vector>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#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<torch::Tensor> rasterize_image_gpu(torch::Tensor V, torch::Tensor F, torch::Tensor D,
int width, int height, float occlusion_truncation, int use_depth_prior);
#endif

View File

@@ -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<torch::Tensor> 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<float>(), F.data_ptr<int>(), 0,
(INT64*)z_min.data_ptr<int64_t>(), occlusion_truncation, width, height, num_vertices, num_faces);
} else {
rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr<float>(), F.data_ptr<int>(), D.data_ptr<float>(),
(INT64*)z_min.data_ptr<int64_t>(), 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<float>(), F.data_ptr<int>(),
findices.data_ptr<int>(), (INT64*)z_min.data_ptr<int64_t>(), width, height, num_vertices, num_faces, barycentric.data_ptr<float>());
return {findices, barycentric};
}

View File

@@ -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"]

View File

@@ -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 <vector>
#include <queue>
#include <cmath>
#include <algorithm>
#include <torch/extension.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
namespace py = pybind11;
using namespace std;
std::pair<py::array_t<float>,
py::array_t<uint8_t>> meshVerticeInpaint_smooth(py::array_t<float> texture,
py::array_t<uint8_t> mask,
py::array_t<float> vtx_pos, py::array_t<float> vtx_uv,
py::array_t<int> pos_idx, py::array_t<int> 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<float*>(texture_buf.ptr);
uint8_t* mask_ptr = static_cast<uint8_t*>(mask_buf.ptr);
int vtx_num = vtx_pos_buf.shape[0];
float* vtx_pos_ptr = static_cast<float*>(vtx_pos_buf.ptr);
float* vtx_uv_ptr = static_cast<float*>(vtx_uv_buf.ptr);
int* pos_idx_ptr = static_cast<int*>(pos_idx_buf.ptr);
int* uv_idx_ptr = static_cast<int*>(uv_idx_buf.ptr);
vector<float> vtx_mask(vtx_num, 0.0f);
vector<vector<float>> vtx_color(vtx_num, vector<float>(texture_channel, 0.0f));
vector<int> uncolored_vtxs;
vector<vector<int>> 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<float> sum_color(texture_channel, 0.0f);
float total_weight = 0.0f;
array<float, 3> 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<float, 3> 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<float> new_texture(texture_buf.size);
py::array_t<uint8_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<float*>(new_texture_buf.ptr);
uint8_t* new_mask_ptr = static_cast<uint8_t*>(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<float>, py::array_t<uint8_t>> meshVerticeInpaint(py::array_t<float> texture,
py::array_t<uint8_t> mask,
py::array_t<float> vtx_pos, py::array_t<float> vtx_uv,
py::array_t<int> pos_idx, py::array_t<int> 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");
}

View File

@@ -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

View File

@@ -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] = []

View File

@@ -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

View File

@@ -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)
)

View File

@@ -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)

View File

@@ -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__

View File

@@ -45,3 +45,4 @@ class AsyncDictStore:
# [request_id, dict]
VIDEO_STORE = AsyncDictStore()
IMAGE_STORE = AsyncDictStore()
MESH_STORE = AsyncDictStore()

View File

@@ -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:

File diff suppressed because it is too large Load Diff

View File

@@ -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,
]

View File

@@ -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,

View File

@@ -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,
]

File diff suppressed because it is too large Load Diff

View File

@@ -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

View File

@@ -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",
]

File diff suppressed because it is too large Load Diff

View File

@@ -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",
]

View File

@@ -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,

File diff suppressed because it is too large Load Diff

View File

@@ -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,

View File

@@ -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)

View File

@@ -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

View File

@@ -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(