From 4a4f063b792b1ed90d2168efe8c595cccbcf8961 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 24 Jun 2026 23:54:01 +0800 Subject: [PATCH] [diffusion] fix: paint multiview vae must follow unit dtype (#29041) --- .../model_specific_stages/hunyuan3d/paint.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py index a031da099..15101a371 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py @@ -571,15 +571,7 @@ class Hunyuan3DPaintTexGenStage(PipelineStage): else: raise FileNotFoundError(f"No VAE weights in {vae_dir}") self.vae.load_state_dict(state_dict) - # Resolve VAE/DiT dtypes from config with simple CPU/MPS fallback - vae_dtype = PRECISION_TO_TYPE.get( - getattr(self.config, "vae_precision", "fp32"), torch.float32 - ) - if self.device.type in ("cpu", "mps") and vae_dtype in ( - torch.float16, - torch.bfloat16, - ): - vae_dtype = torch.float32 + # Resolve the DiT (multiview UNet) dtype from config, with CPU/MPS fallback. dit_dtype = PRECISION_TO_TYPE.get( getattr(self.config, "dit_precision", "fp16"), torch.float16 ) @@ -588,6 +580,13 @@ class Hunyuan3DPaintTexGenStage(PipelineStage): torch.bfloat16, ): dit_dtype = torch.float32 + # The multiview (Stable-Diffusion) AutoencoderKL must share the UNet dtype. + # Reference attention feeds its VAE-encoded ref_latents straight into the + # fp16 UNet, and the official HunyuanPaint pipeline runs VAE+UNet entirely + # in fp16. The `vae_precision` knob targets the 3D ShapeVAE (geometry + # precision) — applying it to this 2D texture VAE produces an + # fp32-input / fp16-weight mismatch that crashes the paint UNet. + vae_dtype = dit_dtype self.vae = self.vae.to(device=self.device, dtype=vae_dtype).eval() self.transformer = UNet2p5DConditionModel.from_pretrained(