Update app.py

app.py

@@ -3,9 +3,11 @@ import torch
 import logging
 import gc
 import time
+import signal
 from pathlib import Path
 from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 import librosa
+from functools import wraps
 
 # Try to import flash attention, but don't fail if not available
 try:
@@ -20,6 +22,25 @@ except ImportError:
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+def timeout_handler(signum, frame):
+    raise TimeoutError("Operation timed out")
+
+def with_timeout(seconds):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Set the signal handler and a timeout alarm
+            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(seconds)
+            try:
+                result = func(*args, **kwargs)
+            finally:
+                signal.alarm(0)  # Disable the alarm
+                signal.signal(signal.SIGALRM, old_handler)
+            return result
+        return wrapper
+    return decorator
+
 class OptimizedWhisperApp:
     def __init__(self):
         self.pipe = None
@@ -37,30 +58,34 @@ class OptimizedWhisperApp:
         ]
     
     def create_pipe(self, model_name, use_flash_attention=True):
-        """Create pipeline like the successful space"""
+        """Create pipeline with better error handling"""
         try:
+            logger.info(f"Starting to load model: {model_name}")
+            
             # Device selection
             if torch.cuda.is_available():
                 device = "cuda:0"
                 torch_dtype = torch.float16
+                logger.info("Using CUDA device")
             else:
                 device = "cpu"
                 torch_dtype = torch.float32
+                logger.info("Using CPU device")
             
-            logger.info(f"Loading {model_name} on {device} with {torch_dtype}")
-            
-            # Attention implementation - gracefully handle missing flash attention
+            # Simpler attention implementation selection
+            attn_implementation = "eager"  # Start with most compatible
             if use_flash_attention and FLASH_ATTN_AVAILABLE and is_flash_attn_2_available() and torch.cuda.is_available():
-                attn_implementation = "flash_attention_2"
-                logger.info("Using Flash Attention 2")
-            else:
-                attn_implementation = "sdpa"
-                if use_flash_attention and not FLASH_ATTN_AVAILABLE:
-                    logger.info("Flash Attention requested but not available, using SDPA")
-                else:
-                    logger.info(f"Using {attn_implementation}")
+                try:
+                    attn_implementation = "flash_attention_2"
+                    logger.info("Attempting Flash Attention 2")
+                except:
+                    attn_implementation = "eager"
+                    logger.info("Flash Attention 2 failed, using eager")
+            
+            logger.info(f"Using attention implementation: {attn_implementation}")
             
-            # Load model directly
+            # Load model with timeout protection
+            logger.info("Loading model...")
             model = AutoModelForSpeechSeq2Seq.from_pretrained(
                 model_name,
                 torch_dtype=torch_dtype,
@@ -69,12 +94,15 @@ class OptimizedWhisperApp:
                 attn_implementation=attn_implementation,
                 cache_dir="./cache"
             )
+            logger.info("Model loaded, moving to device...")
             model.to(device)
             
             # Load processor
+            logger.info("Loading processor...")
             processor = AutoProcessor.from_pretrained(model_name)
             
