Update app.py

app.py

@@ -1,17 +1,21 @@
 import os
 
+# Set environment variables BEFORE any imports
 os.environ["TORCHDYNAMO_DISABLE"] = "1"
 os.environ["TORCH_COMPILE_DISABLE"] = "1"
 os.environ["PYTORCH_DISABLE_CUDNN_BENCHMARK"] = "1"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+# Set CUDA environment to help with unsloth GPU detection
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Force GPU visibility
+os.environ["FORCE_CUDA"] = "1"            # Force CUDA usage
+
 import torch
 import gradio as gr
 import numpy as np
 import spaces
 import logging
 from huggingface_hub import login
-import threading
 import time
 
 torch._dynamo.config.disable = True
@@ -24,75 +28,81 @@ hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-
-_tts_model = None
-_speakers_dict = None
-_model_initialized = False
-_initialization_in_progress = False
+# Check GPU availability
+if torch.cuda.is_available():
+    device = "cuda"
+    logger.info("Using CUDA for inference.")
+elif torch.backends.mps.is_available():
+    device = "mps"
+    logger.info("Using MPS for inference.")
+else:
+    device = "cpu"
+    logger.info("Using CPU for inference.")
 
 def get_speakers_dict():
-    """Get speakers dictionary - moved to function to avoid import issues"""
+    """Get speakers dictionary using the new package structure"""
     try:
-        from maliba_ai.config.speakers import Adame, Moussa, Bourama, Modibo, Seydou
+        from maliba_ai.config.settings import Speakers
         return {
-            "Adama": Adame,
-            "Moussa": Moussa, 
-            "Bourama": Bourama,
-            "Modibo": Modibo,
-            "Seydou": Seydou
+            "Adama": Speakers.Adama,
+            "Moussa": Speakers.Moussa, 
+            "Bourama": Speakers.Bourama,
+            "Modibo": Speakers.Modibo,
+            "Seydou": Speakers.Seydou,
+            "Amadou": Speakers.Amadou,
+            "Bakary": Speakers.Bakary,
+            "Ngolo": Speakers.Ngolo,
+            "Ibrahima": Speakers.Ibrahima,
+            "Amara": Speakers.Amara
         }
     except Exception as e:
-        logger.error(f"Failed to import speakers: {e}")
-        return {}
-
-@spaces.GPU()
-def initialize_model_once():
-    global _tts_model, _speakers_dict, _model_initialized, _initialization_in_progress
-    
-    if _model_initialized:
-        logger.info("Model already initialized, returning existing instance")
-        return _tts_model, _speakers_dict
-    
-    if _initialization_in_progress:
-        logger.info("Initialization already in progress, waiting...")
+        logger.error(f"Failed to import all speakers: {e}")
+        # Fallback to core speakers only
+        try:
+            from maliba_ai.config.settings import Speakers
+            return {
+                "Adama": Speakers.Adama,
+                "Moussa": Speakers.Moussa, 
+                "Bourama": Speakers.Bourama,
+                "Modibo": Speakers.Modibo,
+                "Seydou": Speakers.Seydou
+            }
+        except:
+            logger.error("Failed to import even core speakers")
+            return {}
 
-        for _ in range(50): 
-            time.sleep(0.1)
-            if _model_initialized:
-                return _tts_model, _speakers_dict
- 
-    
-    _initialization_in_progress = True
-    
+def initialize_tts_model():
+    """Initialize TTS model globally - similar to ASR space pattern"""
     try:
-        logger.info("Initializing Bambara TTS model...")
+        logger.info("Initializing Bambara TTS model globally...")
         start_time = time.time()
         
-        from maliba_ai.tts.inference import BambaraTTSInference
+        # Import and initialize the TTS model
+        from maliba_ai.tts import BambaraTTSInference
         
+        # Initialize model 
         model = BambaraTTSInference()
-        speakers = get_speakers_dict()
-        
-        if not speakers:
-            raise ValueError("Failed to load speakers dictionary")
-        
-        _tts_model = model
-        _speakers_dict = speakers
-        _model_initialized = True
         
         elapsed = time.time() - start_time
-        logger.info(f"Model initialized successfully in {elapsed:.2f} seconds!")
+        logger.info(f"TTS Model initialized successfully in {elapsed:.2f} seconds!")
         
