Update streamlit_app.py

streamlit_app.py

@@ -6,18 +6,13 @@ import hashlib
 import traceback
 from glob import glob
 from pathlib import Path
-import json
 import logging
-import mimetypes
-
 import yt_dlp
 import ffmpeg
 import streamlit as st
 from dotenv import load_dotenv
 from difflib import SequenceMatcher
-import requests
 
-# Try import google.generativeai, support multiple SDK shapes
 try:
     import google.generativeai as genai  # type: ignore
     genai_responses = getattr(genai, "responses", None) or getattr(genai, "Responses", None)
@@ -32,7 +27,6 @@ except Exception:
     HAS_GENAI = False
 
 load_dotenv()
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("video_ai")
 logger.propagate = False
@@ -41,6 +35,7 @@ st.set_page_config(page_title="Generate the story of videos", layout="wide")
 DATA_DIR = Path("./data")
 DATA_DIR.mkdir(exist_ok=True)
 
+# session defaults
 st.session_state.setdefault("videos", "")
 st.session_state.setdefault("loop_video", False)
 st.session_state.setdefault("uploaded_file", None)
@@ -57,6 +52,8 @@ st.session_state.setdefault("processing_timeout", 900)
 st.session_state.setdefault("generation_timeout", 300)
 st.session_state.setdefault("preferred_model", "gemini-2.0-flash-lite")
 st.session_state.setdefault("compression_threshold_mb", 200)
+st.session_state.setdefault("model_select", st.session_state.get("preferred_model"))
+st.session_state.setdefault("custom_model", "")
 
 MODEL_OPTIONS = [
     "gemini-2.5-flash",
@@ -177,7 +174,6 @@ def wait_for_processed(file_obj, timeout: int = None, progress_callback=None):
             time.sleep(backoff)
             backoff = min(backoff * 2, 8.0)
             continue
-
         state = getattr(obj, "state", None)
         state_name = getattr(state, "name", None) if state else None
         if progress_callback:
@@ -187,10 +183,8 @@ def wait_for_processed(file_obj, timeout: int = None, progress_callback=None):
                 progress_callback(min(100, pct), elapsed, state_name)
             except Exception:
                 pass
-
         if not state_name or state_name != "PROCESSING":
             return obj
-
         if time.time() - start > timeout:
             raise TimeoutError(f"File processing timed out after {int(time.time() - start)}s")
         time.sleep(backoff)
@@ -199,11 +193,11 @@ def wait_for_processed(file_obj, timeout: int = None, progress_callback=None):
 def _normalize_genai_response(response):
     if response is None:
         return ""
-    if not isinstance(response, dict):
-        try:
+    try:
+        if not isinstance(response, dict):
             response = json.loads(str(response))
-        except Exception:
-            pass
+    except Exception:
+        pass
     candidate_lists = []
     if isinstance(response, dict):
         for key in ("output", "candidates", "items", "responses", "choices"):
@@ -270,52 +264,38 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
         except Exception:
             pass
     fname = file_name_or_id(processed) or None
-
     system_msg = {"role": "system", "content": prompt_text}
     user_msg = {"role": "user", "content": f"Please summarize the attached video: {fname or '[uploaded file]'}."}
-
     call_variants = [
         ("responses.generate", {"model": model_used, "messages": [system_msg, user_msg], "files": [{"name": fname}] if fname else None, "max_output_tokens": max_tokens}),
         ("responses.generate_alt", {"model": model_used, "input": [{"text": prompt_text, "files": [{"name": fname}]}] if fname else None, "max_output_tokens": max_tokens}),
         ("legacy_create", {"model": model_used, "input": prompt_text, "file": fname, "max_output_tokens": max_tokens}),
     ]
-
     def is_transient_error(e_text: str):
         txt = str(e_text).lower()
         return any(k in txt for k in ("internal", "unavailable", "deadlineexceeded", "deadline exceeded", "timeout", "rate limit", "503", "502", "500"))
-
     start = time.time()
     last_exc = None
     backoff = 1.0
     attempts = 0
-
     while True:
         for method_name, payload in call_variants:
             attempts += 1
             try:
-                if progress_callback:
-                    progress_callback("starting", int(time.time() - start), {"model": model_used, "attempt": attempts, "method": method_name})
-
                 if genai_responses is not None and hasattr(genai_responses, "generate"):
                     payload = {k: v for k, v in payload.items() if v is not None}
                     resp = genai_responses.generate(**payload)
                     text = _normalize_genai_response(resp)
-                    if progress_callback:
-                        progress_callback("done", int(time.time() - start), {"method": method_name})
                     if text and ("please provide the video" in text.lower() or "upload the video" in text.lower()):
                         raise RuntimeError("Model indicates it didn't receive the file")
                     return text
-
                 if hasattr(genai, "Responses") and hasattr(genai.Responses, "create"):
                     payload = {k: v for k, v in payload.items() if v is not None}
                     resp = genai.Responses.create(**payload)  # type: ignore
                     text = _normalize_genai_response(resp)
-                    if progress_callback:
-                        progress_callback("done", int(time.time() - start), {"method": method_name})
                     if text and ("please provide the video" in text.lower() or "upload the video" in text.lower()):
                         raise RuntimeError("Model indicates it didn't receive the file")
                     return text
-
                 if hasattr(genai, "GenerativeModel"):
                     try:
                         model_obj = genai.GenerativeModel(model_name=model_used)
@@ -330,14 +310,11 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
                                 resp = send(prompt_text)
                             text = getattr(resp, "text", None) or str(resp)
                             text = text if text else _normalize_genai_response(resp)
-                            if progress_callback:
-                                progress_callback("done", int(time.time() - start), {"method": "GenerativeModel.chat"})
                             if text and ("please provide the video" in text.lower() or "upload the video" in text.lower()):
                                 raise RuntimeError("Model indicates it didn't receive the file")
                             return text
                     except Exception:
                         logger.exception("GenerativeModel.chat fallback failed")
-
                 raise RuntimeError("No supported response generation method available in installed google-generativeai package.")
             except Exception as e:
                 last_exc = e
@@ -374,8 +351,8 @@ def remove_prompt_echo(prompt: str, text: str, check_len: int = 600, ratio_thres
             return b_full[len(ph):].lstrip(" \n:-")
     return text
 
-# UI reset on URL change
-current_url = st.session_state.get("url", "")
+# reset on URL change
+current_url = st.session_state.get("url_input", "")
 if current_url != st.session_state.get("last_url_value"):
     st.session_state["videos"] = ""
     st.session_state["last_loaded_path"] = ""
@@ -392,46 +369,26 @@ if current_url != st.session_state.get("last_url_value"):
     st.session_state["last_url_value"] = current_url
 
 st.sidebar.header("Video Input")
-st.sidebar.text_input("Video URL", key="url_input", placeholder="https://", value=st.session_state.get("url", ""))
+st.sidebar.text_input("Video URL", key="url_input", placeholder="https://", value=st.session_state.get("url_input", ""))
 
 settings_exp = st.sidebar.expander("Settings", expanded=False)
-chosen = settings_exp.selectbox("Gemini model", MODEL_OPTIONS, index=MODEL_OPTIONS.index(st.session_state.get("preferred_model", "gemini-2.0-flash-lite")), key="model_select")
-custom_model = ""
-if settings_exp.session_state.get("model_select") == "custom":
-    custom_model = settings_exp.text_input("Custom model name", value=st.session_state.get("preferred_model", "gemini-2.0-flash-lite"), key="custom_model")
-model_input_value = (custom_model.strip() if custom_model else settings_exp.session_state.get("model_select")).strip()
+st.session_state["model_select"] = settings_exp.selectbox("Gemini model", MODEL_OPTIONS, index=MODEL_OPTIONS.index(st.session_state.get("model_select", "gemini-2.0-flash-lite")), key="model_select")
+if st.session_state.get("model_select") == "custom":
+    st.session_state["custom_model"] = settings_exp.text_input("Custom model name", value=st.session_state.get("custom_model", ""), key="custom_model")
+model_input_value = (st.session_state.get("custom_model") or st.session_state.get("model_select")).strip()
 
-settings_exp.text_input("Google API Key", key="api_key_input", value=st.session_state.get("api_key", ""), type="password")
-st.session_state["api_key"] = settings_exp.session_state.get("api_key_input", st.session_state.get("api_key", ""))
+st.session_state["api_key"] = settings_exp.text_input("Google API Key", key="api_key_input", value=st.session_state.get("api_key", ""), type="password")
 
 default_prompt = (
     "Watch the video and provide a detailed behavioral report focusing on human actions, interactions, posture, movement, and apparent intent. Keep language professional. Include a list of observations for notable events."
 )
-analysis_prompt = settings_exp.text_area("Enter analysis prompt", value=st.session_state.get("analysis_prompt", default_prompt), height=140, key="analysis_prompt")
-st.session_state["analysis_prompt"] = settings_exp.session_state.get("analysis_prompt", default_prompt)
+st.session_state["analysis_prompt"] = settings_exp.text_area("Enter analysis prompt", value=st.session_state.get("analysis_prompt", default_prompt), height=140, key="analysis_prompt")
 
-settings_exp.text_input("Video Password (if needed)", key="video_password_input", placeholder="password", type="password")
+st.session_state["video_password"] = settings_exp.text_input("Video Password (if needed)", key="video_password_input", placeholder="password", type="password")
 
-settings_exp.number_input(
-    "Processing timeout (s)", min_value=60, max_value=3600,
-    value=st.session_state.get("processing_timeout", 900), step=30,
-    key="processing_timeout_input",
-)
-st.session_state["processing_timeout"] = settings_exp.session_state.get("processing_timeout_input", st.session_state.get("processing_timeout", 900))
-
-settings_exp.number_input(
-    "Generation timeout (s)", min_value=30, max_value=1800,
-    value=st.session_state.get("generation_timeout", 300), step=10,
-    key="generation_timeout_input",
-)
-st.session_state["generation_timeout"] = settings_exp.session_state.get("generation_timeout_input", st.session_state.get("generation_timeout", 300))
-
-settings_exp.number_input(
-    "Compression threshold (MB)", min_value=10, max_value=2000,
-    value=st.session_state.get("compression_threshold_mb", 200), step=10,
-    key="compression_threshold_input",
-)
-st.session_state["compression_threshold_mb"] = settings_exp.session_state.get("compression_threshold_input", st.session_state.get("compression_threshold_mb", 200))
+st.session_state["processing_timeout"] = settings_exp.number_input("Processing timeout (s)", min_value=60, max_value=3600, value=st.session_state.get("processing_timeout", 900), step=30, key="processing_timeout_input")
+st.session_state["generation_timeout"] = settings_exp.number_input("Generation timeout (s)", min_value=30, max_value=1800, value=st.session_state.get("generation_timeout", 300), step=10, key="generation_timeout_input")
+st.session_state["compression_threshold_mb"] = settings_exp.number_input("Compression threshold (MB)", min_value=10, max_value=2000, value=st.session_state.get("compression_threshold_mb", 200), step=10, key="compression_threshold_input")
 
 key_source = "session" if st.session_state.get("api_key") else ".env" if os.getenv("GOOGLE_API_KEY") else "none"
 settings_exp.caption(f"Using API key from: {key_source}")
@@ -447,8 +404,8 @@ with col2:
 
 if st.sidebar.button("Load Video", use_container_width=True, key="load_video_btn"):
     try:
-        vpw = settings_exp.session_state.get("video_password_input", "")
-        path = download_video_ytdlp(st.session_state.get("url", settings_exp.session_state.get("url_input", "")), str(DATA_DIR), vpw)
+        vpw = st.session_state.get("video_password", "")
+        path = download_video_ytdlp(st.session_state.get("url_input", ""), str(DATA_DIR), vpw)
         st.session_state["videos"] = path
         st.session_state["last_loaded_path"] = path
         st.session_state.pop("uploaded_file", None)
@@ -468,7 +425,7 @@ if st.session_state["videos"]:
 
     with st.sidebar.expander("Options", expanded=False):
         loop_checkbox = st.checkbox("Enable Loop", value=st.session_state.get("loop_video", False), key="loop_checkbox")
-        st.session_state["loop_video"] = settings_exp.session_state.get("loop_checkbox", st.session_state.get("loop_video", False))
+        st.session_state["loop_video"] = st.session_state.get("loop_checkbox", st.session_state.get("loop_video", False))
 
         if st.button("Clear Video(s)", key="clear_videos_btn"):
             st.session_state["videos"] = ""
@@ -501,7 +458,7 @@ if st.session_state["videos"]:
     except Exception:
         pass
 
-# Generation flow
+# generation flow
 if generate_now and not st.session_state.get("busy"):
     if not st.session_state.get("videos"):
         st.error("No video loaded. Use 'Load Video' in the sidebar.")
@@ -537,7 +494,6 @@ if generate_now and not st.session_state.get("busy"):
                     if not HAS_GENAI or upload_file is None:
                         raise RuntimeError("google.generativeai SDK or upload support unavailable; cannot upload video. Use SDK with upload_file support.")
                     local_path = current_path
-
                     try:
                         file_size_mb = os.path.getsize(local_path) / (1024 * 1024)
                     except Exception:
@@ -583,7 +539,7 @@ if generate_now and not st.session_state.get("busy"):
                     st.session_state["last_loaded_path"] = current_path
                     st.session_state["file_hash"] = current_hash
 
-                prompt_text = (st.session_state.get("analysis_prompt", "") or default_prompt).strip()
+                prompt_text = (st.session_state.get("analysis_prompt", "") or "").strip() or default_prompt
                 out = ""
                 model_used = model_id
                 max_tokens = 2048 if "2.5" in model_used else 1024