-            # Create pipeline manually
+            # Create pipeline
+            logger.info("Creating pipeline...")
             pipe = pipeline(
                 "automatic-speech-recognition",
                 model=model,
@@ -89,25 +117,38 @@ class OptimizedWhisperApp:
             
         except Exception as e:
             logger.error(f"Failed to create pipeline: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
             return None
     
     def load_model(self, model_name, use_flash_attention=True):
-        """Load model if different from current"""
+        """Load model with timeout protection"""
         if self.current_model != model_name or self.pipe is None:
             logger.info(f"Loading new model: {model_name}")
             
             # Clear previous model
             if self.pipe is not None:
+                logger.info("Clearing previous model...")
                 del self.pipe
                 if torch.cuda.is_available():
                     torch.cuda.empty_cache()
                 gc.collect()
             
-            # Create new pipeline
-            self.pipe = self.create_pipe(model_name, use_flash_attention)
-            self.current_model = model_name if self.pipe else None
-            
-            return self.pipe is not None
+            try:
+                # Create new pipeline with timeout
+                self.pipe = self.create_pipe(model_name, use_flash_attention)
+                self.current_model = model_name if self.pipe else None
+                
+                if self.pipe:
+                    logger.info(f"Model {model_name} loaded successfully")
+                    return True
+                else:
+                    logger.error(f"Failed to load model {model_name}")
+                    return False
+                    
+            except Exception as e:
+                logger.error(f"Error loading model: {e}")
+                return False
         else:
             logger.info("Model already loaded")
             return True
@@ -116,32 +157,33 @@ class OptimizedWhisperApp:
                         language="Automatic Detection", task="transcribe",
                         chunk_length_s=30, batch_size=16, use_flash_attention=True,
                         return_timestamps=True):
-        """Transcribe using the optimized approach"""
+        """Transcribe with comprehensive error handling"""
         
         if audio_file is None:
             return "Please upload an audio file", "", ""
         
         try:
+            logger.info("=== Starting transcription ===")
             start_time = time.time()
             
-            # Load model if needed
+            # Load model
+            logger.info(f"Loading model: {model_name}")
             success = self.load_model(model_name, use_flash_attention)
             if not success:
-                return "Failed to load model", "", ""
+                return "Failed to load model - check logs for details", "", ""
             
-            logger.info(f"Processing: {audio_file}")
-            logger.info(f"Settings: {model_name}, {language}, {task}")
+            logger.info(f"Processing audio file: {audio_file}")
+            logger.info(f"Settings - Model: {model_name}, Language: {language}, Task: {task}")
             logger.info(f"Chunk length: {chunk_length_s}s, Batch size: {batch_size}")
             
             # Prepare generation kwargs
             generate_kwargs = {}
             
-            # Only set language if not auto-detection and model supports multilingual
+            # Language handling
             if language != "Automatic Detection" and not model_name.endswith(".en"):
-                # Map common language names
                 language_map = {
                     "Greek": "greek",
-                    "English": "english",
+                    "English": "english", 
                     "Spanish": "spanish",
                     "French": "french",
                     "German": "german",
@@ -151,29 +193,37 @@ class OptimizedWhisperApp:
                 generate_kwargs["language"] = lang_code
                 logger.info(f"Set language: {lang_code}")
             
-            # Set task if model supports it
+            # Task handling
             if not model_name.endswith(".en"):
                 generate_kwargs["task"] = task
                 logger.info(f"Set task: {task}")
             
-            # Transcribe
-            logger.info("Starting transcription...")
-            outputs = self.pipe(
-                audio_file,
-                chunk_length_s=chunk_length_s,
-                batch_size=batch_size,
-                generate_kwargs=generate_kwargs,
-                return_timestamps=return_timestamps,
-            )
+            # Transcribe with error handling
+            logger.info("Starting transcription process...")
+            try:
+                outputs = self.pipe(
+                    audio_file,
+                    chunk_length_s=chunk_length_s,
+                    batch_size=batch_size,
+                    generate_kwargs=generate_kwargs,
+                    return_timestamps=return_timestamps,
+                )
+                logger.info("Transcription completed")
+            except Exception as transcribe_error:
+                logger.error(f"Transcription failed: {transcribe_error}")
+                return f"Transcription failed: {str(transcribe_error)}", "", ""
             
             transcription_time = time.time() - start_time
-            logger.info(f"Transcription completed in {transcription_time:.2f} seconds")
+            logger.info(f"Total processing time: {transcription_time:.2f} seconds")
+            
+            # Extract and validate results
+            transcription = outputs.get("text", "") if outputs else ""
+            chunks = outputs.get("chunks", []) if outputs else []
             
-            # Extract results
-            transcription = outputs.get("text", "")
-            chunks = outputs.get("chunks", [])
+            logger.info(f"Extracted transcription length: {len(transcription)} chars")
+            logger.info(f"Number of chunks: {len(chunks)}")
             
-            # Format timestamps with error handling
+            # Handle timestamps safely
             timestamp_text = ""
             if return_timestamps:
                 try:
@@ -181,9 +231,9 @@ class OptimizedWhisperApp:
                         timestamp_text = self._format_timestamps(chunks)
                     else:
                         timestamp_text = "=== TIMESTAMPS ===\nNo chunks returned by the model.\n"
-                except Exception as e:
-                    logger.warning(f"Error formatting timestamps: {e}")
-                    timestamp_text = f"=== TIMESTAMPS ===\nError formatting timestamps: {str(e)}\n"
+                except Exception as ts_error:
+                    logger.warning(f"Error formatting timestamps: {ts_error}")
+                    timestamp_text = f"=== TIMESTAMPS ===\nError formatting timestamps: {str(ts_error)}\n"
             else:
                 timestamp_text = "=== TIMESTAMPS ===\nTimestamp output disabled.\n"
             
@@ -194,44 +244,50 @@ class OptimizedWhisperApp:
                 use_flash_attention, len(chunks)
             )
             
+            logger.info("=== Transcription completed successfully ===")
             return transcription.strip(), timestamp_text, detailed_output
             
         except Exception as e:
             error_msg = f"Transcription error: {str(e)}"
             logger.error(error_msg)
+            import traceback
+            logger.error(traceback.format_exc())
             return error_msg, "", error_msg
     
     def _format_timestamps(self, chunks):
-        """Format timestamp information with proper error handling"""
+        """Format timestamp information with comprehensive error handling"""
         timestamp_text = "=== TIMESTAMPS ===\n"
         
         if not chunks:
             timestamp_text += "No timestamp information available.\n"
             return timestamp_text
         
+        logger.info(f"Formatting {len(chunks)} chunks")
+        
         for i, chunk in enumerate(chunks):
-            timestamp = chunk.get('timestamp', None)
-            text = chunk.get('text', '')
-            
-            # Handle various timestamp formats and None cases
-            if timestamp is None:
-                timestamp_text += f"[No timestamp]: {text}\n"
-            elif isinstance(timestamp, (list, tuple)) and len(timestamp) >= 2:
-                start, end = timestamp[0], timestamp[1]
-                # Handle None values in timestamp array
-                if start is None or end is None:
-                    timestamp_text += f"[Invalid timestamp]: {text}\n"
+            try:
+                timestamp = chunk.get('timestamp', None)
+                text = chunk.get('text', '')
+                
+                if timestamp is None:
+                    timestamp_text += f"[No timestamp]: {text}\n"
+                elif isinstance(timestamp, (list, tuple)) and len(timestamp) >= 2:
+                    start, end = timestamp[0], timestamp[1]
+                    if start is None or end is None:
+                        timestamp_text += f"[Invalid timestamp]: {text}\n"
+                    else:
+                        try:
+                            start_f = float(start) if start is not None else 0.0
+                            end_f = float(end) if end is not None else 0.0
+                            timestamp_text += f"[{start_f:.1f}s - {end_f:.1f}s]: {text}\n"
+                        except (ValueError, TypeError) as ve:
+                            timestamp_text += f"[Format error]: {text}\n"
                 else:
-                    try:
-                        # Convert to float if needed and format
-                        start_f = float(start) if start is not None else 0.0
-                        end_f = float(end) if end is not None else 0.0
-                        timestamp_text += f"[{start_f:.1f}s - {end_f:.1f}s]: {text}\n"
-                    except (ValueError, TypeError):
-                        timestamp_text += f"[Invalid timestamp format]: {text}\n"
-            else:
-                # Handle unexpected timestamp format
-                timestamp_text += f"[Timestamp format error]: {text}\n"
+                    timestamp_text += f"[Unexpected format]: {text}\n"
+                    
+            except Exception as chunk_error:
+                logger.warning(f"Error processing chunk {i}: {chunk_error}")
+                timestamp_text += f"[Chunk {i} error]: Processing failed\n"
         
         return timestamp_text
     
@@ -254,21 +310,24 @@ class OptimizedWhisperApp:
         output += f"Batch size: {batch_size}\n"
         output += f"Flash Attention: {'Enabled' if use_flash_attention else 'Disabled'}\n"
         
-        if self.pipe:
-            device = next(self.pipe.model.parameters()).device
-            dtype = next(self.pipe.model.parameters()).dtype
-            output += f"Device: {device}\n"
-            output += f"Data type: {dtype}\n"
+        if self.pipe and hasattr(self.pipe, 'model'):
+            try:
+                device = next(self.pipe.model.parameters()).device
+                dtype = next(self.pipe.model.parameters()).dtype
+                output += f"Device: {device}\n"
+                output += f"Data type: {dtype}\n"
+            except:
+                output += "Device info: Unable to retrieve\n"
         
         output += f"Flash Attention 2 available: {FLASH_ATTN_AVAILABLE and is_flash_attn_2_available()}\n"
         
-        output += "\n=== OPTIMIZATIONS ===\n"
-        output += "• Direct model loading (not pipeline abstraction)\n"
-        output += "• Manual pipeline construction\n"
-        output += "• Optimized attention mechanism\n"
-        output += "• Batch processing\n"
-        output += "• Conservative language handling\n"
-        output += "• Proper memory management\n"
+        output += "\n=== SYSTEM STATUS ===\n"
+        if torch.cuda.is_available():
+            output += f"GPU Available: Yes\n"
+            output += f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB\n"
+            output += f"GPU Memory Used: {torch.cuda.memory_allocated() / 1e9:.1f}GB\n"
+        else:
+            output += "GPU Available: No\n"
         
         return output
     
@@ -277,10 +336,12 @@ class OptimizedWhisperApp:
         if self.pipe is None:
             return "No model loaded"
         
-        device = next(self.pipe.model.parameters()).device
-        dtype = next(self.pipe.model.parameters()).dtype
-        
-        return f"✅ {self.current_model} loaded on {device} ({dtype})"
+        try:
+            device = next(self.pipe.model.parameters()).device
+            dtype = next(self.pipe.model.parameters()).dtype
+            return f"✅ {self.current_model} loaded on {device} ({dtype})"
+        except:
+            return f"✅ {self.current_model} loaded (device info unavailable)"
 
 # Initialize the app
 logger.info("Initializing Optimized Whisper App...")
@@ -288,11 +349,17 @@ whisper_app = OptimizedWhisperApp()
 
 def transcribe_wrapper(audio, model_name, language, task, chunk_length_s, 
                       batch_size, use_flash_attention, return_timestamps):
-    """Wrapper for Gradio interface"""
-    return whisper_app.transcribe_audio(
-        audio, model_name, language, task,
-        chunk_length_s, batch_size, use_flash_attention, return_timestamps
-    )
+    """Wrapper for Gradio interface with additional safety"""
+    try:
+        logger.info(f"Transcribe wrapper called with model: {model_name}")
+        return whisper_app.transcribe_audio(
+            audio, model_name, language, task,
+            chunk_length_s, batch_size, use_flash_attention, return_timestamps
+        )
+    except Exception as e:
+        error_msg = f"Wrapper error: {str(e)}"
+        logger.error(error_msg)
+        return error_msg, "", error_msg
 
 def get_model_status():
     """Get current model status"""
@@ -306,13 +373,13 @@ def create_interface():
             """
             # 🚀 Optimized Whisper Transcription
             
-            **High-Performance Speech-to-Text Based on Successful Implementation**
+            **High-Performance Speech-to-Text with Enhanced Error Handling**
             
-            Uses the same optimizations as high-performing Whisper spaces:
-            - Direct model loading for better control
-            - Flash Attention 2 support (when available)
-            - Optimized chunking and batching
-            - Conservative parameter handling
+            Features:
+            - Comprehensive error handling and timeout protection
+            - Better logging for debugging
+            - Graceful handling of model loading issues
+            - Robust timestamp processing
             """
         )
         
