Replace hard-coded cuda generator with current_omni_platform.device_type (#1677)

Signed-off-by: Daniel Huang <daniel1.huang@intel.com>

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -9,6 +9,7 @@ from safetensors.torch import save_file
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -89,7 +90,7 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
                 width=width,
                 num_inference_steps=2,
                 guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                 num_outputs_per_prompt=1,
             ),
         )
@@ -119,7 +120,7 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
                     width=width,
                     num_inference_steps=2,
                     guidance_scale=0.0,
-                    generator=torch.Generator("cuda").manual_seed(42),
+                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                     num_outputs_per_prompt=1,
                     lora_request=lora_request,
                     lora_scale=2.0,

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -8,6 +8,7 @@ import torch
 from tests.utils import hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -41,7 +42,7 @@ def test_stable_audio_model(model_name: str):
         sampling_params_list=OmniDiffusionSamplingParams(
             num_inference_steps=4,  # Minimal steps for speed
             guidance_scale=7.0,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
             num_outputs_per_prompt=1,
             extra_args={
                 "audio_start_in_s": audio_start_in_s,

tests/e2e/offline_inference/test_teacache.py

@@ -17,6 +17,7 @@ import torch
 
 from tests.utils import hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -63,7 +64,7 @@ def test_teacache(model_name: str):
                 width=width,
                 num_inference_steps=num_inference_steps,
                 guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )