Update app.py

app.py

@@ -28,48 +28,55 @@ try:
     gpu_timeout = int(os.getenv("GPU_TIMEOUT", 180))
 except (ValueError, TypeError) as e:
     logging.error(f"Configuration error: {e}")
-    raise  # Re-raise the exception to halt the application
+    raise
 
 SPACE_INTRO_TEXT = """\
 # Sesame CSM 1B - Conversational Demo
 
-This demo allows you to have a conversation with Sesame CSM 1B, leveraging WhisperX for speech-to-text and Gemma for generating responses. This is an experimental integration and may require significant resources.
+This demo allows you to have a conversation with Sesame CSM 1B, leveraging WhisperX for speech-to-text and Gemma for generating responses.  This is an experimental integration and may require significant resources.
 
-*Disclaimer: This demo relies on several large models. Expect longer processing times, and potential resource limitations.*
+*Disclaimer: This demo relies on several large models.  Expect longer processing times, and potential resource limitations.*
 """
 
 # Constants
 SPEAKER_ID = 0  # Arbitrary speaker ID
 MAX_CONTEXT_SEGMENTS = 5
-MAX_GEMMA_LENGTH = 300  #Reduce for the 1.1 2b model
+MAX_GEMMA_LENGTH = 300
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Global conversation history (important: keep it inside app scope)
+# Global conversation history
 conversation_history = []
 
 # Global variables to hold loaded models
 global_generator = None
 global_whisper_model = None
 global_model_a = None
-global_whisper_metadata = None
+# global_whisper_metadata = None  # No longer needed at the global level
 global_tokenizer_gemma = None
 global_model_gemma = None
 
+
 # --- HELPER FUNCTIONS ---
-def transcribe_audio(audio_path: str, whisper_model, model_a, whisper_metadata) -> str:
-    """Transcribes audio using WhisperX."""
+
+def transcribe_audio(audio_path: str, whisper_model, model_a) -> str:  # Removed whisper_metadata
+    """Transcribes audio using WhisperX and aligns it."""
     try:
         audio = whisperx.load_audio(audio_path)
         result = whisper_model.transcribe(audio, batch_size=16)
+        # Get language from the result.  Much more reliable.
+        language = result["language"]
+
 
         # Align Whisper output
-        result_aligned = whisperx.align(result["segments"], model_a, whisper_metadata, audio, whisper_model, device, return_char_alignments=False)
+        model_a, metadata = whisperx.load_align_model(language_code=language, device=device) #Load it here to ensure metadata is extracted.
+        result_aligned = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
 
         return result_aligned["segments"][0]["text"]
     except Exception as e:
         logging.error(f"WhisperX transcription error: {e}")
         return "Error: Could not transcribe audio."
 
+
 def generate_response(text: str, tokenizer_gemma, model_gemma) -> str:
     """Generates a response using Gemma."""
     try:
@@ -81,12 +88,13 @@ def generate_response(text: str, tokenizer_gemma, model_gemma) -> str:
         logging.error(f"Gemma response generation error: {e}")
         return "I'm sorry, I encountered an error generating a response."
 
+
 def load_audio(audio_path: str) -> torch.Tensor:
     """Loads audio from file and returns a torch tensor."""
     try:
         audio_tensor, sample_rate = torchaudio.load(audio_path)
         audio_tensor = audio_tensor.mean(dim=0)  # Mono audio
-        if sample_rate != global_generator.sample_rate: #Access via global generator
+        if sample_rate != global_generator.sample_rate:
             audio_tensor = torchaudio.functional.resample(
                 audio_tensor, orig_freq=sample_rate, new_freq=global_generator.sample_rate
             )
@@ -95,6 +103,7 @@ def load_audio(audio_path: str) -> torch.Tensor:
         logging.error(f"Audio loading error: {e}")
         raise gr.Error("Could not load or process the audio file.") from e
 
+
 def clear_history():
     """Clears the conversation history"""
     global conversation_history
@@ -102,11 +111,13 @@ def clear_history():
     logging.info("Conversation history cleared.")
     return "Conversation history cleared."
 
+
 # --- MAIN INFERENCE FUNCTION ---
+
 @spaces.GPU(gpu_timeout=gpu_timeout)
 def infer(user_audio) -> tuple:
     """Infers a response from the user audio."""
-    global global_generator, global_whisper_model, global_model_a, global_whisper_metadata, global_tokenizer_gemma, global_model_gemma, device
+    global global_generator, global_whisper_model, global_model_a, global_tokenizer_gemma, global_model_gemma, device
 
     try:
         if not user_audio:
@@ -119,27 +130,26 @@ def infer(user_audio) -> tuple:
             logging.info("Sesame CSM 1B loaded successfully on GPU.")
 
         if global_whisper_model is None:
-            global_whisper_model, global_whisper_metadata = whisperx.load_model("large-v2", device)
-            global_model_a, _ = whisperx.load_align_model(language_code=global_whisper_metadata.language, device=device)
+            global_whisper_model = whisperx.load_model("large-v2", device) # No unpacking
             logging.info("WhisperX model loaded successfully on GPU.")
 
         if global_tokenizer_gemma is None:
             global_tokenizer_gemma = AutoTokenizer.from_pretrained("google/gemma-3-1b-pt")
             global_model_gemma = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-pt").to(device)
             logging.info("Gemma 3 1B pt model loaded successfully on GPU.")
-
-        return _infer(user_audio, global_generator, global_whisper_model, global_model_a, global_whisper_metadata, global_tokenizer_gemma, global_model_gemma)
+        return _infer(user_audio, global_generator, global_whisper_model, global_model_a, global_tokenizer_gemma, global_model_gemma) #Removed Metadata
     except Exception as e:
         logging.exception(f"Inference error: {e}")
         raise gr.Error(f"An error occurred during processing: {e}")
 
-def _infer(user_audio, generator, whisper_model, model_a, whisper_metadata, tokenizer_gemma, model_gemma) -> tuple:
+
+def _infer(user_audio, generator, whisper_model, model_a, tokenizer_gemma, model_gemma) -> tuple:
     """Processes the user input, generates a response, and returns audio."""
     global conversation_history
 
     try:
         # 1. ASR: Transcribe user audio using WhisperX
-        user_text = transcribe_audio(user_audio, whisper_model, model_a, whisper_metadata)
+        user_text = transcribe_audio(user_audio, whisper_model, model_a) #Removed Metadata
         logging.info(f"User: {user_text}")
 
         # 2. LLM: Generate a response using Gemma
@@ -147,17 +157,13 @@ def _infer(user_audio, generator, whisper_model, model_a, whisper_metadata, toke
         logging.info(f"AI: {ai_text}")
 
         # 3. Generate audio using the CSM model
-        try:
-            ai_audio = generator.generate(
-                text=ai_text,
-                speaker=SPEAKER_ID,
-                context=conversation_history,
-                max_audio_length_ms=30_000,
-            )
-            logging.info("Audio generated successfully.")
-        except Exception as e:
-             logging.error(f"CSM response generation error: {e}")
-             raise gr.Error(f"CSM response generation error: {e}")
+        ai_audio = generator.generate(
+            text=ai_text,
+            speaker=SPEAKER_ID,
+            context=conversation_history,
+            max_audio_length_ms=30_000,
+        )
+        logging.info("Audio generated successfully.")
 
         #Update conversation history with user input and ai response.
         user_segment = Segment(speaker = SPEAKER_ID, text = 'User Audio', audio = load_audio(user_audio))
@@ -184,7 +190,9 @@ def _infer(user_audio, generator, whisper_model, model_a, whisper_metadata, toke
         logging.exception(f"Error in _infer: {e}")
         raise gr.Error(f"An error occurred during processing: {e}")
 
+
 # --- GRADIO INTERFACE ---
+
 with gr.Blocks() as app:
     gr.Markdown(SPACE_INTRO_TEXT)
     audio_input = gr.Audio(label="Your Input", type="filepath")
@@ -196,4 +204,4 @@ with gr.Blocks() as app:
     btn.click(infer, inputs=[audio_input], outputs=[audio_output])
     clear_button.click(clear_history, outputs=[status_display])
 
-app.launch(share=False) #Add share = True for public link
+app.launch(share=False)