Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,7 +3,6 @@ import os
 # --- ENVIRONMENT SETUP ---
 # Set cache and config directories to writable paths
 os.environ['HF_HOME'] = '/tmp/huggingface_cache'
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 os.environ['STREAMLIT_CONFIG_DIR'] = '/tmp/.streamlit'  # Fix permission error
 
 import streamlit as st
@@ -30,12 +29,13 @@ st.set_page_config(
 # --- CSS STYLING ---
 st.markdown("""
 <style>
-    /* Hide Streamlit header, footer, and menu */
+    /* Hide Streamlit header, footer, and menu - more robust selectors */
     #MainMenu {visibility: hidden;}
     footer {visibility: hidden;}
     header {visibility: hidden;}
-    .st-emotion-cache-1yfpk7 {display: none;}
-    .st-emotion-cache-1a2p8c {display: none;}
+    /* Target common Streamlit emotion cache classes (update if version changes) */
+    [data-testid="stHeader"] {display: none;}
+    [data-testid="stFooter"] {display: none;}
     /* Full-screen body */
     body {
         background-color: #f0f2f6;
@@ -137,59 +137,63 @@ def get_session_state():
 
 session_state = get_session_state()
 
-# --- MODEL LOADING ---
-def load_models():
-    if not st.session_state.models_loaded:
-        # Hide spinner for seamless experience (CSS hides it anyway)
-        with st.spinner(text=""):
-            st.session_state.models = {}
-            device = 0 if torch.cuda.is_available() else -1
-            dtype = torch.float16 if device == 0 else torch.float32
-            
-            # Speech-to-Text: Load Whisper manually to ignore mismatched sizes
-            whisper_model = WhisperForConditionalGeneration.from_pretrained(
-                "openai/whisper-small",
-                device_map="auto" if device == 0 else None,
-                dtype=dtype,
-                ignore_mismatched_sizes=True,
-                low_cpu_mem_usage=True
-            )
-            whisper_tokenizer = AutoTokenizer.from_pretrained("openai/whisper-small")
-            whisper_feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")
-            st.session_state.models['stt'] = pipeline(
-                "automatic-speech-recognition",
-                model=whisper_model,
-                tokenizer=whisper_tokenizer,
-                feature_extractor=whisper_feature_extractor,
-                device=device
-            )
+# --- MODEL LOADING WITH CACHE ---
+@st.cache_resource
+def load_stt_model():
+    device = 0 if torch.cuda.is_available() else -1
+    dtype = torch.float16 if device == 0 else torch.float32
+    
+    # Speech-to-Text: Load Whisper manually to ignore mismatched sizes
+    whisper_model = WhisperForConditionalGeneration.from_pretrained(
+        "openai/whisper-small",
+        device_map="auto" if device == 0 else None,
+        dtype=dtype,
+        ignore_mismatched_sizes=True,
+        low_cpu_mem_usage=True
+    )
+    whisper_tokenizer = AutoTokenizer.from_pretrained("openai/whisper-small")
+    whisper_feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")
+    return pipeline(
+        "automatic-speech-recognition",
+        model=whisper_model,
+        tokenizer=whisper_tokenizer,
+        feature_extractor=whisper_feature_extractor,
+        device=device
+    )
 
-            # Large Language Model: Load Qwen manually
-            qwen_model = AutoModelForCausalLM.from_pretrained(
-                "Qwen/Qwen2-0.5B-Instruct",
-                device_map="auto" if device == 0 else None,
-                dtype=dtype,
-                ignore_mismatched_sizes=True,
-                low_cpu_mem_usage=True,
-                trust_remote_code=True
-            )
-            qwen_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct", trust_remote_code=True)
-            st.session_state.models['llm'] = pipeline(
-                "text-generation",
-                model=qwen_model,
-                tokenizer=qwen_tokenizer,
-                device=device
-            )
+@st.cache_resource
+def load_llm_model():
+    device = 0 if torch.cuda.is_available() else -1
+    dtype = torch.float16 if device == 0 else torch.float32
+    
+    # Large Language Model: Load Qwen manually
+    qwen_model = AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen2-0.5B-Instruct",
+        device_map="auto" if device == 0 else None,
+        dtype=dtype,
+        ignore_mismatched_sizes=True,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True
+    )
+    qwen_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct", trust_remote_code=True)
+    return pipeline(
+        "text-generation",
+        model=qwen_model,
+        tokenizer=qwen_tokenizer,
+        device=device
+    )
 
-            # Text-to-Speech: Load TTS manually if needed, but usually no mismatch
-            st.session_state.models['tts'] = pipeline(
-                "text-to-speech", model="nineninesix/kani-tts-450m-0.1-pt", device=device
-            )
-            st.session_state.models_loaded = True
-            
-    return st.session_state.models['stt'], st.session_state.models['llm'], st.session_state.models['tts']
+@st.cache_resource
+def load_tts_model():
+    device = 0 if torch.cuda.is_available() else -1
+    return pipeline(
+        "text-to-speech", model="nineninesix/kani-tts-450m-0.1-pt", device=device
+    )
 
-stt_pipe, llm_pipe, tts_pipe = load_models()
+# Load models using cache
+stt_pipe = load_stt_model()
+llm_pipe = load_llm_model()
+tts_pipe = load_tts_model()
 
 # --- CORE FUNCTIONS ---
 def transcribe_audio(audio_bytes):
@@ -246,6 +250,9 @@ def update_ui_status(status):
     circle_class = status
     circle_placeholder.markdown(f'<div class="voice-circle {circle_class}"></div>', unsafe_allow_html=True)
 
+# Debug: Temporary message to confirm rendering (remove after testing)
+st.write("App loaded - if you see this, basic rendering works. Circle should appear below.")
+
 # --- APP LOGIC: Auto mode ---
 # Always listening - no chat history shown