Update app.py

app.py

@@ -7,14 +7,14 @@ import torchaudio
 from generator import Segment, load_csm_1b
 from huggingface_hub import hf_hub_download, login
 from watermarking import watermark
-import whisperx
+import whisper
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Authentication and Configuration
+# --- Authentication and Configuration ---  (Moved BEFORE model loading)
 try:
     api_key = os.getenv("HF_TOKEN")
     if not api_key:
@@ -30,55 +30,37 @@ except (ValueError, TypeError) as e:
     logging.error(f"Configuration error: {e}")
     raise
 
-SPACE_INTRO_TEXT = """\
+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 Whisper 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
+# --- Model Loading --- (Moved INSIDE infer function)
+
+# --- Constants ---  (Constants can stay outside)
+SPEAKER_ID = 0
 MAX_CONTEXT_SEGMENTS = 5
 MAX_GEMMA_LENGTH = 300
-device = "cuda"  # if torch.cuda.is_available() else "cpu"
 
-# Global conversation history
+# --- 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  # No longer needed at the global level
-global_tokenizer_gemma = None
-global_model_gemma = None
-
-
-# --- HELPER FUNCTIONS ---
+# --- Helper Functions ---
 
-def transcribe_audio(audio_path: str, whisper_model, model_a) -> str:  # Removed whisper_metadata
-    """Transcribes audio using WhisperX and aligns it."""
+def transcribe_audio(audio_path: str, whisper_model) -> str:  # Pass whisper_model
     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
-        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"]
+        audio = whisper.load_audio(audio_path)
+        audio = whisper.pad_or_trim(audio)
+        result = whisper_model.transcribe(audio)
+        return result["text"]
     except Exception as e:
-        logging.error(f"WhisperX transcription error: {e}")
+        logging.error(f"Whisper transcription error: {e}")
         return "Error: Could not transcribe audio."
 
-
-def generate_response(text: str, tokenizer_gemma, model_gemma) -> str:
-    """Generates a response using Gemma."""
+def generate_response(text: str, model_gemma, tokenizer_gemma, device) -> str: # Pass model and tokenizer
     try:
         input_text = "Here is a response for the user. " + text
         input = tokenizer_gemma(input_text, return_tensors="pt").to(device)
@@ -88,94 +70,88 @@ 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."""
+def load_audio(audio_path: str, generator) -> torch.Tensor: #Pass generator
     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:
-            audio_tensor = torchaudio.functional.resample(
-                audio_tensor, orig_freq=sample_rate, new_freq=global_generator.sample_rate
-            )
+        audio_tensor = audio_tensor.mean(dim=0)
+        if sample_rate != generator.sample_rate:
+            audio_tensor = torchaudio.functional.resample(audio_tensor, orig_freq=sample_rate, new_freq=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
 
-
 def clear_history():
-    """Clears the conversation history"""
     global conversation_history
     conversation_history = []
     logging.info("Conversation history cleared.")
     return "Conversation history cleared."
 
+# --- Main Inference Function ---
 
-# --- 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_tokenizer_gemma, global_model_gemma, device
+@spaces.GPU(duration=gpu_timeout)  # Decorator FIRST
+def infer(user_audio) -> tuple[int, np.ndarray]:
+    # --- CUDA Availability Check (INSIDE infer) ---
+    if torch.cuda.is_available():
+        print(f"CUDA is available! Device count: {torch.cuda.device_count()}")
+        print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+        print(f"CUDA version: {torch.version.cuda}")
+        device = "cuda"
+    else:
+        print("CUDA is NOT available.  Using CPU.")  #  Use CPU, don't raise
+        device = "cpu"
 
     try:
+      # --- Model Loading (INSIDE infer, after device is set) ---
+        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.load_model("large-v2", device=device)
+        logging.info("Whisper model loaded successfully.")
+
+        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.")
+
         if not user_audio:
             raise ValueError("No audio input received.")
-
-        # 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 = 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_tokenizer_gemma, global_model_gemma) #Removed Metadata
+        return _infer(user_audio, generator, whisper_model, tokenizer_gemma, model_gemma, device) #Pass all models
     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, tokenizer_gemma, model_gemma) -> tuple:
-    """Processes the user input, generates a response, and returns audio."""
+def _infer(user_audio, generator, whisper_model, tokenizer_gemma, model_gemma, device) -> tuple[int, np.ndarray]:
     global conversation_history
 
     try:
-        # 1. ASR: Transcribe user audio using WhisperX
-        user_text = transcribe_audio(user_audio, whisper_model, model_a) #Removed Metadata
+        user_text = transcribe_audio(user_audio, whisper_model) # Pass whisper_model
         logging.info(f"User: {user_text}")
 
-        # 2. LLM: Generate a response using Gemma
-        ai_text = generate_response(user_text, tokenizer_gemma, model_gemma)
+        ai_text = generate_response(user_text, model_gemma, tokenizer_gemma, device)  # Pass model and tokenizer
         logging.info(f"AI: {ai_text}")
 
-        # 3. Generate audio using the CSM model
-        ai_audio = generator.generate(
-            text=ai_text,
-            speaker=SPEAKER_ID,
-            context=conversation_history,
-            max_audio_length_ms=30_000,
-        )
-        logging.info("Audio generated successfully.")
+        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"Sesame response generation error: {e}")
+             raise gr.Error(f"Sesame 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))
+
+        user_segment = Segment(speaker = SPEAKER_ID, text = 'User Audio', audio = load_audio(user_audio, generator)) #Pass Generator
         ai_segment =  Segment(speaker = SPEAKER_ID, text = 'AI Audio', audio = ai_audio)
         conversation_history.append(user_segment)
         conversation_history.append(ai_segment)
 
-        #Limit Conversation History
         if len(conversation_history) > MAX_CONTEXT_SEGMENTS:
             conversation_history.pop(0)
 
-        # 4. Watermarking and Audio Conversion
         audio_tensor, wm_sample_rate = watermark(
             generator._watermarker, ai_audio, generator.sample_rate, CSM_1B_HF_WATERMARK
         )
@@ -190,8 +166,7 @@ def _infer(user_audio, generator, whisper_model, model_a, tokenizer_gemma, model
         logging.exception(f"Error in _infer: {e}")
         raise gr.Error(f"An error occurred during processing: {e}")
 
-
-# --- GRADIO INTERFACE ---
+# --- Gradio Interface ---
 
 with gr.Blocks() as app:
     gr.Markdown(SPACE_INTRO_TEXT)
@@ -204,4 +179,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)
+app.launch(ssr_mode=False)