Update app.py

app.py

@@ -1,22 +1,17 @@
 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 (only if not ZeroGPU)
-if not os.getenv("ZERO_GPU"):
-    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
@@ -29,23 +24,16 @@ hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-# Check GPU availability (but don't initialize CUDA yet in ZeroGPU)
-if os.getenv("ZERO_GPU"):
-    device = "cuda"  # Assume CUDA in ZeroGPU
-    logger.info("ZeroGPU environment detected - CUDA will be available in decorated functions")
-elif 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.")
+# Global variables for model caching (like your old working version)
+_tts_model = None
+_speakers_dict = None
+_model_initialized = False
+_initialization_in_progress = False
 
 def get_speakers_dict():
-    """Get speakers dictionary using the new package structure"""
+    """Get speakers dictionary using the correct import structure"""
     try:
+        # Try new structure first
         from maliba_ai.config.settings import Speakers
         return {
             "Adama": Speakers.Adama,
@@ -60,62 +48,70 @@ def get_speakers_dict():
             "Amara": Speakers.Amara
         }
     except Exception as e:
-        logger.error(f"Failed to import all speakers: {e}")
-        # Fallback to core speakers only
+        logger.error(f"Failed to import from settings: {e}")
+        # Fallback to old structure (like your working version)
         try:
-            from maliba_ai.config.settings import Speakers
+            from maliba_ai.config.speakers import Adame, Moussa, Bourama, Modibo, Seydou
             return {
-                "Adama": Speakers.Adama,
-                "Moussa": Speakers.Moussa, 
-                "Bourama": Speakers.Bourama,
-                "Modibo": Speakers.Modibo,
-                "Seydou": Speakers.Seydou
+                "Adama": Adame,
+                "Moussa": Moussa, 
+                "Bourama": Bourama,
+                "Modibo": Modibo,
+                "Seydou": Seydou
             }
-        except:
-            logger.error("Failed to import even core speakers")
+        except Exception as e2:
+            logger.error(f"Failed to import speakers: {e2}")
             return {}
 
-def initialize_tts_model():
-    """Initialize TTS model globally - only if we're not in ZeroGPU environment"""
+@spaces.GPU()
+def initialize_model_once():
+    """Initialize model exactly like your old working version"""
+    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...")
+        for _ in range(50): 
+            time.sleep(0.1)
+            if _model_initialized:
+                return _tts_model, _speakers_dict
+    
+    _initialization_in_progress = True
+    
     try:
-        # Check if we're in ZeroGPU environment - don't initialize globally
-        if os.getenv("ZERO_GPU") or "zero" in str(os.getenv("SPACE_ID", "")).lower():
-            logger.info("ZeroGPU environment detected - skipping global initialization")
-            return None
-            
-        # Only try global init if CUDA is actually available and initialized
-        if not torch.cuda.is_available():
-            logger.info("CUDA not available - skipping global initialization")
-            return None
-            
-        logger.info("Attempting global TTS model initialization...")
+        logger.info("Initializing Bambara TTS model...")
         start_time = time.time()
         
-        # Import and initialize the TTS model
-        from maliba_ai.tts.inference import BambaraTTSInference
+        # Use the same import as your old working version
+        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"TTS Model initialized successfully in {elapsed:.2f} seconds!")
+        logger.info(f"Model initialized successfully in {elapsed:.2f} seconds!")
         
-        return model
+        return _tts_model, _speakers_dict
         
     except Exception as e:
-        logger.error(f"Failed to initialize TTS model globally: {e}")
-        logger.info("Model will be initialized on first request with GPU decorator")
-        return None
-
-# Initialize speakers dictionary (this doesn't require GPU)
-speakers_dict = get_speakers_dict()
-logger.info(f"Available speakers: {list(speakers_dict.keys())}")
-
-# Try to initialize model globally only if not in ZeroGPU environment
-tts_model = initialize_tts_model()
+        logger.error(f"Failed to initialize model: {e}")
+        _initialization_in_progress = False
+        raise e
+    finally:
+        _initialization_in_progress = False
 
 def validate_inputs(text, temperature, top_k, top_p, max_tokens):
-    """Validate user inputs"""
+    """Same validation as your old version"""
     if not text or not text.strip():
         return False, "Please enter some Bambara text."
     
@@ -128,44 +124,33 @@ 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
-    
+    """Generate speech - exactly like your old working version"""
     if not text.strip():
         return None, "Please enter some Bambara text."
     
     try:
-        # 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!")
+        tts, speakers = initialize_model_once()
         
-        if not speakers_dict:
-            return None, "❌ Speakers not properly loaded"
+        if not tts or not speakers:
+            return None, "❌ Model not properly initialized"
         
-        if speaker_name not in speakers_dict:
-            available_speakers = list(speakers_dict.keys())
+        if speaker_name not in speakers:
+            available_speakers = list(speakers.keys())
             return None, f"❌ Speaker '{speaker_name}' not found. Available: {available_speakers}"
         
-        speaker = speakers_dict[speaker_name]
-        logger.info(f"Generating speech with speaker: {speaker_name}")
+        speaker = speakers[speaker_name]
+        logger.info(f"Using 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_model.generate_speech(
+            waveform = tts.generate_speech(
                 text=text.strip(),
                 speaker_id=speaker,
                 temperature=temperature,
@@ -174,104 +159,85 @@ def generate_speech(text, speaker_name, use_advanced, temperature, top_k, top_p,
                 max_new_audio_tokens=int(max_tokens)
             )
         else:
-            # Use default settings
-            waveform = tts_model.generate_speech(
+            waveform = tts.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 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
+            return None, "Failed to generate audio. Please try again."
         
         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}", exc_info=True)
+        logger.error(f"Speech generation failed: {e}")
         return None, f"❌ Error: {str(e)}"
 
-# Get available speakers for dropdown
-SPEAKER_NAMES = list(speakers_dict.keys()) if speakers_dict else ["Adama", "Moussa", "Bourama", "Modibo", "Seydou"]
+# Use available speakers (try to get 10, fallback to 5)
+def get_speaker_names():
+    speakers = get_speakers_dict()
+    if speakers:
+        return list(speakers.keys())
+    return ["Adama", "Moussa", "Bourama", "Modibo", "Seydou"]
+
+SPEAKER_NAMES = get_speaker_names()
 
 # Examples with variety of lengths and speakers matched to content
 examples = [
     ["Aw ni ce", "Adama"],  # Natural conversational greeting
-    ["Mali bɛna diya kɔsɛbɛ, ka a da a kan baara bɛ ka kɛ.", "Bakary"],  # Authoritative tone for serious topic
-    ["Ne bɛ se ka sɛbɛnni yɛlɛma ka kɛ kuma ye", "Moussa"],  # Clear pronunciation for education
-    ["I ka kɛnɛ wa?", "Ngolo"],  # Youthful energy for casual question
-    ["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"],  # Most stable for long educational text
-    ["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"],  # Calm and measured for formal text
-    ["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"],  # Melodic and smooth for heartfelt message
-    ["An dɔlakelen bɛ masike bilenman don ka tɔw gɛn.", "Modibo"],  # Expressive delivery for dramatic statement
-    ["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"],  # Warm and friendly greeting
-    ["Bamanankan ye kan ɲuman ye", "Seydou"],  # Balanced characteristics for simple statement
+    ["Mali bɛna diya kɔsɛbɛ, ka a da a kan baara bɛ ka kɛ.", "Bakary" if "Bakary" in SPEAKER_NAMES else "Moussa"],  
+    ["Ne bɛ se ka sɛbɛnni yɛlɛma ka kɛ kuma ye", "Moussa"],  
+    ["I ka kɛnɛ wa?", "Ngolo" if "Ngolo" in SPEAKER_NAMES else "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.", "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" if "Ibrahima" in SPEAKER_NAMES else "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.", "Amara" if "Amara" in SPEAKER_NAMES else "Moussa"],  
+    ["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" if "Amadou" in SPEAKER_NAMES else "Modibo"],
 ]
 
 def build_interface():
-    """Build the Gradio interface for Bambara TTS"""
+    """Build the Gradio interface - simplified like your old working version"""
     
-    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:
+    with gr.Blocks(title="Bambara TTS - MALIBA-AI") as demo:
+        gr.Markdown("""
+        # 🎤 Bambara Text-to-Speech
         
-        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 pre-loaded' if tts_model is not None else '⏳ Model loads on first request (ZeroGPU optimized)'}
-            """, elem_classes=["main-header"])
+        **Powered by MALIBA-AI**
+        
+        Convert Bambara text to speech using our state-of-the-art TTS model.
+        
+        **Bambara** is spoken by millions of people in Mali and West Africa.
+        """)
 
         with gr.Row():
             with gr.Column(scale=2):
                 text_input = gr.Textbox(
                     label="📝 Bambara Text",
-                    placeholder="I ni ce... (Type your Bambara text here)",
-                    lines=4,
-                    max_lines=8,
+                    placeholder="Type your Bambara text here...",
+                    lines=3,
+                    max_lines=10,
                     value="I ni ce"
                 )
                 
                 speaker_dropdown = gr.Dropdown(
                     choices=SPEAKER_NAMES,
-                    value=SPEAKER_NAMES[0] if SPEAKER_NAMES else "Bourama",  # Default to most stable speaker
+                    value="Bourama" if "Bourama" in SPEAKER_NAMES else SPEAKER_NAMES[0],
                     label="🗣️ Speaker Voice",
-                    info=f"Choose from {len(SPEAKER_NAMES)} authentic voices (Bourama recommended for best quality)"
+                    info=f"Choose from {len(SPEAKER_NAMES)} authentic voices"
                 )
                 
-                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="⚙️ Advanced Settings", 
+                    label="⚙️ Use Advanced Settings", 
                     value=False,
-                    info="Customize generation parameters"
+                    info="Enable to 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, 
@@ -279,7 +245,7 @@ def build_interface():
                         value=0.8, 
                         step=0.1,
                         label="Temperature",
-                        info="Higher = more varied speech"
+                        info="Higher = more varied"
                     )
                     
                     top_k = gr.Slider(
@@ -287,8 +253,7 @@ def build_interface():
                         maximum=100, 
                         value=50, 
                         step=5,
-                        label="Top-K",
-                        info="Vocabulary selection size"
+                        label="Top-K"
                     )
                     
                     top_p = gr.Slider(
@@ -296,8 +261,7 @@ def build_interface():
                         maximum=1.0, 
                         value=0.9, 
                         step=0.05,
-                        label="Top-P",
-                        info="Nucleus sampling threshold"
+                        label="Top-P"
                     )
                     
                     max_tokens = gr.Slider(
@@ -305,8 +269,7 @@ def build_interface():
                         maximum=4096, 
                         value=2048, 
                         step=256,
-                        label="Max Audio Length",
-                        info="Maximum audio duration"
+                        label="Max Length"
                     )
         
         gr.Markdown("### 🔊 Generated Audio")
@@ -314,70 +277,41 @@ def build_interface():
         audio_output = gr.Audio(
             label="Generated Speech",
             type="numpy",
-            interactive=False,
-            show_download_button=True
+            interactive=False
         )
             
         status_output = gr.Textbox(
             label="Status",
             interactive=False,
             show_label=False,
-            container=False,
-            elem_classes=["status-box"]
+            container=False
         )
         
-        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 to try it:**")
+            gr.Markdown("**Click any example below:**")
             
-            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]
-                    )
+            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.Accordion("ℹ️ About", open=False):
+        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
+            - **⚡ Performance**: Model loads once and stays in memory
             - **📱 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 Architecture**: Built on state-of-the-art neural TTS with Bambara-specific optimizations
+            ### 🎭 Available Speakers:
+            {', '.join(SPEAKER_NAMES)}
             
             **License**: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 (CC BY-NC-SA 4.0)
             
@@ -386,7 +320,6 @@ def build_interface():
             **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)
         
@@ -396,7 +329,6 @@ 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],
@@ -404,7 +336,6 @@ 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],
@@ -416,20 +347,17 @@ def build_interface():
 
 def main():
     """Main function to launch the Gradio interface"""
-    logger.info("Starting MALIBA-AI Bambara TTS Gradio interface...")
+    logger.info("Starting Bambara TTS Gradio interface.")
     
-    # Build interface
+    # DO NOT preload - let it initialize on first request only (like your working version)
     interface = build_interface()
-    
-    # Launch interface
     interface.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,
-        show_error=True
+        share=False
     )
     
-    logger.info("Gradio interface launched successfully!")
+    logger.info("Gradio interface launched successfully.")
 
 if __name__ == "__main__":
     main()