Update app.py

app.py

@@ -5,7 +5,7 @@ import tempfile
 import base64
 import json
 from io import BytesIO
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Set # Import Set
 import requests
 from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
@@ -85,18 +85,20 @@ 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 ---
-_temp_files_to_delete = []
+# --- Temporary File Cleanup (Changed to use a set) ---
+_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."""
-    for f_path in list(_temp_files_to_delete):
+    # Create a copy to iterate while modifying the original set
+    for f_path in _temp_files_to_delete.copy():
         if os.path.exists(f_path):
             try:
                 os.remove(f_path)
+                _temp_files_to_delete.discard(f_path) # Remove from set after deletion
             except Exception as e:
                 print(f"Error during final cleanup of {f_path}: {e}")
-    _temp_files_to_delete.clear()
+    _temp_files_to_delete.clear() # Ensure the set is empty
 
 atexit.register(_cleanup_all_temp_files)
 
@@ -156,7 +158,7 @@ def _temp_file(data: bytes, suffix: str) -> str:
     os.close(fd)
     with open(path, "wb") as f:
         f.write(data)
-    _temp_files_to_delete.append(path)
+    _temp_files_to_delete.add(path) # Add to set
     return path
 
 def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD_BYTES, timeout: int = 60, progress=None) -> bytes:
@@ -176,10 +178,15 @@ def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD_BYTES, timeou
                     try:
                         with requests.get(src, timeout=timeout, stream=True, headers=DEFAULT_HEADERS) as r:
                             r.raise_for_status()
-                            with open(p, "wb") as fh:
-                                for chunk in r.iter_content(8192):
-                                    if chunk:
-                                        fh.write(chunk)
+                            total_size = int(r.headers.get("content-length", 0))
+                            downloaded_size = 0
+                            for chunk in r.iter_content(8192):
+                                if chunk:
+                                    fh.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:
                             return fh.read()
                     finally:
@@ -236,11 +243,13 @@ def _ffprobe_streams(path: str) -> Optional[dict]:
     """Uses ffprobe to get stream information for a media file."""
     if not FFMPEG_BIN:
         return None
-    ffprobe_path = FFMPEG_BIN.replace("ffmpeg", "ffprobe") if "ffmpeg" in FFMPEG_BIN else "ffprobe"
+    # Use FFMPEG_BIN for ffprobe if it's explicitly provided and ends with ffmpeg
+    ffprobe_path = FFMPEG_BIN.replace("ffmpeg", "ffprobe") if "ffmpeg" in os.path.basename(FFMPEG_BIN) else "ffprobe"
     if not shutil.which(ffprobe_path):
-        ffprobe_path = "ffprobe" # Try raw 'ffprobe' if 'ffmpeg' replacement isn't found
+        # Fallback to just 'ffprobe' if the path replacement didn't find it
+        ffprobe_path = "ffprobe"
         if not shutil.which(ffprobe_path):
-            return None
+            return None # ffprobe is not available
 
     cmd = [
         ffprobe_path, "-v", "error", "-print_format", "json", "-show_streams", "-show_format", path
@@ -248,7 +257,8 @@ def _ffprobe_streams(path: str) -> Optional[dict]:
     try:
         out = subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
         return json.loads(out)
-    except Exception:
+    except Exception as e:
+        print(f"Error running ffprobe on {path}: {e}")
         return None
 
 def _get_video_info_and_timestamps(media_path: str, sample_count: int) -> Tuple[Optional[dict], List[float]]:
@@ -269,7 +279,7 @@ def _get_video_info_and_timestamps(media_path: str, sample_count: int) -> Tuple[
             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]
+        timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count] # Try fixed early timestamps
 
     return info, timestamps
 
@@ -281,9 +291,13 @@ 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 or not os.path.exists(media_path):
-        print(f"Warning: FFMPEG not found or media path does not exist: {media_path}. Cannot extract frames.")
+    if not FFMPEG_BIN: # Check FFMPEG_BIN here
+        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):
+        print(f"Warning: Media path does not exist: {media_path}. Cannot extract frames.")
+        return frames_for_model, frame_paths_for_gallery
+
 
     if progress is not None:
         progress(0.05, desc="Preparing frame extraction...")
@@ -473,6 +487,10 @@ def analyze_video_cohesive(client, video_path: str, prompt: str, progress=None)
     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.", []
+
     try:
         if progress is not None:
             progress(0.3, desc="Uploading video for full analysis...")
@@ -558,8 +576,7 @@ def _convert_video_for_preview_if_needed(path: str) -> str:
     except Exception as e:
         print(f"Error converting video for preview: {e}")
         # If conversion fails, remove the failed temp file and return original path
-        if out_path in _temp_files_to_delete:
-            _temp_files_to_delete.remove(out_path)
+        _temp_files_to_delete.discard(out_path) # Remove from set
         try: os.remove(out_path)
         except Exception: pass
         return path
@@ -613,7 +630,23 @@ 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 features will fail. Run `pip install mistralai`."
+    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
+
         with gr.Row():
             with gr.Column(scale=1):
                 preview_image = gr.Image(label="Preview Image", type="filepath", elem_classes="preview_media", visible=False)
@@ -641,13 +674,15 @@ def create_demo():
             """
             Cleans up all tracked temporary files and resets all relevant UI components and states.
             """
-            for f_path in list(_temp_files_to_delete):
+            # Iterate a copy of the set while potentially modifying the original
+            for f_path in _temp_files_to_delete.copy():
                 if os.path.exists(f_path):
                     try:
                         os.remove(f_path)
+                        _temp_files_to_delete.discard(f_path) # Remove from set
                     except Exception as e:
                         print(f"Error during proactive cleanup of {f_path}: {e}")
-            _temp_files_to_delete.clear()
+            _temp_files_to_delete.clear() # Ensure set is empty
 
             return "", \
                    gr.update(value=None, visible=False), \
@@ -692,20 +727,19 @@ def create_demo():
             Also handles cleanup of previously loaded media.
             """
             # --- Proactive cleanup of old files related to previous load ---
