Update app.py

app.py

@@ -1,22 +1,33 @@
+import os
+import warnings
+
+# Set environment variables BEFORE any imports to prevent CUDA initialization
+os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Hide CUDA during startup
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"  # For debugging
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
 import gradio as gr
 import numpy as np
-import os
 import spaces
 from huggingface_hub import login
+
+# These imports should now work without CUDA errors
 from maliba_ai.tts.inference import BambaraTTSInference
 from maliba_ai.config.speakers import Adame, Moussa, Bourama, Modibo, Seydou
 
-
 hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-
+# Initialize TTS model (this will use CPU during startup)
 print("Loading Bambara TTS model...")
 tts = BambaraTTSInference()
 print("Model loaded successfully!")
 
-
 SPEAKERS = {
     "Adame": Adame,
     "Moussa": Moussa, 
@@ -48,7 +59,13 @@ def generate_speech(text, speaker_name, use_advanced, temperature, top_k, top_p,
         return None, "Please enter some Bambara text."
     
     try:
-
+        # Re-enable CUDA for GPU context
+        import torch
+        if torch.cuda.is_available():
+            # Remove CUDA visibility restriction for GPU execution
+            if "CUDA_VISIBLE_DEVICES" in os.environ:
+                os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        
         speaker = SPEAKERS[speaker_name]
         
         if use_advanced:
@@ -79,18 +96,13 @@ def generate_speech(text, speaker_name, use_advanced, temperature, top_k, top_p,
     except Exception as e:
         return None, f"❌ Error: {str(e)}"
 
-
 examples = [
-
     ["Aw ni ce", "Adame"],
     ["I ni ce", "Moussa"],
     ["Aw ni tile", "Bourama"],
     ["I ka kene wa?", "Modibo"],
     ["Ala ka Mali suma", "Adame"],
-    
     ["sigikafɔ kɔnɔ jamanaw ni ɲɔgɔn cɛ, olu ye a haminankow ye, wa o ko ninnu ka kan ka kɛ sariya ani tilennenya kɔnɔ", "Seydou"],
-
-
     ["Aw ni ce. Ne tɔgɔ ye Kaya Magan. Aw Sanbe Sanbe.", "Moussa"],
     ["An dɔlakelen bɛ masike bilenman don ka tɔw gɛn.", "Bourama"],
     ["Aw ni ce. Seidu bɛ aw fo wa aw ka yafa a ma, ka da a kan tuma dɔw la kow ka can.", "Modibo"],
@@ -104,6 +116,8 @@ with gr.Blocks(title="Bambara TTS - EXPERIMENTAL", theme=gr.themes.Soft()) as de
     Convert Bambara text to speech using AI. This model is currently experimental.
     
     **Bambara** is spoken by millions of people in Mali and West Africa.
+    
+    ⚡ **Note**: Model loads on CPU during startup, then uses GPU for generation.
     """)
     
     with gr.Row():
@@ -126,7 +140,6 @@ with gr.Blocks(title="Bambara TTS - EXPERIMENTAL", theme=gr.themes.Soft()) as de
             generate_btn = gr.Button("🎵 Generate Speech", variant="primary", size="lg")
             
         with gr.Column(scale=1):
-
             use_advanced = gr.Checkbox(
                 label="⚙️ Use Advanced Settings", 
                 value=False,
@@ -203,6 +216,8 @@ with gr.Blocks(title="Bambara TTS - EXPERIMENTAL", theme=gr.themes.Soft()) as de
         gr.Markdown("""
         **⚠️ This is an experimental Bambara TTS model.**
         
+        The model loads on CPU during startup to avoid CUDA initialization errors,
+        then switches to GPU during speech generation for optimal performance.
         """)
     
     def toggle_advanced(use_adv):