Update app.py

app.py

@@ -73,25 +73,21 @@ class ProfessionalVoiceAgent:
     def load_whisper(self):
         """Load Whisper model for speech recognition"""
         try:
-            if self.use_large_models:
-                # Use larger Whisper for better accuracy
-                model_name = "openai/whisper-small"
-                logger.info(f"Loading Whisper Small for best accuracy...")
-            else:
-                model_name = "openai/whisper-tiny"
-                logger.info(f"Loading Whisper Tiny...")
+            # Use tiny model for speed - small is too slow
+            model_name = "openai/whisper-tiny"
+            logger.info(f"Loading Whisper Tiny for fast processing...")
 
             self.whisper_processor = WhisperProcessor.from_pretrained(model_name)
             self.whisper_model = WhisperForConditionalGeneration.from_pretrained(
                 model_name,
-                torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32
+                torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
+                low_cpu_mem_usage=True
             ).to(self.device)
 
-            # Enable Flash Attention if available
-            if hasattr(self.whisper_model.config, "use_flash_attention_2"):
-                self.whisper_model.config.use_flash_attention_2 = True
+            # Set to eval mode for inference
+            self.whisper_model.eval()
 
-            logger.info("✓ Whisper loaded successfully")
+            logger.info(f"✓ Whisper loaded on {self.device}")
 
         except Exception as e:
             logger.error(f"Failed to load Whisper: {e}")
@@ -123,7 +119,10 @@ class ProfessionalVoiceAgent:
             # Add padding token
             self.chat_tokenizer.pad_token = self.chat_tokenizer.eos_token
 
-            logger.info("✓ Chat model loaded successfully")
+            # Set to eval mode
+            self.chat_model.eval()
+
+            logger.info(f"✓ Chat model loaded on {self.device}")
 
         except Exception as e:
             logger.error(f"Failed to load chat model: {e}")
@@ -149,6 +148,10 @@ class ProfessionalVoiceAgent:
                 torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32
             ).to(self.device)
 
+            # Set to eval mode
+            self.tts_model.eval()
+            self.vocoder.eval()
+
             # Load speaker embeddings for voice
             try:
                 logger.info("Loading speaker embeddings dataset...")
@@ -212,6 +215,12 @@ class ProfessionalVoiceAgent:
                 import librosa
                 audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=16000)
 
+            # Trim silence and limit audio length for speed (max 30 seconds)
+            max_samples = 16000 * 30  # 30 seconds at 16kHz
+            if len(audio_data) > max_samples:
+                logger.warning(f"Audio trimmed from {len(audio_data)/16000:.1f}s to 30s")
+                audio_data = audio_data[:max_samples]
+
             if self.whisper_processor and hasattr(self.whisper_model, 'generate'):
                 # Use loaded model
                 input_features = self.whisper_processor(
@@ -220,14 +229,15 @@ class ProfessionalVoiceAgent:
                     return_tensors="pt"
                 ).input_features.to(self.device)
 
-                # Generate token ids
+                # Generate token ids - optimized for speed
                 with torch.cuda.amp.autocast(enabled=self.device.type == "cuda"):
-                    predicted_ids = self.whisper_model.generate(
-                        input_features,
-                        max_new_tokens=128,
-                        num_beams=5,  # Better accuracy
-                        temperature=0.0  # Deterministic
-                    )
+                    with torch.no_grad():
+                        predicted_ids = self.whisper_model.generate(
+                            input_features,
+                            max_new_tokens=64,  # Reduced for faster processing
+                            num_beams=1,  # Greedy decoding for speed
+                            do_sample=False  # Deterministic
+                        )
 
                 # Decode token ids to text
                 transcription = self.whisper_processor.batch_decode(
@@ -274,18 +284,18 @@ class ProfessionalVoiceAgent:
                     max_length=512
                 ).to(self.device)
 
-                # Generate response
+                # Generate response - optimized for speed
                 with torch.cuda.amp.autocast(enabled=self.device.type == "cuda"):
                     with torch.no_grad():
                         outputs = self.chat_model.generate(
                             inputs,
-                            max_new_tokens=100,
+                            max_new_tokens=50,  # Shorter for faster response
                             temperature=temperature,
                             top_p=0.9,
-                            do_sample=True,
+                            do_sample=True if temperature > 0 else False,
                             pad_token_id=self.chat_tokenizer.eos_token_id,
                             eos_token_id=self.chat_tokenizer.eos_token_id,
-                            num_beams=3
+                            num_beams=1  # Greedy for speed
                         )
 
                 # Decode response