valenti/spark-vibevoice
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add Apple MPS device support (dtype + attention handling) to demos (#67)

Browse Source 
* Enhance model loading with device support and error handling

Updated device handling for model loading and added support for MPS. Improved error handling and fallback mechanisms for attention implementations.

* Improve device handling and model loading logic

Updated device argument handling to support MPS and added validation for MPS availability. Enhanced model loading logic based on the selected device type.

* fallback only when flash_attention_2 and add some comments back

---------

Co-authored-by: YaoyaoChang <cyy574006791@qq.com>
This commit is contained in:
gregory-fanous
2025-09-01 17:16:01 +08:00
committed by GitHub
co-authored by YaoyaoChang
parent 3074f898ca
commit 6dd49cb720
2 changed files with 130 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo/gradio_demo.py
+64 -22
View File
@@ -47,30 +47,67 @@ class VibeVoiceDemo:
def load_model(self):
"""Load the VibeVoice model and processor."""
print(f"Loading processor & model from {self.model_path}")
# Normalize potential 'mpx'
if self.device.lower() == "mpx":
print("Note: device 'mpx' detected, treating it as 'mps'.")
self.device = "mps"
if self.device == "mps" and not torch.backends.mps.is_available():
print("Warning: MPS not available. Falling back to CPU.")
self.device = "cpu"
print(f"Using device: {self.device}")
# Load processor
self.processor = VibeVoiceProcessor.from_pretrained(
self.model_path,
)
self.processor = VibeVoiceProcessor.from_pretrained(self.model_path)
# Decide dtype & attention
if self.device == "mps":
load_dtype = torch.float32
attn_impl_primary = "sdpa"
elif self.device == "cuda":
load_dtype = torch.bfloat16
attn_impl_primary = "flash_attention_2"
else:
load_dtype = torch.float32
attn_impl_primary = "sdpa"
print(f"Using device: {self.device}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}")
# Load model
try:
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=torch.bfloat16,
device_map='cuda',
attn_implementation='flash_attention_2' # flash_attention_2 is recommended
)
if self.device == "mps":
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=load_dtype,
attn_implementation=attn_impl_primary,
device_map=None,
)
self.model.to("mps")
elif self.device == "cuda":
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=load_dtype,
device_map="cuda",
attn_implementation=attn_impl_primary,
)
else:
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=load_dtype,
device_map="cpu",
attn_implementation=attn_impl_primary,
)
except Exception as e:
print(f"[ERROR] : {type(e).__name__}: {e}")
print(traceback.format_exc())
print("Error loading the model. Trying to use SDPA. However, note that only flash_attention_2 has been fully tested, and using SDPA may result in lower audio quality.")
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=torch.bfloat16,
device_map='cuda',
attn_implementation='sdpa'
)
if attn_impl_primary == 'flash_attention_2':
print(f"[ERROR] : {type(e).__name__}: {e}")
print(traceback.format_exc())
fallback_attn = "sdpa"
print(f"Falling back to attention implementation: {fallback_attn}")
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=load_dtype,
device_map=(self.device if self.device in ("cuda", "cpu") else None),
attn_implementation=fallback_attn,
)
if self.device == "mps":
self.model.to("mps")
else:
raise e
self.model.eval()
# Use SDE solver by default
@@ -238,6 +275,11 @@ class VibeVoiceDemo:
return_tensors="pt",
return_attention_mask=True,
)
# Move tensors to device
target_device = self.device if self.device in ("cuda", "mps") else "cpu"
for k, v in inputs.items():
if torch.is_tensor(v):
inputs[k] = v.to(target_device)
# Create audio streamer
audio_streamer = AudioStreamer(
@@ -1133,8 +1175,8 @@ def parse_args():
parser.add_argument(
"--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device for inference",
default=("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")),
help="Device for inference: cuda | mps | cpu",
)
parser.add_argument(
"--inference_steps",