@@ -335,9 +402,9 @@ def create_interface():
                 # Model selection
                 model_dropdown = gr.Dropdown(
                     choices=whisper_app.available_models,
-                    value="openai/whisper-medium",
+                    value="openai/whisper-small",  # Start with smaller model
                     label="Model",
-                    info="Medium often works best for real-world usage"
+                    info="Start with 'small' model for testing"
                 )
                 
                 # Basic settings
@@ -367,18 +434,17 @@ def create_interface():
                     
                     batch_size = gr.Slider(
                         minimum=1,
-                        maximum=32,
-                        value=16,
+                        maximum=8,  # Reduced default
+                        value=4,    # Smaller batch size
                         step=1,
                         label="Batch Size",
-                        info="Higher = faster, more memory"
+                        info="Start with smaller batch size"
                     )
                     
                     use_flash_attention = gr.Checkbox(
                         label="Flash Attention 2",
-                        value=FLASH_ATTN_AVAILABLE and torch.cuda.is_available(),
-                        info="Faster processing (requires compatible GPU and flash-attn package)",
-                        interactive=FLASH_ATTN_AVAILABLE
+                        value=False,  # Disabled by default for stability
+                        info="Enable only if you have compatible GPU and flash-attn installed"
                     )
                     
                     return_timestamps = gr.Checkbox(
@@ -432,24 +498,19 @@ def create_interface():
         # Footer
         gr.Markdown(
             """
-            ### 🎯 Model Recommendations
-            
-            **For Greek dialect of Lesbos:**
-            - `ilsp/whisper_greek_dialect_of_lesbos` - Specialized but may have issues
-            - `openai/whisper-medium` - Often better for real-world usage
-            - `openai/whisper-large-v2` - More accurate but slower
+            ### 🔧 Troubleshooting
             
-            **General recommendations:**
-            - **Medium model** often provides the best balance
-            - **30-second chunks** work well for most audio
-            - **Flash Attention** speeds up processing significantly (when available)
-            - **Automatic language detection** usually works well
+            **If processing hangs:**
+            1. Start with `whisper-small` model
+            2. Disable Flash Attention
+            3. Use smaller batch size (1-4)
+            4. Check the console logs for errors
             
-            ### ⚡ Performance Tips
-            - GPU with Flash Attention 2 = Fastest
-            - Batch size 16-24 optimal for most GPUs  
-            - Lower chunk length for very noisy audio
-            - Use English-only models (.en) for English-only content
+            **For best results:**
+            - Test with small model first
+            - Gradually increase model size if needed
+            - Monitor GPU memory usage
+            - Check logs for any error messages
             """
         )
     
@@ -458,4 +519,4 @@ def create_interface():
 # Launch the app
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch(share=True)
+    interface.launch(share=True, debug=True)