NeuralFalcon commited on
Commit
3a87885
·
verified ·
1 Parent(s): 0ccb57a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -6
app.py CHANGED
@@ -157,12 +157,20 @@ class VibeVoiceDemo:
157
  def load_model(self):
158
  """Load the VibeVoice model and processor."""
159
  print(f"Loading processor & model from {self.model_path}")
160
- self.processor = VibeVoiceProcessor.from_pretrained(self.model_path)
161
- self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
162
- self.model_path,
163
- torch_dtype=torch.bfloat16,
164
- device_map=self.device,
165
- )
 
 
 
 
 
 
 
 
166
  self.model.eval()
167
  self.model.model.noise_scheduler = self.model.model.noise_scheduler.from_config(
168
  self.model.model.noise_scheduler.config,
 
157
  def load_model(self):
158
  """Load the VibeVoice model and processor."""
159
  print(f"Loading processor & model from {self.model_path}")
160
+ self.processor = VibeVoiceProcessor.from_pretrained(self.model_path)
161
+ if self.device == "cuda":
162
+ self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
163
+ self.model_path,
164
+ torch_dtype=torch.bfloat16,
165
+ device_map=self.device,
166
+ )
167
+ else:
168
+ self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
169
+ self.model_path,
170
+ torch_dtype=torch.float32, # Use float32 for CPU
171
+ device_map="cpu",
172
+ )
173
+
174
  self.model.eval()
175
  self.model.model.noise_scheduler = self.model.model.noise_scheduler.from_config(
176
  self.model.model.noise_scheduler.config,