Update app.py

app.py

@@ -9,17 +9,25 @@ def get_available_speakers():
     speakers = list(_DEFAULT_SPEAKERS.keys())
     return speakers
 
+def cleanup_cuda():
+    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
 @spaces.GPU
 def generate_tts(text, temperature, repetition_penalty, speaker_selection, reference_audio):
-    model_config = outetts.HFModelConfig_v2(
-        model_path="OuteAI/OuteTTS-0.3-1B",
-        tokenizer_path="OuteAI/OuteTTS-0.3-1B",
-        dtype=torch.bfloat16,
-        device="cuda"
-    )
-    interface = outetts.InterfaceHF(model_version="0.3", cfg=model_config)
-
+    # Clear CUDA cache before starting
+    cleanup_cuda()
+    
     try:
+        model_config = outetts.HFModelConfig_v2(
+            model_path="OuteAI/OuteTTS-0.3-1B",
+            tokenizer_path="OuteAI/OuteTTS-0.3-1B",
+            dtype=torch.bfloat16,
+            device="cuda"
+        )
+        interface = outetts.InterfaceHF(model_version="0.3", cfg=model_config)
+
         if reference_audio:
             speaker = interface.create_speaker(reference_audio)
         elif speaker_selection and speaker_selection != "None":
@@ -41,8 +49,16 @@ def generate_tts(text, temperature, repetition_penalty, speaker_selection, refer
 
         output_path = "output.wav"
         output.save(output_path)
+        
+        # Cleanup after generation
+        cleanup_cuda()
+        del interface
+        del model_config
+        del output
+        
         return output_path, None
     except Exception as e:
+        cleanup_cuda()
         return None, str(e)
 
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
@@ -114,4 +130,5 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
         outputs=[error_box]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()