-        return _tts_model, _speakers_dict
+        return model
         
     except Exception as e:
-        logger.error(f"Failed to initialize model: {e}")
-        _initialization_in_progress = False
-        raise e
-    finally:
-        _initialization_in_progress = False
+        logger.error(f"Failed to initialize TTS model: {e}")
+        logger.info("Model will be initialized on first request instead")
+        return None
+
+# Initialize speakers dictionary
+speakers_dict = get_speakers_dict()
+logger.info(f"Available speakers: {list(speakers_dict.keys())}")
+
+# Try to initialize model globally (like ASR space)
+# If it fails due to GPU detection, it will be None and we'll init on first request
+tts_model = initialize_tts_model()
 
 def validate_inputs(text, temperature, top_k, top_p, max_tokens):
+    """Validate user inputs"""
     if not text or not text.strip():
         return False, "Please enter some Bambara text."
     
@@ -105,32 +115,44 @@ def validate_inputs(text, temperature, top_k, top_p, max_tokens):
     if not (0.1 <= top_p <= 1.0):
         return False, "Top-P must be between 0.1 and 1.0"
     
+    if len(text.strip()) > 1000:
+        return False, "Text is too long. Please use shorter text (max 1000 characters)."
+    
     return True, ""
 
 @spaces.GPU()
 def generate_speech(text, speaker_name, use_advanced, temperature, top_k, top_p, max_tokens):
+    """Generate speech - with fallback initialization if global init failed"""
+    global tts_model
+    
     if not text.strip():
         return None, "Please enter some Bambara text."
     
     try:
-        tts, speakers = initialize_model_once()
+        # If global initialization failed, try to initialize here with GPU decorator
+        if tts_model is None:
+            logger.info("Global model initialization failed, initializing with GPU decorator...")
+            from maliba_ai.tts import BambaraTTSInference
+            tts_model = BambaraTTSInference()
+            logger.info("Model initialized successfully with GPU decorator!")
         
-        if not tts or not speakers:
-            return None, "❌ Model not properly initialized"
+        if not speakers_dict:
+            return None, "❌ Speakers not properly loaded"
         
-        if speaker_name not in speakers:
-            available_speakers = list(speakers.keys())
+        if speaker_name not in speakers_dict:
+            available_speakers = list(speakers_dict.keys())
             return None, f"❌ Speaker '{speaker_name}' not found. Available: {available_speakers}"
         
-        speaker = speakers[speaker_name]
-        logger.info(f"Using speaker: {speaker_name}")
+        speaker = speakers_dict[speaker_name]
+        logger.info(f"Generating speech with speaker: {speaker_name}")
         
+        # Validate inputs if using advanced settings
         if use_advanced:
             is_valid, error_msg = validate_inputs(text, temperature, top_k, top_p, max_tokens)
             if not is_valid:
                 return None, f"❌ {error_msg}"
             
