Update app.py

app.py

@@ -1,21 +1,36 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
 from __future__ import annotations
-import os, shutil, subprocess, tempfile, base64, requests, time
+
+import os
+import shutil
+import subprocess
+import tempfile
+import base64
+import json
 from io import BytesIO
 from typing import List, Tuple
+
+import requests
 from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
 
+# ----------------------------------------------------------------------
+# Optional: Mistral SDK – fall back to raw HTTP if not installed
+# ----------------------------------------------------------------------
 try:
     from mistralai import Mistral
-except Exception:
+except Exception:  # pragma: no cover
     Mistral = None
 
-# Config
+# ----------------------------------------------------------------------
+# Configuration constants
+# ----------------------------------------------------------------------
 DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
 PIXTRAL_MODEL = "pixtral-12b-2409"
 VIDEO_MODEL = "voxtral-mini-latest"
-STREAM_THRESHOLD = 20 * 1024 * 1024
+STREAM_THRESHOLD = 20 * 1024 * 1024          # 20 MiB
 FFMPEG_BIN = shutil.which("ffmpeg")
 IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
 VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".flv")
@@ -27,9 +42,15 @@ SYSTEM_INSTRUCTION = (
     "Produce a concise, narrative that blends scientific precision with erotic detail. Do not invent sensory information not present in the media."
 )
 
+# ----------------------------------------------------------------------
+# Pillow safety tweaks
+# ----------------------------------------------------------------------
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 Image.MAX_IMAGE_PIXELS = 10000 * 10000
 
+# ----------------------------------------------------------------------
+# Helper: Mistral client
+# ----------------------------------------------------------------------
 def get_client(key: str | None = None):
     api_key = (key or "").strip() or DEFAULT_KEY
     if Mistral is None:
@@ -38,28 +59,39 @@ def get_client(key: str | None = None):
         return Dummy(api_key)
     return Mistral(api_key=api_key)
 
+
+# ----------------------------------------------------------------------
+# URL / file utilities
+# ----------------------------------------------------------------------
 def is_remote(src: str) -> bool:
     return bool(src) and src.startswith(("http://", "https://"))
 
+
 def ext_from_src(src: str) -> str:
-    if not src: return ""
+    if not src:
+        return ""
     _, ext = os.path.splitext((src or "").split("?")[0])
     return ext.lower()
 
+
 def safe_head(url: str, timeout: int = 6):
     try:
         r = requests.head(url, timeout=timeout, allow_redirects=True)
-        if r.status_code >= 400: return None
+        if r.status_code >= 400:
+            return None
         return r
     except Exception:
         return None
 
+
 def safe_get(url: str, timeout: int = 15):
     r = requests.get(url, timeout=timeout)
     r.raise_for_status()
     return r
 
+
 def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD, timeout: int = 60) -> bytes:
+    """Download remote files or read local ones, streaming large objects."""
     if is_remote(src):
         head = safe_head(src)
         if head is not None:
@@ -73,31 +105,44 @@ def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD, timeout: int
                         try:
                             with open(p, "wb") as fh:
                                 for chunk in r.iter_content(8192):
-                                    if chunk: fh.write(chunk)
-                            with open(p, "rb") as fh: return fh.read()
+                                    if chunk:
+                                        fh.write(chunk)
+                            with open(p, "rb") as fh:
+                                return fh.read()
                         finally:
-                            try: os.remove(p)
-                            except Exception: pass
+                            try:
+                                os.remove(p)
+                            except Exception:
+                                pass
             except Exception:
                 pass
         r = safe_get(src, timeout=timeout)
         return r.content
     else:
-        with open(src, "rb") as f: return f.read()
+        with open(src, "rb") as f:
+            return f.read()
+
 
 def save_bytes_to_temp(b: bytes, suffix: str) -> str:
     fd, path = tempfile.mkstemp(suffix=suffix)
     os.close(fd)
-    with open(path, "wb") as f: f.write(b)
+    with open(path, "wb") as f:
+        f.write(b)
     return path
 
+
+# ----------------------------------------------------------------------
+# Image preprocessing
+# ----------------------------------------------------------------------
 def convert_to_jpeg_bytes(img_bytes: bytes, base_h: int = 480) -> bytes:
     img = Image.open(BytesIO(img_bytes))
     try:
-        if getattr(img, "is_animated", False): img.seek(0)
+        if getattr(img, "is_animated", False):
+            img.seek(0)
     except Exception:
         pass
-    if img.mode != "RGB": img = img.convert("RGB")
+    if img.mode != "RGB":
+        img = img.convert("RGB")
     h = base_h
     w = max(1, int(img.width * (h / img.height)))
     img = img.resize((w, h), Image.LANCZOS)
@@ -105,246 +150,115 @@ def convert_to_jpeg_bytes(img_bytes: bytes, base_h: int = 480) -> bytes:
     img.save(buf, format="JPEG", quality=85)
     return buf.getvalue()
 
+
 def b64_bytes(b: bytes, mime: str = "image/jpeg") -> str:
     return f"data:{mime};base64," + base64.b64encode(b).decode("utf-8")
 
-def extract_best_frames_bytes(media_path: str, sample_count: int = 5, timeout_extract: int = 15) -> List[bytes]:
-    frames = []
-    if not FFMPEG_BIN or not os.path.exists(media_path): return frames
-    duration_pts = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count]
-    for i, t in enumerate(duration_pts):
-        fd, tmp = tempfile.mkstemp(suffix=f"_{i}.jpg"); os.close(fd)
-        cmd = [FFMPEG_BIN, "-nostdin", "-y", "-ss", str(t), "-i", media_path, "-frames:v", "1", "-q:v", "2", tmp]
+
+# ----------------------------------------------------------------------
+# Video frame extraction (fallback)
+# ----------------------------------------------------------------------
+def extract_best_frames_bytes(
+    media_path: str, sample_count: int = 5, timeout_extract: int = 15
+) -> List[bytes]:
+    frames: List[bytes] = []
+    if not FFMPEG_BIN or not os.path.exists(media_path):
+        return frames
+    timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count]
+    for i, t in enumerate(timestamps):
+        fd, tmp = tempfile.mkstemp(suffix=f"_{i}.jpg")
+        os.close(fd)
+        cmd = [
+            FFMPEG_BIN,
+            "-nostdin",
+            "-y",
+            "-ss",
+            str(t),
+            "-i",
+            media_path,
+            "-frames:v",
+            "1",
+            "-q:v",
+            "2",
+            tmp,
+        ]
         try:
-            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=timeout_extract)
+            subprocess.run(
+                cmd,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                timeout=timeout_extract,
+            )
             if os.path.exists(tmp) and os.path.getsize(tmp) > 0:
-                with open(tmp, "rb") as f: frames.append(f.read())
+                with open(tmp, "rb") as f:
+                    frames.append(f.read())
         except Exception:
             pass
         finally:
-            try: os.remove(tmp)
-            except Exception: pass
+            try:
+                os.remove(tmp)
+            except Exception:
+                pass
     return frames
 
+
+# ----------------------------------------------------------------------
+# Model interaction helpers
+# ----------------------------------------------------------------------
 def chat_complete(client, model: str, messages, timeout: int = 120) -> str:
+    """Wrap SDK and raw‑HTTP calls to Mistral chat completions."""
     try:
         if hasattr(client, "chat") and hasattr(client.chat, "complete"):
             res = client.chat.complete(model=model, messages=messages, stream=False)
         else:
             api_key = getattr(client, "api_key", "") or DEFAULT_KEY
             url = "https://api.mistral.ai/v1/chat/completions"
-            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} if api_key else {"Content-Type":"application/json"}
-            r = requests.post(url, json={"model": model, "messages": messages}, headers=headers, timeout=timeout)
+            headers = (
+                {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+                if api_key
+                else {"Content-Type": "application/json"}
+            )
+            r = requests.post(
+                url,
+                json={"model": model, "messages": messages},
+                headers=headers,
+                timeout=timeout,
+            )
             r.raise_for_status()
             res = r.json()
         choices = getattr(res, "choices", None) or (res.get("choices") if isinstance(res, dict) else [])
-        if not choices: return str(res)
+        if not choices:
+            return str(res)
         first = choices[0]
-        msg = first.message if hasattr(first, "message") else (first.get("message") if isinstance(first, dict) else first)
+        msg = (
+            first.message
+            if hasattr(first, "message")
+            else (first.get("message") if isinstance(first, dict) else first)
+        )
         content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
-        if isinstance(content, str): return content.strip()
-        return str(content)
+        return content.strip() if isinstance(content, str) else str(content)
     except Exception as e:
         return f"Error during model call: {e}"
 
-def upload_file_to_mistral(client, path: str, filename: str | None = None, purpose: str = "batch") -> str:
+
+def upload_file_to_mistral(
+    client,
+    path: str,
+    filename: str | None = None,
+    purpose: str = "batch",
+    timeout: int = 120,
+) -> str:
+    """Upload a file to Mistral and return its file‑id."""
     fname = filename or os.path.basename(path)
+    # SDK path -------------------------------------------------
     try:
         if hasattr(client, "files") and hasattr(client.files, "upload"):
             with open(path, "rb") as fh:
-                res = client.files.upload(file={"file_name": fname, "content": fh}, purpose=purpose)
+                res = client.files.upload(
+                    file={"file_name": fname, "content": fh}, purpose=purpose
+                )
             fid = getattr(res, "id", None) or (res.get("id") if isinstance(res, dict) else None)
-            if not fid: fid = res["data"][0]["id"]
+            if not fid:
+                fid = res["data"][0]["id"]
             return fid
     except Exception:
-        pass
-    api_key = getattr(client, "api_key", "") or DEFAULT_KEY
-    url = "https://api.mistral.ai/v1/files"
-    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
-    with open(path, "rb") as fh:
-        files = {"file": (fname, fh)}
-        data = {"purpose": purpose}
-        r = requests.post(url, headers=headers, files=files, data=data, timeout=120)
-        r.raise_for_status()
-        jr = r.json()
-        return jr.get("id") or jr.get("data", [{}])[0].get("id")
-
-def analyze_image_structured(client, img_bytes: bytes, prompt: str) -> str:
-    jpeg = convert_to_jpeg_bytes(img_bytes, base_h=1024)
-    data_url = b64_bytes(jpeg, mime="image/jpeg")
-    messages = [
-        {"role": "system", "content": SYSTEM_INSTRUCTION},
-        {"role": "user", "content": [
-            {"type": "text", "text": prompt},
-            {"type": "image_url", "image_url": data_url}
-        ]}
-    ]
-    return chat_complete(client, PIXTRAL_MODEL, messages)
-
-def analyze_video_cohesive(client, video_path: str, prompt: str) -> str:
-    try:
-        file_id = upload_file_to_mistral(client, video_path, filename=os.path.basename(video_path))
-        extra_msg = (
-            f"Uploaded video file id: {file_id}\n\n"
-            "Instruction: Analyze the entire video and produce a single cohesive narrative describing consistent observations."
-        )
-        messages = [
-            {"role": "system", "content": SYSTEM_INSTRUCTION},
-            {"role": "user", "content": extra_msg + "\n\n" + prompt}
-        ]
-        return chat_complete(client, VIDEO_MODEL, messages)
-    except Exception:
-        frames = extract_best_frames_bytes(video_path, sample_count=6)
-        if not frames:
-            return "Error: could not upload video and no frames extracted (ffmpeg missing or failed)."
-        images_entries = []
-        for i, fb in enumerate(frames, start=1):
-            try:
-                j = convert_to_jpeg_bytes(fb, base_h=720)
-                images_entries.append({"type": "image_url", "image_url": b64_bytes(j, mime="image/jpeg"), "meta": {"frame_index": i}})
-            except Exception:
-                continue
-        content_list = [{"type":"text", "text": prompt + "\n\nPlease consolidate observations across these frames into a single cohesive narrative."}] + images_entries
-        messages = [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": content_list}]
-        return chat_complete(client, PIXTRAL_MODEL, messages)
-
-def determine_media_type(src: str) -> Tuple[bool, bool]:
-    is_image = False; is_video = False
-    ext = ext_from_src(src)
-    if ext in IMAGE_EXTS: is_image = True
-    if ext in VIDEO_EXTS: is_video = True
-    if is_remote(src):
-        head = safe_head(src)
-        if head:
-            ctype = (head.headers.get("content-type") or "").lower()
-            if ctype.startswith("image/"): is_image, is_video = True, False
-            elif ctype.startswith("video/"): is_video, is_image = True, False
-    return is_image, is_video
-
-def process_media(src: str, custom_prompt: str, api_key: str, progress=gr.Progress()) -> str:
-    client = get_client(api_key)
-    prompt = (custom_prompt or "").strip() or "Please provide a detailed visual review."
-    if not src: return "No URL or path provided."
-    progress(0.05, desc="Determining media type")
-    is_image, is_video = determine_media_type(src)
-    if is_image:
-        try:
-            raw = fetch_bytes(src)
-        except Exception as e:
-            return f"Error fetching image: {e}"
-        progress(0.2, desc="Preparing image")
-        try:
-            out = analyze_image_structured(client, raw, prompt)
-            progress(1.0, desc="Done")
-            return out
-        except UnidentifiedImageError:
-            return "Error: provided file is not a valid image."
-        except Exception as e:
-            return f"Error analyzing image: {e}"
-    if is_video:
-        try:
-            raw = fetch_bytes(src, timeout=120)
-        except Exception as e:
-            return f"Error fetching video: {e}"
-        tmp = save_bytes_to_temp(raw, suffix=ext_from_src(src) or ".mp4")
-        try:
-            progress(0.2, desc="Uploading video / extracting frames")
-            out = analyze_video_cohesive(client, tmp, prompt)
-            progress(1.0, desc="Done")
-            return out
-        finally:
-            try: os.remove(tmp)
-            except Exception: pass
-    try:
-        raw = fetch_bytes(src)
-        progress(0.2, desc="Treating as image")
-        out = analyze_image_structured(client, raw, prompt)
-        progress(1.0, desc="Done")
-        return out
-    except Exception as e:
-        return f"Unable to determine media type or fetch file: {e}"
-
-css = ".preview_media img, .preview_media video { max-width: 100%; height: auto; }"
-
-def load_preview(url: str):
-    empty_img = gr.update(value=None, visible=False)
-    empty_vid = gr.update(value=None, visible=False)
-    if not url:
-        return empty_img, empty_vid
-    if not is_remote(url) and os.path.exists(url):
-        ext = ext_from_src(url)
-        if ext in VIDEO_EXTS:
-            return empty_img, gr.update(value=os.path.abspath(url), visible=True)
-        if ext in IMAGE_EXTS:
-            try:
-                img = Image.open(url)
-                if getattr(img, "is_animated", False): img.seek(0)
-                return gr.update(value=img.convert("RGB"), visible=True), empty_vid
-            except Exception:
-                return empty_img, empty_vid
-    head = safe_head(url)
-    if head:
-        ctype = (head.headers.get("content-type") or "").lower()
-        if ctype.startswith("video/") or any(url.lower().split("?")[0].endswith(ext) for ext in VIDEO_EXTS):
-            return empty_img, gr.update(value=url, visible=True)
-    try:
-        r = safe_get(url, timeout=15)
-        img = Image.open(BytesIO(r.content))
-        if getattr(img, "is_animated", False): img.seek(0)
-        return gr.update(value=img.convert("RGB"), visible=True), empty_vid
-    except Exception:
-        return empty_img, empty_vid
-
-def _btn_label_for_status(status: str) -> str:
-    return {"idle": "Submit", "busy": "Processing…", "done": "Submit", "error": "Retry"}.get(status or "idle", "Submit")
-
-def create_demo():
-    with gr.Blocks(title="Flux Multimodal (Pixtral fixed)", css=css) as demo:
-        with gr.Row():
-            with gr.Column(scale=1):
-                url_input = gr.Textbox(label="Image / Video URL or local path", placeholder="https://... or /path/to/file", lines=1)
-                custom_prompt = gr.Textbox(label="Prompt (optional)", lines=2, value="")
-                with gr.Accordion("Mistral API Key (optional)", open=False):
-                    api_key = gr.Textbox(label="API Key", type="password", max_lines=1)
-                submit_btn = gr.Button("Submit")
-                clear_btn = gr.Button("Clear")
-                preview_image = gr.Image(label="Preview Image", type="pil", elem_classes="preview_media", visible=False)
-                preview_video = gr.Video(label="Preview Video", elem_classes="preview_media", visible=False)
-            with gr.Column(scale=2):
-                final_md = gr.Markdown(value="")
-
-        url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video])
-        clear_btn.click(lambda: (gr.update(value=""), gr.update(value=None, visible=False), gr.update(value=None, visible=False)), inputs=[], outputs=[url_input, preview_image, preview_video])
-
-        status = gr.State("idle")
-
-        def start_busy() -> str:
-            return "busy"
-
-        def worker(url: str, prompt: str, key: str, progress=gr.Progress()):
-            try:
-                return process_media(url or "", prompt or "", key or "", progress=progress)
-            except Exception as e:
-                return f"Unhandled error in worker: {e}"
-
-        def finish(status_text: str) -> tuple[str, str]:
-            if (not status_text) or status_text.lower().startswith("error") or status_text.lower().startswith("unhandled"):
-                return "error", status_text
-            return "done", status_text
-
-        def btn_props_from_status(s: str):
-            s = s or "idle"
-            interactive = False if s == "busy" else True
-            label = _btn_label_for_status(s)
-            return gr.update(interactive=interactive, value=label)
-
-        submit_btn.click(fn=start_busy, inputs=[], outputs=[status])
-        submit_btn.click(fn=worker, inputs=[url_input, custom_prompt, api_key], outputs=[final_md], queue=True)
-        submit_btn.click(fn=finish, inputs=[final_md], outputs=[status, final_md])
-        status.change(fn=btn_props_from_status, inputs=[status], outputs=[submit_btn])
-
-    return demo
-
-if __name__ == "__main__":
-    demo = create_demo()
-    demo.queue().launch()