Update app.py

app.py

@@ -4,6 +4,7 @@ import subprocess
 import tempfile
 import base64
 import json
+import mimetypes # New import for mimetype guessing
 from io import BytesIO
 from typing import List, Tuple, Optional, Set
 import requests
@@ -11,60 +12,11 @@ from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
 import time
 import atexit
-from requests.exceptions import RequestException, HTTPError # Import HTTPError for requests fallback
-
-# --- Mistral Client Import & Placeholder for graceful degradation ---
-_MISTRAL_CLIENT_INSTALLED = False
-try:
-    from mistralai import Mistral
-    from mistralai.exceptions import MistralAPIException
-    _MISTRAL_CLIENT_INSTALLED = True
-except ImportError:
-    print(
-        "Warning: Mistral AI client library ('mistralai') not found. "
-        "Please install it with 'pip install mistralai' to enable full AI analysis features. " # Updated message
-        "The application will launch, but API calls will fall back to direct HTTP requests "
-        "if an API key is provided." # Updated message to reflect fallback
-    )
-    # Define placeholder classes to prevent NameErrors and provide clear messages
-    class MistralAPIException(Exception):
-        """A placeholder for mistralai.exceptions.MistralAPIException."""
-        def __init__(self, message: str, status_code: Optional[int] = None):
-            super().__init__(message)
-            self.message = message
-            self.status_code = status_code or 500
-        def __str__(self):
-            return f"MistralAPIException (Status: {self.status_code}): {self.message}"
-
-    class _DummyMistralChatClient:
-        """Placeholder for Mistral client's chat interface."""
-        def complete(self, *args, **kwargs):
-            # This method will typically not be called if _MISTRAL_CLIENT_INSTALLED is False,
-            # as the `chat_complete` function will use the requests fallback instead.
-            raise MistralAPIException(
-                "Mistral AI chat client is unavailable. "
-                "Please install 'mistralai' with 'pip install mistralai'.",
-                status_code=500
-            )
-    class _DummyMistralFilesClient:
-        """Placeholder for Mistral client's files interface."""
-        def upload(self, *args, **kwargs):
-            # This method will typically not be called if _MISTRAL_CLIENT_INSTALLED is False.
-            raise MistralAPIException(
-                "Mistral AI files client is unavailable. "
-                "Please install 'mistralai' with 'pip install mistralai'.",
-                status_code=500
-            )
-    class Mistral:
-        """A placeholder for the Mistral client if the library is not installed."""
-        def __init__(self, api_key: str = "", *args, **kwargs): # Added api_key to store it for fallback
-            self.api_key = api_key # Store the API key
-        @property
-        def chat(self):
-            return _DummyMistralChatClient()
-        @property
-        def files(self):
-            return _DummyMistralFilesClient()
+from requests.exceptions import RequestException # Keep for general network errors in fetch_bytes
+
+# --- Mistral Client Import (Assume installed as requested) ---
+from mistralai import Mistral
+from mistralai.exceptions import MistralAPIException
 
 # --- Configuration and Globals ---
 DEFAULT_MISTRAL_KEY = os.getenv("MISTRAL_API_KEY", "")
@@ -89,13 +41,12 @@ Image.MAX_IMAGE_PIXELS = 10000 * 10000
 
 DEFAULT_HEADERS = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"}
 
-# --- Temporary File Cleanup (Changed to use a set) ---
+# --- Temporary File Cleanup ---
 _temp_files_to_delete: Set[str] = set() # Use a set for better management
 
 def _cleanup_all_temp_files():
     """Removes all temporary files created upon application exit."""
-    # Create a copy to iterate while modifying the original set
-    for f_path in _temp_files_to_delete.copy():
+    for f_path in list(_temp_files_to_delete): # Iterate over a copy to allow modification
         if os.path.exists(f_path):
             try:
                 os.remove(f_path)
@@ -107,11 +58,10 @@ def _cleanup_all_temp_files():
 atexit.register(_cleanup_all_temp_files)
 
 # --- Mistral Client and API Helpers ---
-def get_client(api_key: Optional[str] = None):
+def get_client(api_key: Optional[str] = None) -> Mistral:
     """
     Returns a Mistral client instance. If the API key is missing, a MistralAPIException is raised.
-    If the client library is not installed, a placeholder client is returned, and API calls
-    will fall back to direct HTTP requests.
+    Assumes mistralai client library is installed.
     """
     key_to_use = (api_key or "").strip() or DEFAULT_MISTRAL_KEY
     if not key_to_use:
@@ -119,10 +69,6 @@ def get_client(api_key: Optional[str] = None):
             "Mistral API key is not set. Please provide it in the UI or as MISTRAL_API_KEY environment variable.",
             status_code=401 # Unauthorized
         )
-
-    # Always return a Mistral client instance.
-    # If _MISTRAL_CLIENT_INSTALLED is True, this will be the real Mistral client.
-    # Otherwise, it's the placeholder that has the api_key stored.
     return Mistral(api_key=key_to_use)
 
 def is_remote(src: str) -> bool:
@@ -176,7 +122,6 @@ def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD_BYTES, timeou
                     fd, p = tempfile.mkstemp(suffix=ext_from_src(src) or ".tmp")
                     os.close(fd)
                     try:
-                        # FIX: Open file for writing before the streaming loop
                         with open(p, "wb") as fh_write:
                             with requests.get(src, timeout=timeout, stream=True, headers=DEFAULT_HEADERS) as r:
                                 r.raise_for_status()
@@ -187,12 +132,11 @@ def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD_BYTES, timeou
                                         fh_write.write(chunk)
                                         downloaded_size += len(chunk)
                                         if progress is not None and total_size > 0:
-                                            # Scale progress from 0.1 to 0.25 for streaming phase
                                             progress(0.1 + (downloaded_size / total_size) * 0.15)
-                        with open(p, "rb") as fh_read: # Open again to read full content
+                        with open(p, "rb") as fh_read:
                             return fh_read.read()
                     finally:
-                        try: os.remove(p)
+                        try: _temp_files_to_delete.discard(p); os.remove(p)
                         except Exception as e: print(f"Error during streaming temp file cleanup {p}: {e}")
             except Exception as e:
                 print(f"Warning: Streaming download failed for {src}: {e}. Falling back to non-streaming.")
@@ -223,7 +167,7 @@ def convert_to_jpeg_bytes(img_bytes: bytes, base_h: int = 480) -> bytes:
         return b""
 
     try:
-        if getattr(img, "is_animated", False): # Handle animated images (e.g., GIFs) by taking the first frame
+        if getattr(img, "is_animated", False):
             img.seek(0)
     except Exception:
         pass
@@ -253,11 +197,11 @@ def _ffprobe_streams(path: str) -> Optional[dict]:
         if os.path.exists(potential_ffprobe_in_dir) and os.access(potential_ffprobe_in_dir, os.X_OK):
             ffprobe_path = potential_ffprobe_in_dir
 
-    if not ffprobe_path: # Fallback to checking PATH if not found next to ffmpeg
+    if not ffprobe_path:
         ffprobe_path = shutil.which("ffprobe")
 
     if not ffprobe_path:
-        return None # ffprobe is not available
+        return None
 
     cmd = [
         ffprobe_path, "-v", "error", "-print_format", "json", "-show_streams", "-show_format", path
@@ -281,13 +225,13 @@ def _get_video_info_and_timestamps(media_path: str, sample_count: int) -> Tuple[
 
     timestamps: List[float] = []
     if duration > 0 and sample_count > 0:
-        actual_sample_count = min(sample_count, max(1, int(duration))) # Ensure sample_count doesn't exceed duration
+        actual_sample_count = min(sample_count, max(1, int(duration)))
         if actual_sample_count > 0:
             step = duration / (actual_sample_count + 1)
             timestamps = [step * (i + 1) for i in range(actual_sample_count)]
 
-    if not timestamps: # Fallback for very short videos or if duration couldn't be determined
-        timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count] # Try fixed early timestamps
+    if not timestamps:
+        timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count]
 
     return info, timestamps
 