-            files_to_clean_up_now = []
             if current_main_preview_path and os.path.exists(current_main_preview_path):
-                files_to_clean_up_now.append(current_main_preview_path)
+                _temp_files_to_delete.discard(current_main_preview_path) # Remove from set
+                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):
-                files_to_clean_up_now.append(current_raw_media_path)
+                _temp_files_to_delete.discard(current_raw_media_path) # Remove from set
+                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):
-                    files_to_clean_up_now.append(path)
-
-            for f_path in files_to_clean_up_now:
-                if f_path in _temp_files_to_delete:
-                    _temp_files_to_delete.remove(f_path)
-                try: os.remove(f_path)
-                except Exception as e: print(f"Error cleaning up old temp file {f_path}: {e}")
+                    _temp_files_to_delete.discard(path) # Remove from set
+                    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)
@@ -723,7 +757,7 @@ 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)
+                raw_bytes_for_analysis = fetch_bytes(url, timeout=60, progress=progress) # Pass progress here
                 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), \
@@ -736,14 +770,13 @@ def create_demo():
                            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.05, desc="Generating playable preview...")
+                progress(0.25, desc="Generating playable preview...") # Adjusted progress start
                 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
-                    if temp_raw_path_for_analysis in _temp_files_to_delete:
-                        _temp_files_to_delete.remove(temp_raw_path_for_analysis)
+                    _temp_files_to_delete.discard(temp_raw_path_for_analysis) # Remove from set
                     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}")
 
@@ -765,13 +798,11 @@ def create_demo():
                            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
-                    if local_playable_path in _temp_files_to_delete:
-                        _temp_files_to_delete.remove(local_playable_path)
+                    _temp_files_to_delete.discard(local_playable_path) # Remove from set
                     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
-                    if temp_raw_path_for_analysis in _temp_files_to_delete:
-                        _temp_files_to_delete.remove(temp_raw_path_for_analysis)
+                    _temp_files_to_delete.discard(temp_raw_path_for_analysis) # Remove from set
                     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}")
 
@@ -782,8 +813,7 @@ def create_demo():
             except Exception as e:
                 # If an error occurred during loading, clear all relevant paths.
                 if os.path.exists(temp_raw_path_for_analysis):
-                    if temp_raw_path_for_analysis in _temp_files_to_delete:
-                        _temp_files_to_delete.remove(temp_raw_path_for_analysis)
+                    _temp_files_to_delete.discard(temp_raw_path_for_analysis) # Remove from set
                     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}")
 
@@ -808,6 +838,14 @@ def create_demo():
                 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, []
 
+                # 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.
+                    is_image_check, is_video_check = determine_media_type(url)
+                    if is_video_check:
+                         return "error", "**Error:** FFmpeg is not found in your system PATH. Video analysis is unavailable. Please install FFmpeg.", current_main_preview_path, []
+
                 with open(raw_media_path, "rb") as f:
                     raw_bytes_for_analysis = f.read()