-            waveform = tts.generate_speech(
+            waveform = tts_model.generate_speech(
                 text=text.strip(),
                 speaker_id=speaker,
                 temperature=temperature,
@@ -139,76 +161,104 @@ def generate_speech(text, speaker_name, use_advanced, temperature, top_k, top_p,
                 max_new_audio_tokens=int(max_tokens)
             )
         else:
-            waveform = tts.generate_speech(
+            # Use default settings
+            waveform = tts_model.generate_speech(
                 text=text.strip(),
                 speaker_id=speaker
             )
         
         if waveform is None or waveform.size == 0:
-            return None, "Failed to generate audio. Please try again."
+            return None, "❌ Failed to generate audio. Please try again with different text."
+        
+        # Ensure waveform is in correct format
+        if isinstance(waveform, torch.Tensor):
+            waveform = waveform.cpu().numpy()
+        
+        # Normalize audio to prevent clipping
+        if np.max(np.abs(waveform)) > 0:
+            waveform = waveform / np.max(np.abs(waveform)) * 0.9
         
         sample_rate = 16000
         return (sample_rate, waveform), f"✅ Audio generated successfully for speaker {speaker_name}"
         
     except Exception as e:
-        logger.error(f"Speech generation failed: {e}")
+        logger.error(f"Speech generation failed: {e}", exc_info=True)
         return None, f"❌ Error: {str(e)}"
 
-SPEAKER_NAMES = ["Adama", "Moussa", "Bourama", "Modibo", "Seydou"]
+# Get available speakers for dropdown
+SPEAKER_NAMES = list(speakers_dict.keys()) if speakers_dict else ["Adama", "Moussa", "Bourama", "Modibo", "Seydou"]
+
 
 examples = [
-    ["Aw ni ce", "Adama"],
-    ["Mali bɛna diya kɔsɛbɛ,  ka a da a kan baara bɛ ka kɛ.", "Moussa"],
-    ["Ne bɛ se ka sɛbɛnni yɛlɛma ka kɛ kuma ye", "Bourama"],
-    ["I ka kɛnɛ wa?", "Modibo"],
-    ["Lakɔli karamɔgɔw tun tɛ ka se ka sɛbɛnni kɛ ka ɲɛ walanda kan wa denmisɛnw tun tɛ ka se ka o sɛbɛnni ninnu ye,  kuma tɛ ka u kalan. Denmisɛnw kɛra kunfinw ye.", "Adama"],
-    ["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 Adama. Awɔ,  ne ye maliden de ye. Aw Sanbɛ Sanbɛ. San min tɛ ɲinan ye, an bɛɛ ka jɛ ka o seli ɲɔgɔn fɛ,  hɛɛrɛ  ni lafiya la. Ala ka Mali suma. Ala ka Mali yiriwa. Ala ka Mali taa ɲɛ. Ala ka an ka seliw caya. Ala ka yafa an bɛɛ ma.", "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"],
+    ["Aw ni ce", "Adama"],  
+    ["Mali bɛna diya kɔsɛbɛ, ka a da a kan baara bɛ ka kɛ.", "Bakary"],  
+    ["Ne bɛ se ka sɛbɛnni yɛlɛma ka kɛ kuma ye", "Moussa"],  
+    ["I ka kɛnɛ wa?", "Ngolo"],  
+    ["Lakɔli karamɔgɔw tun tɛ ka se ka sɛbɛnni kɛ ka ɲɛ walanda kan wa denmisɛnw tun tɛ ka se ka o sɛbɛnni ninnu ye, kuma tɛ ka u kalan. Denmisɛnw kɛra kunfinw ye.", "Bourama"],  
+    ["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ɔ.", "Ibrahima"],  
+    ["Aw ni ce. Ne tɔgɔ ye Adama. Awɔ, ne ye maliden de ye. Aw Sanbɛ Sanbɛ. San min tɛ ɲinan ye, an bɛɛ ka jɛ ka o seli ɲɔgɔn fɛ, hɛɛrɛ ni lafiya la. Ala ka Mali suma. Ala ka Mali yiriwa. Ala ka Mali taa ɲɛ. Ala ka an ka seliw caya. Ala ka yafa an bɛɛ ma.", "Amara"],  
+    ["An dɔlakelen bɛ masike bilenman don ka tɔw gɛn.", "Modibo"],  
+    ["Aw ni ce. Seidu bɛ aw fo wa aw ka yafa a ma, ka da a kan tuma dɔw la kow ka can.", "Amadou"],  
+    ["Bamanankan ye kan ɲuman ye", "Seydou"],  
 ]
 
 def build_interface():
     """Build the Gradio interface for Bambara TTS"""
     
-    with gr.Blocks(title="Bambara TTS - EXPERIMENTAL") as demo:
-        gr.Markdown("""
-        # 🎤 Bambara Text-to-Speech ⚠️ EXPERIMENTAL
-        
-        **Powered by MALIBA-AI**
+    with gr.Blocks(
+        title="Bambara TTS - MALIBA-AI", 
+        theme=gr.themes.Soft(),
+        css="""
+        .main-header { text-align: center; margin-bottom: 2rem; }
+        .status-box { margin-top: 1rem; }
+        """
+    ) as demo:
         
-        Convert Bambara text to speech. This model is currently experimental.
-        
-        **Bambara** is spoken by millions of people in Mali and West Africa.
-        """)
+        with gr.Row():
+            gr.Markdown(f"""
+            # 🎤 Bambara Text-to-Speech
+            
+            **Powered by MALIBA-AI** | *First Open-Source Bambara TTS*
+            
+            Convert Bambara text to natural-sounding speech using our state-of-the-art neural TTS system.
+            
+            **Bambara** is spoken by millions of people in Mali and West Africa 🌍
+            
+            **Status**: {'✅ Model loaded' if tts_model is not None else '⏳ Model will load on first request'}
+            """, elem_classes=["main-header"])
 
         with gr.Row():
             with gr.Column(scale=2):
                 text_input = gr.Textbox(
                     label="📝 Bambara Text",
-                    placeholder="Type your Bambara text here...",
-                    lines=3,
-                    max_lines=10,
+                    placeholder="I ni ce... (Type your Bambara text here)",
+                    lines=4,
+                    max_lines=8,
                     value="I ni ce"
                 )
                 
                 speaker_dropdown = gr.Dropdown(
                     choices=SPEAKER_NAMES,
-                    value="Adama",
-                    label="🗣️ Speaker Voice"
+                    value=SPEAKER_NAMES[0] if SPEAKER_NAMES else "Bourama",  # Default to most stable speaker
+                    label="🗣️ Speaker Voice",
+                    info=f"Choose from {len(SPEAKER_NAMES)} authentic voices (Bourama recommended for best quality)"
                 )
                 
-                generate_btn = gr.Button("🎵 Generate Speech", variant="primary", size="lg")
+                generate_btn = gr.Button(
+                    "🎵 Generate Speech", 
+                    variant="primary", 
+                    size="lg"
+                )
                 
             with gr.Column(scale=1):
                 use_advanced = gr.Checkbox(
-                    label="⚙️ Use Advanced Settings", 
+                    label="⚙️ Advanced Settings", 
                     value=False,
-                    info="Enable to customize generation parameters"
+                    info="Customize generation parameters"
                 )
                 
                 with gr.Group(visible=False) as advanced_group:
-                    gr.Markdown("**Advanced Parameters:**")
+                    gr.Markdown("**🔧 Advanced Parameters:**")
                     
                     temperature = gr.Slider(
                         minimum=0.1, 
@@ -216,7 +266,7 @@ def build_interface():
                         value=0.8, 
                         step=0.1,
                         label="Temperature",
-                        info="Higher = more varied"
+                        info="Higher = more varied speech"
                     )
                     
                     top_k = gr.Slider(
@@ -224,7 +274,8 @@ def build_interface():
                         maximum=100, 
                         value=50, 
                         step=5,
-                        label="Top-K"
+                        label="Top-K",
+                        info="Vocabulary selection size"
                     )
                     
                     top_p = gr.Slider(
@@ -232,7 +283,8 @@ def build_interface():
                         maximum=1.0, 
                         value=0.9, 
                         step=0.05,
-                        label="Top-P"
+                        label="Top-P",
+                        info="Nucleus sampling threshold"
                     )
                     
                     max_tokens = gr.Slider(
@@ -240,7 +292,8 @@ def build_interface():
                         maximum=4096, 
                         value=2048, 
                         step=256,
-                        label="Max Length"
+                        label="Max Audio Length",
+                        info="Maximum audio duration"
                     )
         
         gr.Markdown("### 🔊 Generated Audio")
@@ -248,39 +301,79 @@ def build_interface():
         audio_output = gr.Audio(
             label="Generated Speech",
             type="numpy",
-            interactive=False
+            interactive=False,
+            show_download_button=True
         )
             
         status_output = gr.Textbox(
             label="Status",
             interactive=False,
             show_label=False,
-            container=False
+            container=False,
+            elem_classes=["status-box"]
         )
         
-        with gr.Accordion("Try These Examples", open=True):
+        with gr.Accordion("📚 Try These Examples", open=True):
             def load_example(text, speaker):
                 return text, speaker, False, 0.8, 50, 0.9, 2048
             
-            gr.Markdown("**Click any example below:**")
+            gr.Markdown("**Click any example below to try it:**")
             
-            for i, (text, speaker) in enumerate(examples):
-                btn = gr.Button(f"{text[:30]}{'...' if len(text) > 30 else ''}", size="sm")
-                btn.click(
-                    fn=lambda t=text, s=speaker: load_example(t, s),
-                    outputs=[text_input, speaker_dropdown, use_advanced, temperature, top_k, top_p, max_tokens]
-                )
+            with gr.Row():
+                for i, (text, speaker) in enumerate(examples[:5]):
+                    btn = gr.Button(
+                        f"🔹 {text[:25]}{'...' if len(text) > 25 else ''}", 
+                        size="sm"
+                    )
+                    btn.click(
+                        fn=lambda t=text, s=speaker: load_example(t, s),
+                        outputs=[text_input, speaker_dropdown, use_advanced, temperature, top_k, top_p, max_tokens]
+                    )
+            
+            with gr.Row():
+                for i, (text, speaker) in enumerate(examples[5:]):
+                    btn = gr.Button(
+                        f"🔹 {text[:25]}{'...' if len(text) > 25 else ''}", 
+                        size="sm"
+                    )
+                    btn.click(
+                        fn=lambda t=text, s=speaker: load_example(t, s),
+                        outputs=[text_input, speaker_dropdown, use_advanced, temperature, top_k, top_p, max_tokens]
+                    )
         
-        with gr.Accordion("About", open=False):
-            gr.Markdown("""
-            **⚠️ This is an experimental Bambara TTS model.**
-            - **Languages**: Bambara (bm)  
-            - **Speakers**: 5 different voice options
-            - **Sample Rate**: 16kHz
+        with gr.Accordion("ℹ️ About", open=False):
+            gr.Markdown(f"""
+            ## About MALIBA-AI Bambara TTS
+            
+            - **🎯 Purpose**: First open-source Text-to-Speech system for Bambara language
+            - **🗣️ Speakers**: {len(SPEAKER_NAMES)} different authentic voices  
+            - **🔊 Quality**: 16kHz neural speech synthesis
+            - **⚡ Performance**: Optimized for real-time generation
+            - **📱 Usage**: Educational, accessibility, and cultural preservation
+            
+            ### 🎭 Speaker Characteristics:
+            
+            - **Bourama**: Most stable and accurate (recommended)
+            - **Adama**: Natural conversational tone
+            - **Moussa**: Clear pronunciation for educational content
+            - **Modibo**: Expressive delivery for storytelling
+            - **Seydou**: Balanced characteristics for general use
+            - **Amadou**: Warm and friendly voice
+            - **Bakary**: Deep, authoritative tone
+            - **Ngolo**: Youthful and energetic
+            - **Ibrahima**: Calm and measured delivery
+            - **Amara**: Melodic and smooth
             
-            **Model loads once on first request and stays in memory**
+            **Model Architecture**: Built on state-of-the-art neural TTS with Bambara-specific optimizations
+            
+            **License**: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 (CC BY-NC-SA 4.0)
+            
+            ---
+            
+            **MALIBA-AI Mission**: Ensuring no Malian is left behind by technological advances 🇲🇱
             """)
         
+        # Event handlers
         def toggle_advanced(use_adv):
             return gr.Group(visible=use_adv)
         
@@ -290,6 +383,7 @@ def build_interface():
             outputs=[advanced_group]
         )
         
+        # Generate speech on button click
         generate_btn.click(
             fn=generate_speech,
             inputs=[text_input, speaker_dropdown, use_advanced, temperature, top_k, top_p, max_tokens],
@@ -297,6 +391,7 @@ def build_interface():
             show_progress=True
         )
         
+        # Generate speech on Enter key
         text_input.submit(
             fn=generate_speech,
             inputs=[text_input, speaker_dropdown, use_advanced, temperature, top_k, top_p, max_tokens],
@@ -308,17 +403,20 @@ def build_interface():
 
 def main():
     """Main function to launch the Gradio interface"""
-    logger.info("Starting Bambara TTS Gradio interface.")
+    logger.info("Starting MALIBA-AI Bambara TTS Gradio interface...")
     
-    # DO NOT preload - let it initialize on first request only
+    # Build interface
     interface = build_interface()
+    
+    # Launch interface
     interface.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False
+        share=False,
+        show_error=True
     )
     
-    logger.info("Gradio interface launched successfully.")
+    logger.info("Gradio interface launched successfully!")
 
 if __name__ == "__main__":
     main()