@@ -299,7 +243,7 @@ def extract_frames_for_model_and_gallery(media_path: str, sample_count: int = 5,
     frames_for_model: List[bytes] = []
     frame_paths_for_gallery: List[str] = []
 
-    if not FFMPEG_BIN: # Check FFMPEG_BIN here
+    if not FFMPEG_BIN:
         print(f"Warning: FFMPEG not found. Cannot extract frames for {media_path}.")
         return frames_for_model, frame_paths_for_gallery
     if not os.path.exists(media_path):
@@ -361,7 +305,7 @@ def extract_frames_for_model_and_gallery(media_path: str, sample_count: int = 5,
         progress(0.45, desc=f"Extracted {len(frames_for_model)} frames for analysis and gallery")
     return frames_for_model, frame_paths_for_gallery
 
-def chat_complete(client, model: str, messages, timeout: int = 120, progress=None) -> str:
+def chat_complete(client: Mistral, model: str, messages, timeout: int = 120, progress=None) -> str:
     """Sends messages to the Mistral chat completion API with retry logic."""
     max_retries = 5
     initial_delay = 1.0
@@ -370,39 +314,25 @@ def chat_complete(client, model: str, messages, timeout: int = 120, progress=Non
             if progress is not None:
                 progress(0.6 + 0.01 * attempt, desc=f"Sending request to model (attempt {attempt+1}/{max_retries})...")
 
-            res = None
-            if _MISTRAL_CLIENT_INSTALLED:
-                # Use the real Mistral client's chat.complete method
-                res = client.chat.complete(model=model, messages=messages, stream=False, timeout_ms=timeout * 1000)
-            else:
-                # Fallback to direct HTTP request if client library not installed
-                api_key = getattr(client, "api_key", "") # Get key from client, should always be present now
-                if not api_key: # Double check, though get_client already ensures this
-                    return "Error: Mistral API key is not set for fallback."
-
-                url = "https://api.mistral.ai/v1/chat/completions"
-                headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-                payload = {"model": model, "messages": messages, "stream": False} # Removed timeout_ms for requests
-                r = requests.post(url, json=payload, headers=headers, timeout=timeout) # requests timeout in seconds
-                r.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
-                res = r.json()
+            # Always use the real Mistral client's chat.complete method
+            res = client.chat.complete(model=model, messages=messages, stream=False, timeout_ms=timeout * 1000)
 
             if progress is not None:
                 progress(0.8, desc="Model responded, parsing...")
 
-            # Handle both object-style (from mistralai client) and dict-style (from requests.json()) responses
-            choices = getattr(res, "choices", None) or (res.get("choices") if isinstance(res, dict) else [])
+            # Access attributes directly from the client's response object
+            choices = getattr(res, "choices", [])
 
             if not choices:
                 return f"Empty response from model: {res}"
 
             first = choices[0]
-            msg = getattr(first, "message", None) or (first.get("message") if isinstance(first, dict) else first)
-            content = getattr(msg, "content", None) or (msg.get("content") if isinstance(msg, dict) else None)
+            msg = getattr(first, "message", None)
+            content = getattr(msg, "content", None)
             return content.strip() if isinstance(content, str) else str(content)
 
-        except (MistralAPIException, HTTPError) as e: # Catch both client lib and requests HTTP errors
-            status_code = getattr(e, "status_code", None) or (e.response.status_code if isinstance(e, HTTPError) else None)
+        except MistralAPIException as e:
+            status_code = getattr(e, "status_code", None)
             message = getattr(e, "message", str(e))
 
             if status_code == 429 and attempt < max_retries - 1:
@@ -411,7 +341,7 @@ def chat_complete(client, model: str, messages, timeout: int = 120, progress=Non
                 time.sleep(delay)
             else:
                 return f"Error: Mistral API error occurred ({status_code if status_code else 'unknown'}): {message}"
-        except RequestException as e: # Catch other requests errors (e.g., connection issues, timeout)
+        except RequestException as e: # Catch network issues that MistralAPIException might not fully wrap
             if attempt < max_retries - 1:
                 delay = initial_delay * (2 ** attempt)
                 print(f"Network/API request failed: {e}. Retrying in {delay:.2f}s...")
@@ -423,7 +353,7 @@ def chat_complete(client, model: str, messages, timeout: int = 120, progress=Non
 
     return "Error: Maximum retries reached for API call."
 
-def upload_file_to_mistral(client, path: str, filename: str | None = None, purpose: str = "batch", timeout: int = 120, progress=None) -> str:
+def upload_file_to_mistral(client: Mistral, path: str, filename: str | None = None, purpose: str = "batch", timeout: int = 120, progress=None) -> str:
     """Uploads a file to the Mistral API, returning its file ID."""
     fname = filename or os.path.basename(path)
     max_retries = 3
@@ -433,38 +363,23 @@ def upload_file_to_mistral(client, path: str, filename: str | None = None, purpo
             if progress is not None:
                 progress(0.5 + 0.01 * attempt, desc=f"Uploading file to model service (attempt {attempt+1}/{max_retries})...")
 
-            fid = None
-            if _MISTRAL_CLIENT_INSTALLED:
-                with open(path, "rb") as fh:
-                    # Mistral client's file upload expects (filename, file_like_object) for 'file' param
-                    res = client.files.upload(file=(fname, fh), purpose=purpose) 
-                fid = getattr(res, "id", None) or (res.get("id") if isinstance(res, dict) else None)
-                if not fid:
-                    raise RuntimeError(f"Mistral API upload response missing file ID from client: {res}")
-            else:
-                # Fallback to direct HTTP request
-                api_key = getattr(client, "api_key", "")
-                if not api_key:
-                    raise RuntimeError("Mistral API key is not set for file upload fallback.")
-
-                url = "https://api.mistral.ai/v1/files"
-                headers = {"Authorization": f"Bearer {api_key}"}
-                with open(path, "rb") as fh:
-                    files = {"file": (fname, fh)}
-                    data = {"purpose": purpose}
-                    r = requests.post(url, headers=headers, files=files, data=data, timeout=timeout)
-                    r.raise_for_status()
-                    jr = r.json()
-                    fid = jr.get("id") or jr.get("data", [{}])[0].get("id") # Handle potential nested 'data' from older API
-                    if not fid:
-                        raise RuntimeError(f"Mistral API upload response missing file ID from direct request: {jr}")
+            # Guess mimetype for robust upload
+            mimetype, _ = mimetypes.guess_type(fname)
+            mimetype = mimetype or "application/octet-stream"
+
+            with open(path, "rb") as fh:
+                # Mistral client's file upload expects (filename, file_like_object, mimetype) for 'file' param
+                res = client.files.upload(file=(fname, fh, mimetype), purpose=purpose)
+            fid = getattr(res, "id", None)
+            if not fid:
+                raise RuntimeError(f"Mistral API upload response missing file ID: {res}")
 
             if progress is not None:
                 progress(0.6, desc="Upload complete")
             return fid
 
-        except (MistralAPIException, HTTPError) as e:
-            status_code = getattr(e, "status_code", None) or (e.response.status_code if isinstance(e, HTTPError) else None)
+        except MistralAPIException as e:
+            status_code = getattr(e, "status_code", None)
             message = getattr(e, "message", str(e))
             if status_code == 429 and attempt < max_retries - 1:
                 delay = initial_delay * (2 ** attempt)
@@ -507,7 +422,7 @@ def determine_media_type(src: str, progress=None) -> Tuple[bool, bool]:
         progress(0.02, desc="Determined media type (initial hint)")
     return is_image, is_video
 
-def analyze_image_structured(client, img_bytes: bytes, prompt: str, progress=None) -> str:
+def analyze_image_structured(client: Mistral, img_bytes: bytes, prompt: str, progress=None) -> str:
     """Analyzes an image using the PixTRAL model."""
     try:
         if progress is not None:
@@ -529,14 +444,13 @@ def analyze_image_structured(client, img_bytes: bytes, prompt: str, progress=Non
     except Exception as e:
         return f"Error analyzing image: {e}"
 
-def analyze_video_cohesive(client, video_path: str, prompt: str, progress=None) -> Tuple[str, List[str]]:
+def analyze_video_cohesive(client: Mistral, video_path: str, prompt: str, progress=None) -> Tuple[str, List[str]]:
     """
     Analyzes a video using the VoxTRAL model (if available) or by extracting frames
     and using PixTRAL as a fallback.
     Returns: (analysis result text, list of paths to gallery frames)
     """
     gallery_frame_paths: List[str] = []
-    # If FFmpeg is not available, we can't do video analysis at all
     if not FFMPEG_BIN:
         return "Error: FFmpeg is not found in your system PATH. Video analysis and preview are unavailable.", []
 
@@ -548,12 +462,12 @@ def analyze_video_cohesive(client, video_path: str, prompt: str, progress=None)
         messages = [
             {"role": "system", "content": SYSTEM_INSTRUCTION},
             {"role": "user", "content": [
-                {"type": "text", "text": f"Uploaded video file id: {file_id}\n\nInstruction: Analyze the entire video and produce a single cohesive narrative describing consistent observations.\n\n{prompt}"},
+                {"type": "video", "id": file_id}, # Correct format for video input
+                {"type": "text", "text": f"Instruction: Analyze the entire video and produce a single cohesive narrative describing consistent observations.\n\n{prompt}"},
             ]},
         ]
         result = chat_complete(client, VIDEO_MODEL, messages, progress=progress)
 
-        # Still extract frames for gallery even if full video upload was successful
         _, gallery_frame_paths = extract_frames_for_model_and_gallery(
             video_path, sample_count=6, gallery_base_h=1080, model_base_h=1024, progress=progress
         )
@@ -581,7 +495,6 @@ def analyze_video_cohesive(client, video_path: str, prompt: str, progress=None)
                     "meta": {"frame_index": i},
                 }
             )
-        # Consolidate frames for a cohesive narrative, as per requirement
         content = [{"type": "text", "text": prompt + "\n\nPlease consolidate observations across these frames into a single cohesive narrative."}] + image_entries
         messages = [
             {"role": "system", "content": SYSTEM_INSTRUCTION},
@@ -599,13 +512,12 @@ def _convert_video_for_preview_if_needed(path: str) -> str:
     if not FFMPEG_BIN or not os.path.exists(path):
         return path
 
-    # Check if it's already a web-friendly MP4 (H.264/AVC1, H.265)
     if path.lower().endswith((".mp4", ".m4v")):
         info = _ffprobe_streams(path)
         if info:
             video_streams = [s for s in info.get("streams", []) if s.get("codec_type") == "video"]
             if video_streams and any(s.get("codec_name") in ("h264", "h265", "avc1") for s in video_streams):
-                return path # Already compatible, no conversion needed
+                return path
 
     out_path = _temp_file(b"", suffix=".mp4")
     if not out_path:
@@ -614,18 +526,17 @@ def _convert_video_for_preview_if_needed(path: str) -> str:
 
     cmd = [
         FFMPEG_BIN, "-y", "-i", path,
-        "-c:v", "libx264", "-preset", "veryfast", "-crf", "28", # H.264 codec
-        "-c:a", "aac", "-b:a", "128k", # AAC audio
-        "-movflags", "+faststart", out_path, # Optimize for web streaming
-        "-map_metadata", "-1" # Remove metadata
+        "-c:v", "libx264", "-preset", "veryfast", "-crf", "28",
+        "-c:a", "aac", "-b:a", "128k",
+        "-movflags", "+faststart", out_path,
+        "-map_metadata", "-1"
     ]
     try:
         subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=60)
         return out_path
     except Exception as e:
         print(f"Error converting video for preview: {e}")
-        # If conversion fails, remove the failed temp file and return original path
-        _temp_files_to_delete.discard(out_path) # Remove from set
+        _temp_files_to_delete.discard(out_path)
         try: os.remove(out_path)
         except Exception: pass
         return path
@@ -645,14 +556,14 @@ def _get_playable_preview_path_from_raw(src_url: str, raw_bytes: bytes, is_image
         Image.open(BytesIO(raw_bytes)).verify()
         is_actually_image = True
     except (UnidentifiedImageError, Exception):
-        pass # Not a verifiable image, proceed to video or fallback logic
+        pass
 
     if is_actually_image:
         jpeg_bytes = convert_to_jpeg_bytes(raw_bytes, base_h=1024)
         if jpeg_bytes:
             return _temp_file(jpeg_bytes, suffix=".jpg")
-        return "" # Failed image conversion
-    elif is_video_hint: # If it's not an image, and was hinted as video
+        return ""
+    elif is_video_hint:
         temp_raw_video_path = _temp_file(raw_bytes, suffix=ext_from_src(src_url) or ".mp4")
         if not temp_raw_video_path:
             print(f"Error: Failed to create temporary raw video file for {src_url}.")
@@ -660,7 +571,7 @@ def _get_playable_preview_path_from_raw(src_url: str, raw_bytes: bytes, is_image
 
         playable_path = _convert_video_for_preview_if_needed(temp_raw_video_path)
         return playable_path
-    elif is_image_hint: # Fallback: if hinted image but not verifiable, still try conversion
+    elif is_image_hint:
          jpeg_bytes = convert_to_jpeg_bytes(raw_bytes, base_h=1024)
          if jpeg_bytes:
              return _temp_file(jpeg_bytes, suffix=".jpg")
@@ -679,22 +590,15 @@ def _get_button_label_for_status(status: str) -> str:
 
 def create_demo():
     """Creates the Gradio interface for Flux Multimodal analysis."""
-    # Determine FFMPEG status once for UI message
     ffmpeg_status_message = ""
     if not FFMPEG_BIN:
         ffmpeg_status_message = "🔴 FFmpeg not found! Video analysis and preview will be limited/unavailable."
     else:
         ffmpeg_status_message = "🟢 FFmpeg found. Video features enabled."
 
-    mistral_client_status_message = ""
-    if not _MISTRAL_CLIENT_INSTALLED:
-        mistral_client_status_message = "🟡 Mistral AI client ('mistralai') not installed. AI analysis will fall back to direct HTTP requests. Run `pip install mistralai` for full features." # Updated message
-    else:
-        mistral_client_status_message = "🟢 Mistral AI client found."
-
     with gr.Blocks(title="Flux Multimodal", css=GRADIO_CSS) as demo:
         gr.Markdown("# Flux Multimodal AI Assistant")
-        gr.Markdown(f"{mistral_client_status_message}<br>{ffmpeg_status_message}") # Display dependency status
+        gr.Markdown(f"🟢 Mistral AI client found.<br>{ffmpeg_status_message}") # Simplified status
 
         with gr.Row():
             with gr.Column(scale=1):
@@ -715,23 +619,22 @@ def create_demo():
                 output_markdown = gr.Markdown("")
 
                 status_state = gr.State("idle")
-                main_preview_path_state = gr.State("") # Path to the playable preview file (image or video)
-                screenshot_paths_state = gr.State([]) # List of paths to extracted frames for gallery
-                raw_media_path_state = gr.State("") # Path to the raw downloaded media file for analysis
+                main_preview_path_state = gr.State("")
+                screenshot_paths_state = gr.State([])
+                raw_media_path_state = gr.State("")
 
         def clear_all_ui_and_files_handler():
             """
             Cleans up all tracked temporary files and resets all relevant UI components and states.
             """
-            # Iterate a copy of the set while potentially modifying the original
-            for f_path in _temp_files_to_delete.copy():
+            for f_path in list(_temp_files_to_delete):
                 if os.path.exists(f_path):
                     try:
                         os.remove(f_path)
-                        _temp_files_to_delete.discard(f_path) # Remove from set
+                        _temp_files_to_delete.discard(f_path)
                     except Exception as e:
                         print(f"Error during proactive cleanup of {f_path}: {e}")
-            _temp_files_to_delete.clear() # Ensure set is empty
+            _temp_files_to_delete.clear()
 
             return "", \
                    gr.update(value=None, visible=False), \
@@ -775,22 +678,20 @@ def create_demo():
             Loads media from URL, generates a preview, and sets up temporary files for analysis.
             Also handles cleanup of previously loaded media.
             """
-            # --- Proactive cleanup of old files related to previous load ---
             if current_main_preview_path and os.path.exists(current_main_preview_path):
-                _temp_files_to_delete.discard(current_main_preview_path) # Remove from set
+                _temp_files_to_delete.discard(current_main_preview_path)
                 try: os.remove(current_main_preview_path)
                 except Exception as e: print(f"Error cleaning up old temp file {current_main_preview_path}: {e}")
             if current_raw_media_path and os.path.exists(current_raw_media_path):
-                _temp_files_to_delete.discard(current_raw_media_path) # Remove from set
+                _temp_files_to_delete.discard(current_raw_media_path)
                 try: os.remove(current_raw_media_path)
                 except Exception as e: print(f"Error cleaning up old temp file {current_raw_media_path}: {e}")
             for path in current_screenshot_paths:
                 if path and os.path.exists(path):
-                    _temp_files_to_delete.discard(path) # Remove from set
+                    _temp_files_to_delete.discard(path)
                     try: os.remove(path)
                     except Exception as e: print(f"Error cleaning up old temp file {path}: {e}")
 
-            # Default cleared states for UI and backend values to be returned on empty URL or error
             img_update_clear = gr.update(value=None, visible=False)
             video_update_clear = gr.update(value=None, visible=False)
             gallery_update_clear = gr.update(value=[], visible=False)
@@ -806,26 +707,24 @@ def create_demo():
             temp_raw_path_for_analysis = ""
             try:
                 progress(0.01, desc="Downloading media for preview and analysis...")
-                raw_bytes_for_analysis = fetch_bytes(url, timeout=60, progress=progress) # Pass progress here
+                raw_bytes_for_analysis = fetch_bytes(url, timeout=60, progress=progress)
                 if not raw_bytes_for_analysis:
                     return img_update_clear, video_update_clear, gallery_update_clear, \
                            gr.update(value="Preview load failed: No media bytes fetched.", visible=True), \
                            main_path_clear, raw_media_path_clear, screenshot_paths_clear
 
-                # Store raw bytes in a temp file for potential analysis (especially video uploads or large image processing)
                 temp_raw_path_for_analysis = _temp_file(raw_bytes_for_analysis, suffix=ext_from_src(url) or ".tmp")
                 if not temp_raw_path_for_analysis:
                     return img_update_clear, video_update_clear, gallery_update_clear, \
                            gr.update(value="Preview load failed: Could not save raw media to temp file.", visible=True), \
                            main_path_clear, raw_media_path_clear, screenshot_paths_clear
 
-                progress(0.25, desc="Generating playable preview...") # Adjusted progress start
+                progress(0.25, desc="Generating playable preview...")
                 is_img_initial, is_vid_initial = determine_media_type(url)
                 local_playable_path = _get_playable_preview_path_from_raw(url, raw_bytes_for_analysis, is_img_initial, is_vid_initial)
 
                 if not local_playable_path:
-                    # If preview failed, cleanup the temp_raw_path_for_analysis as well
-                    _temp_files_to_delete.discard(temp_raw_path_for_analysis) # Remove from set
+                    _temp_files_to_delete.discard(temp_raw_path_for_analysis)
                     try: os.remove(temp_raw_path_for_analysis)
                     except Exception as e: print(f"Error during cleanup of raw temp file {temp_raw_path_for_analysis}: {e}")
 
@@ -846,12 +745,10 @@ def create_demo():
                            gallery_update_clear, gr.update(value="Video preview loaded.", visible=True), \
                            local_playable_path, temp_raw_path_for_analysis, screenshot_paths_clear
                 else:
-                    # If local_playable_path exists but is not image/video, clean it up
-                    _temp_files_to_delete.discard(local_playable_path) # Remove from set
+                    _temp_files_to_delete.discard(local_playable_path)
                     try: os.remove(local_playable_path)
                     except Exception as e: print(f"Error during cleanup of unplayable temp file {local_playable_path}: {e}")
-                    # Also clean up raw_media_path if the playable path was not generated successfully
-                    _temp_files_to_delete.discard(temp_raw_path_for_analysis) # Remove from set
+                    _temp_files_to_delete.discard(temp_raw_path_for_analysis)
                     try: os.remove(temp_raw_path_for_analysis)
                     except Exception as e: print(f"Error during cleanup of raw temp file {temp_raw_path_for_analysis}: {e}")
 
@@ -860,9 +757,8 @@ def create_demo():
                            main_path_clear, raw_media_path_clear, screenshot_paths_clear
 
             except Exception as e:
-                # If an error occurred during loading, clear all relevant paths.
                 if os.path.exists(temp_raw_path_for_analysis):
-                    _temp_files_to_delete.discard(temp_raw_path_for_analysis) # Remove from set
+                    _temp_files_to_delete.discard(temp_raw_path_for_analysis)
                     try: os.remove(temp_raw_path_for_analysis)
                     except Exception as ex: print(f"Error during cleanup of raw temp file {temp_raw_path_for_analysis} on error: {ex}")
 
@@ -876,7 +772,7 @@ def create_demo():
             outputs=[preview_image, preview_video, screenshot_gallery, preview_status_text, main_preview_path_state, raw_media_path_state, screenshot_paths_state]
         )
 
-        def worker(url: str, prompt: str, key: str, current_main_preview_path: str, raw_media_path: str, progress=gr.Progress()):
+        def worker(url: str, prompt: str, key: str, raw_media_path: str, progress=gr.Progress()):
             """
             The main worker function that performs media analysis using Mistral models.
             """
@@ -885,43 +781,36 @@ def create_demo():
 
             try:
                 if not raw_media_path or not os.path.exists(raw_media_path):
-                    return "error", "**Error:** No raw media file available for analysis. Please load a URL first.", current_main_preview_path, []
+                    return "error", "**Error:** No raw media file available for analysis. Please load a URL first.", [], []
 
-                # Initial check for FFmpeg for video processing
                 if not FFMPEG_BIN:
-                    # If this is a video, and ffmpeg is missing, return early.
-                    # Otherwise, proceed for image analysis.
-                    # Determine media type by extension from raw_media_path
                     ext = ext_from_src(raw_media_path)
                     if ext in VIDEO_EXTENSIONS:
-                         return "error", "**Error:** FFmpeg is not found in your system PATH. Video analysis is unavailable. Please install FFmpeg.", current_main_preview_path, []
+                         return "error", "**Error:** FFmpeg is not found in your system PATH. Video analysis is unavailable. Please install FFmpeg.", [], []
 
                 with open(raw_media_path, "rb") as f:
                     raw_bytes_for_analysis = f.read()
 
                 if not raw_bytes_for_analysis:
-                    return "error", "**Error:** Raw media file is empty for analysis.", current_main_preview_path, []
+                    return "error", "**Error:** Raw media file is empty for analysis.", [], []
 
                 progress(0.01, desc="Starting media analysis...")
 
                 is_actually_image_for_analysis = False
                 is_actually_video_for_analysis = False
 
-                # Determine media type for analysis robustly
                 try:
                     Image.open(BytesIO(raw_bytes_for_analysis)).verify()
                     is_actually_image_for_analysis = True
                 except UnidentifiedImageError:
-                    # If PIL can't identify it as an image, check if it has a video extension.
                     if ext_from_src(raw_media_path) in VIDEO_EXTENSIONS:
                         is_actually_video_for_analysis = True
                 except Exception as e:
-                    # Catch other PIL errors (e.g., truncated, memory, etc.).
                     print(f"Warning: PIL error during image verification for raw analysis media ({raw_media_path}): {e}. Checking for video extension.")
                     if ext_from_src(raw_media_path) in VIDEO_EXTENSIONS:
                         is_actually_video_for_analysis = True
 
-                client = get_client(key) # This will raise MistralAPIException if key missing, but NOT if lib not installed
+                client = get_client(key)
 
                 if is_actually_video_for_analysis:
                     progress(0.25, desc="Running full-video analysis")
@@ -930,21 +819,20 @@ def create_demo():
                     progress(0.20, desc="Running image analysis")
                     result_text = analyze_image_structured(client, raw_bytes_for_analysis, prompt, progress=progress)
                 else:
-                    return "error", "Error: Could not definitively determine media type for analysis after byte inspection and extension check. Please check the URL/file content.", current_main_preview_path, []
+                    return "error", "Error: Could not definitively determine media type for analysis after byte inspection and extension check. Please check the URL/file content.", [], []
 
                 status = "done" if not (isinstance(result_text, str) and result_text.lower().startswith("error")) else "error"
-                return status, result_text, current_main_preview_path, generated_screenshot_paths
+                return status, result_text, generated_screenshot_paths, [] # Ensure the main_preview_path is reset or handled if it should not change after analysis
 
             except MistralAPIException as e:
-                # Catch API key missing or client not installed errors from get_client or client method calls
-                return "error", f"**Mistral API Error:** {e.message}", current_main_preview_path, []
+                return "error", f"**Mistral API Error:** {e.message}", [], []
             except Exception as exc:
-                return "error", f"**Unexpected worker error:** {type(exc).__name__}: {exc}", current_main_preview_path, []
+                return "error", f"**Unexpected worker error:** {type(exc).__name__}: {exc}", [], []
 
         submit_btn.click(
             fn=worker,
-            inputs=[url_input, custom_prompt, api_key_input, main_preview_path_state, raw_media_path_state],
-            outputs=[status_state, output_markdown, main_preview_path_state, screenshot_paths_state],
+            inputs=[url_input, custom_prompt, api_key_input, raw_media_path_state],
+            outputs=[status_state, output_markdown, screenshot_paths_state, main_preview_path_state], # main_preview_path_state should remain unchanged or be updated from worker if needed.
             show_progress="full",
             show_progress_on=progress_markdown,
         )
@@ -973,7 +861,6 @@ def create_demo():
             gallery_update = gr.update(value=current_screenshot_paths, visible=bool(current_screenshot_paths))
             return img_update, video_update, gallery_update
 
-        # These change events use queue=False to ensure UI updates are immediate and don't block
         main_preview_path_state.change(
             fn=_update_preview_components,
             inputs=[main_preview_path_state, screenshot_paths_state],