[Diffusion][NPU] Add support for Hunyuan3D (#20352)

Co-authored-by: Elizaveta Martirosian <elizaveta.martirosian@gmail.com>

python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py

@@ -18,21 +18,24 @@ _abs_path = os.path.dirname(os.path.abspath(__file__))
 _custom_rasterizer_kernel = None
 
 
-def _load_custom_rasterizer():
+def _load_custom_rasterizer(
+    is_cuda: bool = True,
+):
     """JIT compile and load the custom rasterizer kernel."""
     global _custom_rasterizer_kernel
 
     if _custom_rasterizer_kernel is not None:
         return _custom_rasterizer_kernel
-
+    
+    cuda_enabled_flag = ["-DCUDA_ENABLED"] if is_cuda else []
+    
     _custom_rasterizer_kernel = load_extension_with_recovery(
         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"],
+        ] + ([f"{_abs_path}/rasterizer_gpu.cu"] if is_cuda else []),
+        extra_cflags=["-O3"] + cuda_enabled_flag,
+        extra_cuda_cflags=["-O3", "--use_fast_math"] + cuda_enabled_flag,
         verbose=False,
     )
     return _custom_rasterizer_kernel
@@ -46,7 +49,8 @@ def rasterize(
     use_depth_prior: int = 0,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Rasterize mesh to get face indices and barycentric coordinates."""
-    kernel = _load_custom_rasterizer()
+    device = "cpu" if pos.device.type == "npu" else pos.device.type
+    kernel = _load_custom_rasterizer(device == "cuda")
 
     if clamp_depth is None:
         clamp_depth = torch.zeros(0, device=pos.device)
@@ -56,8 +60,12 @@ def rasterize(
         pos = pos[0]
 
     findices, barycentric = kernel.rasterize_image(
-        pos, tri, clamp_depth, resolution[1], resolution[0], 1e-6, use_depth_prior
+        pos.to(device), tri.to(device), clamp_depth.to(device), resolution[1], resolution[0], 1e-6, use_depth_prior
     )
+
+    findices = findices.to(pos.device)
+    barycentric = barycentric.to(pos.device)
+
     return findices, barycentric
 
 

python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp

@@ -128,11 +128,11 @@ std::vector<torch::Tensor> rasterize_image_cpu(torch::Tensor V, torch::Tensor F,
 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);
+#ifdef CUDA_ENABLED
+    return rasterize_image_gpu(V, F, D, width, height, occlusion_truncation, use_depth_prior);
+#else
+    return rasterize_image_cpu(V, F, D, width, height, occlusion_truncation, use_depth_prior);
+#endif
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h

@@ -8,7 +8,13 @@
 #include <torch/extension.h>
 #include <vector>
 #include <ATen/ATen.h>
+
+#ifdef CUDA_ENABLED
 #include <ATen/cuda/CUDAContext.h>
+#else
+#define __host__
+#define __device__
+#endif
 
 #define INT64 unsigned long long
 #define MAXINT 2147483647

python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py

@@ -335,11 +335,11 @@ class Hunyuan3DPaintPreprocessStage(PipelineStage):
             image_array[:, :, 3] = alpha_channel
             image = PILImage.fromarray(image_array)
 
-            image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device)
+            image_tensor = torch.tensor(np.array(image) / 255.0).float().to(self.device)
             alpha = image_tensor[:, :, 3:]
             rgb_target = image_tensor[:, :, :3]
         else:
-            image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device)
+            image_tensor = torch.tensor(np.array(image) / 255.0).float().to(self.device)
             alpha = torch.ones_like(image_tensor)[:, :, :1]
             rgb_target = image_tensor[:, :, :3]
 
@@ -356,7 +356,7 @@ class Hunyuan3DPaintPreprocessStage(PipelineStage):
             guidance_scale=self.config.delight_guidance_scale,
         ).images[0]
 
-        image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device)
+        image_tensor = torch.tensor(np.array(image) / 255.0).float().to(self.device)
         rgb_src = image_tensor[:, :, :3]
         image = _recorrect_rgb(rgb_src, rgb_target, alpha)
         image = image[:, :, :3] * image[:, :, 3:] + torch.ones_like(image[:, :, :3]) * (
@@ -401,6 +401,7 @@ class Hunyuan3DPaintPreprocessStage(PipelineStage):
         self._renderer = MeshRender(
             default_resolution=self.config.paint_render_size,
             texture_size=self.config.paint_texture_size,
+            device=self.device,
         )
         self._renderer_loaded = True
         logger.info("Mesh renderer initialized")