Update app.py

app.py

@@ -146,13 +146,26 @@ def generate_speech(
         # Clean text
         full_script = text.strip().replace("'", "'").replace('"', '"').replace('"', '"')
 
-        # Get voice sample
+        # Get voice sample path
         voice_sample = VOICE_MAPPER.get_voice_path(speaker_name)
         
-        # Load voice sample to GPU
+        # Load voice sample to CPU first, then move to GPU
         all_prefilled_outputs = torch.load(
-            voice_sample, map_location="cuda", weights_only=False
+            voice_sample, map_location="cpu", weights_only=False
         )
+        
+        # Move model to GPU
+        MODEL.to("cuda")
+        
+        # Move voice sample tensors to GPU
+        if isinstance(all_prefilled_outputs, dict):
+            for key in all_prefilled_outputs:
+                if torch.is_tensor(all_prefilled_outputs[key]):
+                    all_prefilled_outputs[key] = all_prefilled_outputs[key].to("cuda")
+                elif isinstance(all_prefilled_outputs[key], dict):
+                    for sub_key in all_prefilled_outputs[key]:
+                        if torch.is_tensor(all_prefilled_outputs[key][sub_key]):
+                            all_prefilled_outputs[key][sub_key] = all_prefilled_outputs[key][sub_key].to("cuda")
 
         progress(0.2, desc="Preparing inputs...")
 
@@ -165,8 +178,7 @@ def generate_speech(
             return_attention_mask=True,
         )
 
-        # Move model and tensors to GPU
-        MODEL.to("cuda")
+        # Move input tensors to GPU
         for k, v in inputs.items():
             if torch.is_tensor(v):
                 inputs[k] = v.to("cuda")