Update app.py

app.py

@@ -6,9 +6,39 @@ os.environ["TORCH_COMPILE_DISABLE"] = "1"
 os.environ["PYTORCH_DISABLE_CUDNN_BENCHMARK"] = "1"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# Disable Unsloth optimizations that cause issues in ZeroGPU
-os.environ["UNSLOTH_DISABLE"] = "1"
-os.environ["DISABLE_UNSLOTH"] = "1"
+# Check if we're in ZeroGPU or similar restricted environment
+def is_restricted_environment():
+    return (
+        os.getenv("ZERO_GPU") or 
+        "zero" in str(os.getenv("SPACE_ID", "")).lower() or
+        os.getenv("SPACES_ZERO_GPU") or
+        "spaces" in str(os.getenv("HOSTNAME", "")).lower()
+    )
+
+# Disable Unsloth optimizations in restricted environments
+if is_restricted_environment():
+    os.environ["UNSLOTH_DISABLE"] = "1"
+    os.environ["DISABLE_UNSLOTH"] = "1" 
+    os.environ["UNSLOTH_IGNORE_ERRORS"] = "1"
+    os.environ["UNSLOTH_NO_COMPILE"] = "1"
+    print("🚀 ZeroGPU detected - Unsloth optimizations disabled for compatibility")
+else:
+    print("🔧 Local environment detected - Unsloth optimizations enabled")
+
+# Check if we're in ZeroGPU or similar restricted environment
+def is_restricted_environment():
+    return (
+        os.getenv("ZERO_GPU") or 
+        "zero" in str(os.getenv("SPACE_ID", "")).lower() or
+        os.getenv("SPACES_ZERO_GPU")
+    )
+
+# Disable Unsloth optimizations in restricted environments
+if is_restricted_environment():
+    os.environ["UNSLOTH_DISABLE"] = "1"
+    os.environ["DISABLE_UNSLOTH"] = "1"
+    os.environ["UNSLOTH_IGNORE_ERRORS"] = "1"
+    print("🚀 ZeroGPU detected - Unsloth optimizations disabled for compatibility")
 
 import torch
 import gradio as gr
@@ -64,7 +94,7 @@ def get_speakers_dict():
 
 @spaces.GPU()
 def initialize_model_once():
-    """Initialize model exactly like your old working version"""
+    """Initialize model with retry logic for Unsloth compilation issues"""
     global _tts_model, _speakers_dict, _model_initialized, _initialization_in_progress
     
     if _model_initialized:
@@ -80,34 +110,44 @@ def initialize_model_once():
     
     _initialization_in_progress = True
     
-    try:
-        logger.info("Initializing Bambara TTS model...")
-        start_time = time.time()
-        
-        # Use the correct import path
-        from maliba_ai.tts.inference import BambaraTTSInference
-        
-        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!")
-        
-        return _tts_model, _speakers_dict
-        
-    except Exception as e:
-        logger.error(f"Failed to initialize model: {e}")
-        _initialization_in_progress = False
-        raise e
+    # Retry logic for Unsloth compilation issues
+    max_retries = 2
+    for attempt in range(max_retries + 1):
+        try:
+            logger.info(f"Initializing Bambara TTS model... (attempt {attempt + 1})")
+            start_time = time.time()
+            
+            # Use the correct import path
+            from maliba_ai.tts.inference import BambaraTTSInference
+            
+            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!")
+            
+            return _tts_model, _speakers_dict
+            
+        except Exception as e:
+            error_msg = str(e)
+            if "unsloth_compiled_module" in error_msg and attempt < max_retries:
+                logger.warning(f"Unsloth compilation failed (attempt {attempt + 1}/{max_retries + 1}), retrying...")
+                time.sleep(2)  # Brief delay before retry
+                continue
+            else:
+                logger.error(f"Failed to initialize model: {e}")
+                _initialization_in_progress = False
+                raise e
     finally:
-        _initialization_in_progress = False
+        if not _model_initialized:
+            _initialization_in_progress = False
 
 def validate_inputs(text, temperature, top_k, top_p, max_tokens):
     """Same validation as your old version"""