[Bugfix] Place the MAGI-2 preview pipeline on the current accelerator

docs/models/supported_models.md

@@ -59,7 +59,7 @@ th {
 | `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
 | `SeedVR2Pipeline` | [SeedVR2 3B](seedvr2.md) video restoration | `ByteDance-Seed/SeedVR2-3B` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/ByteDance/SeedVR2-RTX-5090.md) |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
-| `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) |
+| `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
 | `Pi05Pipeline` | π0.5 VLA | `lerobot/pi05_base` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi05.md) |

tests/diffusion/models/magi2/test_musa_pipeline.py

@@ -60,11 +60,11 @@ def test_constructor_uses_selected_device(monkeypatch, device_type):
 
 
 @pytest.mark.cpu
-@pytest.mark.parametrize("device_type,available", [("cuda", False), ("musa", False), ("cpu", True), ("xpu", True)])
+@pytest.mark.parametrize("device_type,available", [("cuda", False), ("musa", False), ("cpu", True)])
 def test_constructor_rejects_unsupported_or_unavailable(monkeypatch, device_type, available):
     monkeypatch.setattr(pipeline, "current_omni_platform", fake_platform(device_type, available))
     resolver = stop_before_weights(monkeypatch)
-    with pytest.raises(RuntimeError, match="available CUDA or MUSA"):
+    with pytest.raises(RuntimeError, match="requires an accelerator"):
         pipeline.Magi2Pipeline(od_config=config())
     resolver.assert_not_called()
 

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -300,8 +300,9 @@ def test_forward_stops_peak_monitor_when_initial_synchronize_fails(monkeypatch):
     pipe.enable_diffusion_pipeline_profiler = True
     monitor = Mock(peak_bytes=0)
     monkeypatch.setattr(
-        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
-        lambda: True,
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform",
+        "device_type",
+        "cuda",
     )
     monkeypatch.setattr(
         "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",

vllm_omni/diffusion/models/magi2/pipeline_magi2.py

@@ -550,9 +550,8 @@ def __init__(self, od_config: OmniDiffusionConfig, **kwargs: object) -> None:
         if not od_config.model:
             raise ValueError("MAGI-2 requires od_config.model")
         _validate_native_topology(od_config)
-        supported_gpu = current_omni_platform.is_cuda() or current_omni_platform.is_musa()
-        if not supported_gpu or not current_omni_platform.is_available():
-            raise RuntimeError("MAGI-2 Preview requires available CUDA or MUSA GPUs")
+        if current_omni_platform.device_type == "cpu" or not current_omni_platform.is_available():
+            raise RuntimeError("MAGI-2 Preview requires an accelerator")
 
         self.od_config = od_config
         self.dtype = od_config.dtype or torch.bfloat16
@@ -1085,10 +1084,8 @@ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         seed = _resolve_request_seed(sampling)
         _seed_request(seed)
 
-        has_accelerator = (
-            current_omni_platform.is_cuda() or current_omni_platform.is_musa()
-        ) and current_omni_platform.is_available()
-        device_index = torch.accelerator.current_device_index() if has_accelerator else None
+        on_accelerator = current_omni_platform.device_type != "cpu" and current_omni_platform.is_available()
+        device_index = torch.accelerator.current_device_index() if on_accelerator else None
         # Sampling reserved memory is qualification instrumentation, not part
         # of ordinary serving. It starts only when the pipeline profiler is
         # explicitly enabled, so ordinary GPU requests avoid a 20 Hz thread.
@@ -1112,7 +1109,7 @@ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
                 height=height,
                 num_inference_steps=steps,
             )
-            if has_accelerator:
+            if on_accelerator:
                 torch.accelerator.synchronize()
         finally:
             if monitor_started:
@@ -1123,7 +1120,7 @@ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
             video = _resize_video(video, output_width, output_height)
 
         peak_memory_mb = monitor.peak_bytes / 1024**2 if monitor is not None else 0.0
-        if has_accelerator and dist.is_available() and dist.is_initialized() and self._parallel_group.world_size > 1:
+        if on_accelerator and dist.is_available() and dist.is_initialized() and self._parallel_group.world_size > 1:
             peak = torch.tensor(peak_memory_mb, device=self.device_str)
             dist.all_reduce(
                 peak,