Update app.py

app.py

@@ -15,7 +15,7 @@ from transformers import GenerationConfig
 # Configure logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# --- Authentication and Configuration ---  (Moved BEFORE model loading)
+# --- Authentication and Configuration ---
 try:
     api_key = os.getenv("HF_TOKEN")
     if not api_key:
@@ -39,9 +39,7 @@ 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 --- (Moved INSIDE infer function)
-
-# --- Constants ---  (Constants can stay outside)
+# --- Constants ---
 SPEAKER_ID = 0
 MAX_CONTEXT_SEGMENTS = 3
 MAX_GEMMA_LENGTH = 128
@@ -49,6 +47,45 @@ MAX_GEMMA_LENGTH = 128
 # --- Global Conversation History ---
 conversation_history = []
 
+# --- Model Downloading (PRE-DOWNLOAD) ---
+
+# Download Sesame CSM 1B
+csm_1b_model_path = "csm_1b_ckpt.pt"  # Local path
+try:
+    if not os.path.exists(csm_1b_model_path):
+        hf_hub_download(repo_id="sesame/csm-1b", filename="ckpt.pt", local_dir=".", local_dir_use_symlinks=False)
+        os.rename("ckpt.pt", csm_1b_model_path)  # Rename to avoid confusion
+        logging.info("Sesame CSM 1B model downloaded.")
+    else:
+        logging.info("Sesame CSM 1B model already downloaded.")
+except Exception as e:
+    logging.error(f"Error downloading Sesame CSM 1B: {e}")
+    raise
+
+# Download Whisper (using the built-in download mechanism)
+whisper_model_name = "small.en"
+try:
+    whisper.load_model(whisper_model_name)  # This downloads if not already present
+    logging.info(f"Whisper model '{whisper_model_name}' downloaded/loaded.")
+except Exception as e:
+    logging.error(f"Error downloading Whisper model: {e}")
+    raise
+
+# Download Gemma 3 1B (Tokenizer and Model)
+gemma_repo_id = "google/gemma-3-1b-it"
+gemma_local_path = "gemma_model" # Using a directory
+try:
+    if not os.path.exists(gemma_local_path):
+        tokenizer_gemma = AutoTokenizer.from_pretrained(gemma_repo_id, cache_dir=gemma_local_path)  #downloads
+        model_gemma = AutoModelForCausalLM.from_pretrained(gemma_repo_id, cache_dir=gemma_local_path) #downloads
+        logging.info("Gemma 3 1B model and tokenizer downloaded.")
+    else:
+        logging.info("Gemma 3 1B model and tokenizer already downloaded.")
+except Exception as e:
+    logging.error(f"Error downloading Gemma 3 1B: {e}")
+    raise
+
+
 # --- Helper Functions ---
 
 def transcribe_audio(audio_path: str, whisper_model) -> str:  # Pass whisper_model
@@ -85,10 +122,6 @@ def generate_response(text: str, model_gemma, tokenizer_gemma, device) -> str: #
             assistant_response = decoded_output[start_index:].strip()
             return assistant_response
         return decoded_output
-        #input_text = "Reapond to the users prompt: " + 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)
-        #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."
@@ -116,25 +149,26 @@ def clear_history():
 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"
+        logging.info(f"CUDA is available! Using device: {torch.cuda.get_device_name(0)}")
     else:
-        print("CUDA is NOT available.  Using CPU.")  #  Use CPU, don't raise
         device = "cpu"
+        logging.info("CUDA is NOT available.  Using 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)
+        # Load Sesame CSM 1B (from local file)
+        generator = load_csm_1b(csm_1b_model_path, device)
         logging.info("Sesame CSM 1B loaded successfully.")
 
-        whisper_model = whisper.load_model("small.en", device=device)
-        logging.info("Whisper model loaded successfully.")
+        # Load Whisper (from local cache or downloaded)
+        whisper_model = whisper.load_model(whisper_model_name, device=device)
+        logging.info(f"Whisper model '{whisper_model_name}' loaded successfully.")
 
-        tokenizer_gemma = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
-        model_gemma = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it").to(device)
+        # Load Gemma (from local cache)
+        tokenizer_gemma = AutoTokenizer.from_pretrained(gemma_local_path)
+        model_gemma = AutoModelForCausalLM.from_pretrained(gemma_local_path).to(device)
         logging.info("Gemma 3 1B pt model loaded successfully.")
 
         if not user_audio: