Update app.py

app.py

@@ -38,41 +38,29 @@ This demo allows you to have a conversation with Sesame CSM 1B, leveraging Whisp
 *Disclaimer: This demo relies on several large models. Expect longer processing times, and potential resource limitations.*
 """
 
-# Model Loading
-@spaces.GPU()
-try:
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model_path = hf_hub_download(repo_id="sesame/csm-1b", filename="ckpt.pt")
-    generator = load_csm_1b(model_path, device)
-    logging.info("Sesame CSM 1B loaded successfully.")
-
-    whisper_model, whisper_metadata = whisperx.load_model("large-v2", device)
-    model_a, whisper_metadata = whisperx.load_align_model(language_code=whisper_metadata.language, device=device)
-    logging.info("WhisperX model loaded successfully.")
-
-    # Load Gemma 1.1 2B - adjust model name if needed
-    tokenizer_gemma = AutoTokenizer.from_pretrained("google/gemma-3-1b-pt")
-    model_gemma = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-pt").to(device)
-    logging.info("Gemma 3 1B pt model loaded successfully.")
-
-except Exception as e:
-    logging.error(f"Model loading error: {e}")
-    raise  # Re-raise to prevent the app from launching with incomplete models
-
 # 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  #Reduce for the 1.1 2b model
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Global conversation history (important: keep it inside app scope)
 conversation_history = []
 
+# Global variables to hold loaded models
+global_generator = None
+global_whisper_model = None
+global_model_a = None
+global_whisper_metadata = None
+global_tokenizer_gemma = None
+global_model_gemma = None
+
 # --- HELPER FUNCTIONS ---
-def transcribe_audio(audio_path: str) -> str:
+def transcribe_audio(audio_path: str, whisper_model, model_a, whisper_metadata) -> str:
     """Transcribes audio using WhisperX."""
     try:
         audio = whisperx.load_audio(audio_path)
-        result = whisper_model.transcribe(audio, batch_size=16) # Added batch_size
+        result = whisper_model.transcribe(audio, batch_size=16)
 
         # Align Whisper output
         result_aligned = whisperx.align(result["segments"], model_a, whisper_metadata, audio, whisper_model, device, return_char_alignments=False)
@@ -80,32 +68,32 @@ def transcribe_audio(audio_path: str) -> str:
         return result_aligned["segments"][0]["text"]
     except Exception as e:
         logging.error(f"WhisperX transcription error: {e}")
-        return "Error: Could not transcribe audio." # Return an error message
+        return "Error: Could not transcribe audio."
 
-def generate_response(text: str) -> str:
+def generate_response(text: str, tokenizer_gemma, model_gemma) -> str:
     """Generates a response using Gemma."""
     try:
         input_text = "Here is a response for the user. " + text
         input = tokenizer_gemma(input_text, return_tensors="pt").to(device)
-        generated_output = model_gemma.generate(**input, max_length=MAX_GEMMA_LENGTH, early_stopping=True) # Added early_stopping
+        generated_output = model_gemma.generate(**input, max_length=MAX_GEMMA_LENGTH, early_stopping=True)
         return tokenizer_gemma.decode(generated_output[0], skip_special_tokens=True)
     except Exception as e:
         logging.error(f"Gemma response generation error: {e}")
-        return "I'm sorry, I encountered an error generating a response."  # Error fallback
+        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 != generator.sample_rate:
+        if sample_rate != global_generator.sample_rate: #Access via global generator
             audio_tensor = torchaudio.functional.resample(
-                audio_tensor, orig_freq=sample_rate, new_freq=generator.sample_rate
+                audio_tensor, orig_freq=sample_rate, new_freq=global_generator.sample_rate
             )
         return audio_tensor
     except Exception as e:
         logging.error(f"Audio loading error: {e}")
-        raise gr.Error("Could not load or process the audio file.") from e # Re-raise as Gradio error
+        raise gr.Error("Could not load or process the audio file.") from e
 
 def clear_history():
     """Clears the conversation history"""
@@ -115,28 +103,47 @@ def clear_history():
     return "Conversation history cleared."
 
 # --- MAIN INFERENCE FUNCTION ---
-#@spaces.GPU()
-def infer(user_audio) -> tuple[int, np.ndarray]:  # Return sample_rate as int
+@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
+
     try:
         if not user_audio:
             raise ValueError("No audio input received.")
-        return _infer(user_audio)
+
+        # Load models if not already loaded
+        if global_generator is None:
+            model_path = hf_hub_download(repo_id="sesame/csm-1b", filename="ckpt.pt")
+            global_generator = load_csm_1b(model_path, device)
+            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)
+            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)
     except Exception as e:
-        logging.exception(f"Inference error: {e}") # Log the full exception
+        logging.exception(f"Inference error: {e}")
         raise gr.Error(f"An error occurred during processing: {e}")
 
-def _infer(user_audio) -> tuple[int, np.ndarray]:  # Return sample_rate as int
+def _infer(user_audio, generator, whisper_model, model_a, whisper_metadata, tokenizer_gemma, model_gemma) -> tuple:
     """Processes the user input, generates a response, and returns audio."""
-    global conversation_history # Declare to modify the global list
+    global conversation_history
 
     try:
         # 1. ASR: Transcribe user audio using WhisperX
-        user_text = transcribe_audio(user_audio)
+        user_text = transcribe_audio(user_audio, whisper_model, model_a, whisper_metadata)
         logging.info(f"User: {user_text}")
 
         # 2. LLM: Generate a response using Gemma
-        ai_text = generate_response(user_text)
+        ai_text = generate_response(user_text, tokenizer_gemma, model_gemma)
         logging.info(f"AI: {ai_text}")
 
         # 3. Generate audio using the CSM model
@@ -149,8 +156,8 @@ def _infer(user_audio) -> tuple[int, np.ndarray]:  # Return sample_rate as int
             )
             logging.info("Audio generated successfully.")
         except Exception as e:
-             logging.error(f"Gemma response generation error: {e}")
-             raise gr.Error(f"Gemma response generation error: {e}")  # Error fallback
+             logging.error(f"CSM response generation error: {e}")
+             raise gr.Error(f"CSM response generation error: {e}")
 
         #Update conversation history with user input and ai response.
         user_segment = Segment(speaker = SPEAKER_ID, text = 'User Audio', audio = load_audio(user_audio))
@@ -175,8 +182,6 @@ def _infer(user_audio) -> tuple[int, np.ndarray]:  # Return sample_rate as int
 
     except Exception as e:
         logging.exception(f"Error in _infer: {e}")
-        # Log the full exception including stack trace for debugging.
-        # It's crucial to log the *exception*, not just the error message.
         raise gr.Error(f"An error occurred during processing: {e}")
 
 # --- GRADIO INTERFACE ---
@@ -189,6 +194,6 @@ with gr.Blocks() as app:
 
     btn = gr.Button("Generate Response")
     btn.click(infer, inputs=[audio_input], outputs=[audio_output])
-    clear_button.click(clear_history, outputs=[status_display]) # No input needed
+    clear_button.click(clear_history, outputs=[status_display])
 
-app.launch(ssr_mode=True)
+app.launch(share=False) #Add share = True for public link