Update app.py

app.py

@@ -12,7 +12,20 @@ from pathlib import Path
 from mira.model import MiraTTS
 
 MODEL = None
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Safe device detection with fallback
+def get_device():
+    """Safely detect available device."""
+    try:
+        if torch.cuda.is_available():
+            # Try to actually access CUDA to verify it works
+            torch.cuda.current_device()
+            return "cuda"
+    except Exception as e:
+        logging.warning(f"CUDA not available or driver error: {e}")
+    return "cpu"
+
+DEVICE = get_device()
 HISTORY_FILE = "generation_history.json"
 GENERATION_QUEUE = queue.Queue()
 PROCESSING_LOCK = threading.Lock()
@@ -67,18 +80,39 @@ def initialize_model(model_dir="YatharthS/MiraTTS", device=None):
     """Load the MiraTTS model once at the beginning."""
     global DEVICE
     if device:
-        DEVICE = device
+        # Verify the requested device is available
+        if device == "cuda":
+            try:
+                if not torch.cuda.is_available():
+                    logging.warning("CUDA requested but not available, falling back to CPU")
+                    DEVICE = "cpu"
+                else:
+                    torch.cuda.current_device()  # Test CUDA access
+                    DEVICE = device
+            except Exception as e:
+                logging.warning(f"CUDA test failed: {e}, falling back to CPU")
+                DEVICE = "cpu"
+        else:
+            DEVICE = device
     
     logging.info(f"Loading MiraTTS model from: {model_dir}")
     logging.info(f"Using device: {DEVICE}")
     
-    model = MiraTTS(model_dir)
-    
-    # Move model to appropriate device
-    if hasattr(model, 'to'):
-        model = model.to(DEVICE)
-    
-    return model
+    try:
+        model = MiraTTS(model_dir)
+        
+        # Move model to appropriate device
+        if hasattr(model, 'to') and DEVICE == "cuda":
+            try:
+                model = model.to(DEVICE)
+            except Exception as e:
+                logging.warning(f"Failed to move model to CUDA: {e}, using CPU")
+                DEVICE = "cpu"
+        
+        return model
+    except Exception as e:
+        logging.error(f"Error initializing model: {e}")
+        raise
 
 def generate_audio(text, prompt_audio_path):
     """Generate audio from text using MiraTTS with voice cloning."""
@@ -92,12 +126,25 @@ def generate_audio(text, prompt_audio_path):
         context_tokens = MODEL.encode_audio(prompt_audio_path)
         
         # Move context tokens to device if needed
-        if torch.is_tensor(context_tokens):
-            context_tokens = context_tokens.to(DEVICE)
+        if torch.is_tensor(context_tokens) and DEVICE == "cuda":
+            try:
+                context_tokens = context_tokens.to(DEVICE)
+            except Exception as e:
+                logging.warning(f"Failed to move tensors to CUDA: {e}")
         
-        # Generate audio
-        with torch.inference_mode() if DEVICE == "cpu" else torch.cuda.amp.autocast():
-            audio = MODEL.generate(text, context_tokens)
+        # Generate audio with appropriate context
+        try:
+            if DEVICE == "cpu":
+                with torch.inference_mode():
+                    audio = MODEL.generate(text, context_tokens)
+            else:
+                with torch.cuda.amp.autocast():
+                    audio = MODEL.generate(text, context_tokens)
+        except Exception as e:
+            # Fallback to simple generation if autocast fails
+            logging.warning(f"Autocast failed: {e}, using standard generation")
+            with torch.inference_mode():
+                audio = MODEL.generate(text, context_tokens)
         
         # Convert to numpy array if it's a tensor and handle dtype
         if torch.is_tensor(audio):
@@ -235,12 +282,25 @@ def voice_creation_callback(text, temperature, top_p, top_k, progress=gr.Progres
             # Generate audio with dtype conversion
             context_tokens = MODEL.encode_audio(default_audio)
             
-            # Move to device
-            if torch.is_tensor(context_tokens):
-                context_tokens = context_tokens.to(DEVICE)
+            # Move to device safely
+            if torch.is_tensor(context_tokens) and DEVICE == "cuda":
+                try:
+                    context_tokens = context_tokens.to(DEVICE)
+                except Exception as e:
+                    logging.warning(f"Failed to move tensors to CUDA: {e}")
             
-            with torch.inference_mode() if DEVICE == "cpu" else torch.cuda.amp.autocast():
-                audio = MODEL.generate(text, context_tokens)
+            try:
+                if DEVICE == "cpu":
+                    with torch.inference_mode():
+                        audio = MODEL.generate(text, context_tokens)
+                else:
+                    with torch.cuda.amp.autocast():
+                        audio = MODEL.generate(text, context_tokens)
+            except Exception as e:
+                # Fallback to simple generation
+                logging.warning(f"Autocast failed: {e}, using standard generation")
+                with torch.inference_mode():
+                    audio = MODEL.generate(text, context_tokens)
             
             # Handle tensor conversion and dtype
             if torch.is_tensor(audio):
@@ -331,7 +391,12 @@ def build_ui():
         # Device info
         device_info = f"🖥️ Running on: **{DEVICE.upper()}**"
         if DEVICE == "cuda":
-            device_info += f" (GPU: {torch.cuda.get_device_name(0)})"
+            try:
+                device_info += f" (GPU: {torch.cuda.get_device_name(0)})"
+            except:
+                device_info += " (GPU)"
+        else:
+            device_info += " (CPU mode - slower but works without GPU)"
         gr.Markdown(device_info)
         
         # Description
@@ -558,7 +623,21 @@ if __name__ == "__main__":
     
     # Set device if specified
     if args.device:
-        DEVICE = args.device
+        if args.device == "cuda":
+            try:
+                if not torch.cuda.is_available():
+                    logging.warning("CUDA requested but not available, falling back to CPU")
+                    DEVICE = "cpu"
+                else:
+                    torch.cuda.current_device()  # Test CUDA access
+                    DEVICE = args.device
+            except Exception as e:
+                logging.warning(f"CUDA test failed: {e}, falling back to CPU")
+                DEVICE = "cpu"
+        else:
+            DEVICE = args.device
+    
+    logging.info(f"Device selected: {DEVICE}")
     
     # Initialize model
     logging.info("Initializing MiraTTS model...")