From 24c80070d6caceffc33edd6eb931a15adccc8bdc Mon Sep 17 00:00:00 2001 From: Daniel Huang Date: Wed, 4 Mar 2026 22:53:30 -0800 Subject: [PATCH] Replace hard-coded cuda generator with current_omni_platform.device_type (#1677) Signed-off-by: Daniel Huang --- tests/e2e/offline_inference/test_diffusion_lora.py | 5 +++-- tests/e2e/offline_inference/test_stable_audio_model.py | 3 ++- tests/e2e/offline_inference/test_teacache.py | 3 ++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/e2e/offline_inference/test_diffusion_lora.py b/tests/e2e/offline_inference/test_diffusion_lora.py index 2465bc607..a95ab43da 100644 --- a/tests/e2e/offline_inference/test_diffusion_lora.py +++ b/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, diff --git a/tests/e2e/offline_inference/test_stable_audio_model.py b/tests/e2e/offline_inference/test_stable_audio_model.py index ad879d751..2acfe5097 100644 --- a/tests/e2e/offline_inference/test_stable_audio_model.py +++ b/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, diff --git a/tests/e2e/offline_inference/test_teacache.py b/tests/e2e/offline_inference/test_teacache.py index d97991779..96254f5b6 100644 --- a/tests/e2e/offline_inference/test_teacache.py +++ b/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 ), )