From 5d691a44f45815cd4f908c6df0a759800b142072 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Fri, 5 Jun 2026 19:15:08 +0800 Subject: [PATCH] [diffusion] fix: fix LingBot World timestep error on MUSA (#27341) Signed-off-by: Xiaodong Ye --- python/sglang/multimodal_gen/runtime/models/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/models/utils.py b/python/sglang/multimodal_gen/runtime/models/utils.py index af2691b02..7628e5922 100644 --- a/python/sglang/multimodal_gen/runtime/models/utils.py +++ b/python/sglang/multimodal_gen/runtime/models/utils.py @@ -9,6 +9,7 @@ from typing import Any import torch +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.srt.utils import ( get_bool_env_var, is_gfx95_supported, @@ -143,7 +144,10 @@ def pred_noise_to_pred_video( pred_noise = pred_noise.double().to(device) noise_input_latent = noise_input_latent.double().to(device) sigmas = scheduler.sigmas.double().to(device) - timesteps = scheduler.timesteps.double().to(device) + high_dtype = ( + torch.float64 if current_platform.is_float64_supported() else torch.float32 + ) + timesteps = scheduler.timesteps.to(high_dtype).to(device) timestep_id = torch.argmin( (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 )