diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 1a387cdcc39..c67ff811f52 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -49,6 +49,7 @@ unpad_prompt_conditioning, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -198,7 +199,7 @@ def _ensure_s3_model(cls): ) from e model = _s3.load_model("speech_tokenizer_v3_25hz") - device = "cuda" if torch.cuda.is_available() else "cpu" + device = current_omni_platform.get_torch_device() model = model.to(device).eval() cls._s3_model = (model, _s3, device) return cls._s3_model diff --git a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py index cdf0da992fc..69ebd29eeb2 100644 --- a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py +++ b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py @@ -17,6 +17,7 @@ DAC_SAMPLE_RATE, build_dac_codec, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -129,7 +130,7 @@ def encode_reference_audio_codes( (dtype=torch.long). """ if device is None: - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + device = current_omni_platform.get_torch_device() else: device = torch.device(device) dtype = torch.float32 diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index 6e14840e24b..664a794c4c0 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -53,6 +53,7 @@ from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -411,7 +412,7 @@ def _ensure_cached_runtime_components( # Load voice clone components (speech tokenizer + CampPlus) if load_voice_clone: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() self.processor_device = device campplus_path = getattr(self, "_campplus_path", None) or resolve_glm_tts_campplus_path(model_dir) self._campplus_path = campplus_path diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py index 6a14707ee5b..91609e1a4ee 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py @@ -41,6 +41,7 @@ load_higgs_audio_codec, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform __all__ = [ "HiggsAudioV2Code2Wav", @@ -173,7 +174,7 @@ def load_codec_from_disk(self, model_dir: str, device: torch.device | None = Non 2. Otherwise fall back to ``/``. """ if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() audio_tokenizer_path = self._resolve_audio_tokenizer_path(model_dir) quantizer, fc2, acoustic_decoder, _tokenizer_config = load_higgs_audio_codec(audio_tokenizer_path, device) if len(quantizer.quantizers) != self.num_codebooks: diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py index cacf82ee56c..960f6cfef02 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py @@ -29,6 +29,8 @@ import numpy as np import torch +from vllm_omni.platforms import current_omni_platform + __all__ = [ "UnsupportedInputError", "MULTI_SPEAKER_TAG_PATTERN", @@ -239,7 +241,7 @@ def _load_audio_tokenizer(): allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], ) audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) - device = "cuda" if torch.cuda.is_available() else "cpu" + device = current_omni_platform.get_torch_device() model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device) return model.eval() diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py index 81c18a32e92..bc82c474b2d 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py @@ -27,6 +27,7 @@ HiggsAudioV3Config, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform __all__ = [ "HiggsAudioV3Code2Wav", @@ -126,7 +127,7 @@ def _load_from_tokenizer_repo(self) -> None: else: raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() quantizer, fc2, acoustic_decoder, _cfg = load_higgs_audio_codec(tokenizer_dir, device) self.quantizer = quantizer self.fc2 = fc2 @@ -278,7 +279,7 @@ def _load_from_bundled_state( ) if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() # The bundled codec may use boson-ai key naming. Try remapping. needs_remap = any(k.startswith("quantizer.vq.layers.") for k in codec_state) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index bd178afe667..26482c6bce0 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -23,6 +23,7 @@ from vllm_omni.model_executor.models.mimo_audio.cuda_graph_decoder_wrapper import CUDAGraphMiMoDecoderWrapper from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -442,8 +443,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor(config.vocab_size) - device_str = "cuda" if torch.cuda.is_available() else "cpu" - self.device = torch.device(device_str) + self.device = current_omni_platform.get_torch_device() self.sample_rate = getattr(config, "audio_sample_rate", 24000) self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path, trust_remote_code=True) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 7746b423e9e..d9b48807ae4 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -291,7 +291,7 @@ def forward( device = inputs_embeds.device else: num_tokens = 1 - device = torch.device("cuda") + device = current_omni_platform.get_torch_device() hidden_dim = self.config.hidden_size if hasattr(self.config, "hidden_size") else 2560 # Profile/dummy run: both input_ids and inputs_embeds are None. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index c7208752ba8..083f54afe3d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -25,6 +25,8 @@ from vllm.config import VllmConfig from vllm.model_executor.models.interfaces import SupportsPP +from vllm_omni.platforms import current_omni_platform + try: from stepaudio2 import Token2wav as _Token2wav @@ -379,7 +381,7 @@ def _dummy_hidden_states( break else: num_tokens = 1 - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() hidden_size = int(getattr(self, "_hidden_size", 768) or 768) return torch.zeros((num_tokens, hidden_size), device=device, dtype=torch.bfloat16) diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index 87b93288684..1a86c2b82fd 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -38,6 +38,7 @@ reinit_rotary_inv_freq, transformers_keys_to_ignore_compat, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -183,7 +184,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # Eager construction (not in load_weights) -- see module docstring. _patch_torchaudio_load() - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() self._device: torch.device = device if device.type == "cuda" and torch.cuda.is_bf16_supported(): diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index c5026086e23..c37caeef540 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -38,6 +38,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig logger = init_logger(__name__) @@ -495,7 +496,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: try: device = next(self.parameters()).device except StopIteration: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() model_dir = self._resolve_model_dir() diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 4a144bff027..b800d8d143f 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -46,6 +46,7 @@ Qwen3OmniMoeThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape +from vllm_omni.platforms import current_omni_platform # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -1213,7 +1214,7 @@ def _talker_to_code_predictor( else ( talker_hidden_states.device if isinstance(talker_hidden_states, torch.Tensor) - else torch.device("cuda" if torch.cuda.is_available() else "cpu") + else current_omni_platform.get_torch_device() ) ) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 14aa93b20fb..19c569e996c 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -243,6 +243,10 @@ class UnspecifiedOmniPlatform(OmniPlatform): _enum = PlatformEnum.UNSPECIFIED device_type = "cpu" + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + return torch.device("cpu") + @classmethod def get_device_count(cls) -> int: return 0 diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 08551ad6948..8eafd531ef8 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -12,6 +12,7 @@ from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -58,7 +59,7 @@ def init_device(self): f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) - self.device = torch.device(f"cuda:{self.local_rank}") + self.device = current_omni_platform.get_torch_device(self.local_rank) torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 2a85436a007..38a88c84c05 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -11,6 +11,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -57,7 +58,7 @@ def init_device(self): f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) - self.device = torch.device(f"cuda:{self.local_rank}") + self.device = current_omni_platform.get_torch_device(self.local_rank) torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype)