diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py index c74482dae..7689c5ddb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py @@ -838,9 +838,11 @@ class SanaWMDenoisingStage(DenoisingStage): scheduler = getattr( batch, "scheduler", None ) or get_or_create_request_scheduler(batch, self.scheduler) + self._move_scheduler_tensors_to_device(scheduler, device) timesteps = batch.timesteps if timesteps is None: raise ValueError("SANA-WM denoising requires prepared timesteps.") + timesteps = timesteps.to(device=device) latents = batch.latents.to(device=device, dtype=target_dtype) init_latents = latents.clone() @@ -1040,6 +1042,17 @@ class SanaWMDenoisingStage(DenoisingStage): batch.latents = server_args.pipeline_config.post_denoising_loop(latents, batch) return batch + @staticmethod + def _move_scheduler_tensors_to_device(scheduler: object, device) -> None: + for name in ("sigmas", "timesteps"): + for attr_name in (name, f"_{name}"): + value = getattr(scheduler, attr_name, None) + if isinstance(value, torch.Tensor): + try: + setattr(scheduler, attr_name, value.to(device=device)) + except AttributeError: + pass + class SanaWMBeforeDenoisingStage(PipelineStage): """ diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py index 75e594b37..39f847002 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py @@ -1079,6 +1079,44 @@ class TestSanaWMDenoisingStage(unittest.TestCase): self.assertTrue(torch.allclose(combined_from_pos_rank, serial)) self.assertTrue(torch.allclose(combined_from_neg_rank, serial)) + def test_scheduler_tensors_move_to_target_device(self) -> None: + scheduler = SimpleNamespace( + sigmas=torch.tensor([1.0, 0.0]), + timesteps=torch.tensor([1000.0, 0.0]), + untouched="value", + ) + + SanaWMDenoisingStage._move_scheduler_tensors_to_device( + scheduler, torch.device("cpu") + ) + + self.assertEqual(scheduler.sigmas.device.type, "cpu") + self.assertEqual(scheduler.timesteps.device.type, "cpu") + self.assertEqual(scheduler.untouched, "value") + + class PropertyScheduler: + def __init__(self): + self._sigmas = torch.tensor([1.0, 0.0]) + self._timesteps = torch.tensor([1000.0, 0.0]) + self.untouched = "value" + + @property + def sigmas(self): + return self._sigmas + + @property + def timesteps(self): + return self._timesteps + + property_scheduler = PropertyScheduler() + SanaWMDenoisingStage._move_scheduler_tensors_to_device( + property_scheduler, torch.device("cpu") + ) + + self.assertEqual(property_scheduler.sigmas.device.type, "cpu") + self.assertEqual(property_scheduler.timesteps.device.type, "cpu") + self.assertEqual(property_scheduler.untouched, "value") + class TestSanaWMNativeDiTChunking(unittest.TestCase): def test_softmax_chunking_is_disabled_by_default_for_upstream_parity(self) -> None: