[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:
@@ -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 = [
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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"
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
29
python/sglang/multimodal_gen/configs/sample/hunyuan3d.py
Normal file
29
python/sglang/multimodal_gen/configs/sample/hunyuan3d.py
Normal 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__()
|
||||
@@ -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:
|
||||
|
||||
@@ -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"]
|
||||
@@ -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");
|
||||
}
|
||||
@@ -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
|
||||
@@ -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};
|
||||
}
|
||||
@@ -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"]
|
||||
@@ -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");
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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] = []
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -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__
|
||||
|
||||
@@ -45,3 +45,4 @@ class AsyncDictStore:
|
||||
# [request_id, dict]
|
||||
VIDEO_STORE = AsyncDictStore()
|
||||
IMAGE_STORE = AsyncDictStore()
|
||||
MESH_STORE = AsyncDictStore()
|
||||
|
||||
@@ -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:
|
||||
|
||||
1401
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py
Normal file
1401
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
]
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
]
|
||||
1224
python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py
Normal file
1224
python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
|
||||
1114
python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py
Normal file
1114
python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user