[HenryVarro666]

Fixes #180.

Root cause

load_asr_model() passes the target device to transformers.pipeline() via device_map=. A plain device (e.g. torch.device("cuda:1"), "cpu") passed as device_map is not honored for single-device placement — the pipeline auto-selects an accelerator, which in practice is always GPU 0. So the Whisper model lands on cuda:0 no matter where the main model lives, and there is no way to redirect it — which matches the screenshots in #180 ("绑不上去").

The current log line even reports the wrong thing. On master, with the main model pinned to cuda:1 (2x H100 NVL):

m = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="cuda:1", dtype=torch.float16)
m.load_asr_model()          # logs "ASR model loaded on cuda:1."
m._asr_pipe.model.device    # -> device(type='cuda', index=0)   (actually on GPU 0)

Fix

• Use device= (single-device placement, honored by the pipeline) instead of device_map=. The default behavior now does what the log line always claimed: the ASR model follows the main model's device.
• Add an asr_device kwarg to from_pretrained and a device parameter to load_asr_model() (same plumbing as asr_model_name, Honor asr_model_name when the ASR model is loaded lazily #221), so multi-GPU deployments can place Whisper explicitly — e.g. on a second GPU (the 显存泄露问题 #180 use case) or on the CPU:

# spread across GPUs
model = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="cuda:0",
                                  dtype=torch.float16, asr_device="cuda:1")
# or keep Whisper off the GPUs entirely
model.load_asr_model(device="cpu")

The ASR dtype selection (fp16 on cuda/xpu, fp32 otherwise) now keys off the ASR device instead of the main model device, so asr_device="cpu" correctly gets fp32.

Verification (2x NVIDIA H100 NVL, torch 2.8.0 + Apple Silicon MPS)

Before (master 1574e06):

setup	main model	ASR lands on
device_map="cuda:1"	cuda:1	cuda:0 (wrong GPU)
device_map="auto" (sharded 2 GPUs)	shards on 0+1	cuda:0 (no way to move)
asr_device=...	—	TypeError (no such option)

After (this branch):

setup	ASR lands on
device_map="cuda:1", default ASR	cuda:1 (follows main model)
device_map="auto" + asr_device="cuda:1"	cuda:1, fp16
load_asr_model(device="cpu") per-call	cpu, fp32
MPS Mac + asr_device="cpu"	cpu; auto-transcription output unchanged

End-to-end on cuda:1 after the fix: generate(text=..., instruct="male, low pitch") produces normal audio (2.44 s), and auto-transcription of a reference clip returns the same text as before the change.

Single-GPU and Mac setups are unaffected: there the auto-selected device and the main model device coincide, so placement is identical to before.

pre-commit run --all-files passes.

[gemini-code-assist]

Code Review

This pull request introduces the ability to specify a target device (asr_device) for loading the Whisper ASR model, allowing users to control device placement (e.g., offloading to a different GPU or CPU). This is documented in the README and implemented in omnivoice.py by passing the device parameter directly to the Hugging Face pipeline instead of using device_map. The review feedback correctly identifies that the logic for determining the model's data type (torch.float16 vs torch.float32) based on the device name is fragile, as it fails to recognize integer or string-digit GPU indices (e.g., 0 or "1"), and provides a robust code suggestion to fix this.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If device is specified as an integer (e.g., 0 or 1) or a string digit (e.g., "0"), which are valid ways to specify GPU indices in Hugging Face pipelines, str(device).startswith(("cuda", "xpu")) will evaluate to False. This causes the ASR model to load with torch.float32 instead of torch.float16 on GPUs, leading to higher memory usage and slower inference.\n\nWe should update the check to robustly identify integer and string-digit GPU device specifications.

        is_cuda_or_xpu = (\n            (isinstance(device, int) and device >= 0)\n            or (isinstance(device, str) and device.isdigit())\n            or str(device).startswith(("cuda", "xpu"))\n        )\n        asr_dtype = torch.float16 if is_cuda_or_xpu else torch.float32

[zhu-han]

Thanks again for your contribution!