diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py index 6d410130b..a02a80c85 100644 --- a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py +++ b/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 diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp index 72aa005b3..16773e857 100644 --- a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp @@ -128,11 +128,11 @@ std::vector rasterize_image_cpu(torch::Tensor V, torch::Tensor F, std::vector rasterize_image(torch::Tensor V, torch::Tensor F, torch::Tensor D, int width, int height, float occlusion_truncation, int use_depth_prior) { - int device_id = V.get_device(); - if (device_id == -1) - return rasterize_image_cpu(V, F, D, width, height, occlusion_truncation, use_depth_prior); - else - return rasterize_image_gpu(V, F, D, width, height, occlusion_truncation, use_depth_prior); +#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) { diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h index bb1703cf0..84e12ca71 100644 --- a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h @@ -8,7 +8,13 @@ #include #include #include + +#ifdef CUDA_ENABLED #include +#else +#define __host__ +#define __device__ +#endif #define INT64 unsigned long long #define MAXINT 2147483647 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py index 15baaedd6..fe69d1ac9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py +++ b/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")