Update app.py

app.py

@@ -1,131 +1,336 @@
 #!/usr/bin/env python3
-import os
-from io import BytesIO
-from typing import Tuple
+# -*- coding: utf-8 -*-
 
+from __future__ import annotations
+import os, shutil, subprocess, tempfile, base64, json
+from io import BytesIO
+from typing import List, Tuple
+import requests
+from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
-from PIL import Image
 
-css = """
-.preview_media img, .preview_media video { max-width: 100%; height: auto; border-radius: 6px; }
-.top_preview { display:flex; gap:12px; align-items:flex-start; flex-direction:column; }
-.pip_button { margin-top:6px; }
-"""
+# --- Config
+DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
+PIXTRAL_MODEL = "pixtral-12b-2409"
+VIDEO_MODEL = "voxtral-mini-latest"
+STREAM_THRESHOLD = 20 * 1024 * 1024
+FFMPEG_BIN = shutil.which("ffmpeg")
+IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
+VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".flv")
 
-def _btn_label_for_status(status: str) -> str:
-    return {
-        "idle": "Submit",
-        "busy": "Processing…",
-        "done": "Submit",
-        "error": "Retry",
-    }.get(status or "idle", "Submit")
-
-
-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, gr.update(visible=False)
-
-    # Local file
-    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), gr.update(visible=True)
-        if ext in IMAGE_EXTS:
+SYSTEM_INSTRUCTION = (
+    "You are a clinical visual analyst. Only analyze media actually provided (image or video data). If video review entirely, not frame by frame. "
+    "Provide factual, clinical descriptions of observable features. Avoid inventing sensory information."
+)
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+Image.MAX_IMAGE_PIXELS = 10000 * 10000
+
+try:
+    from mistralai import Mistral
+except Exception:
+    Mistral = None
+
+def get_client(key: str | None = None):
+    api_key = (key or "").strip() or DEFAULT_KEY
+    if Mistral is None:
+        class Dummy:
+            def __init__(self, k): self.api_key = k
+        return Dummy(api_key)
+    return Mistral(api_key=api_key)
+
+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 ""
+    _, 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)
+        return None if r.status_code >= 400 else 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:
+    if is_remote(src):
+        head = safe_head(src)
+        if head is not None:
+            cl = head.headers.get("content-length")
             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, gr.update(visible=False)
+                if cl and int(cl) > stream_threshold:
+                    with requests.get(src, timeout=timeout, stream=True) as r:
+                        r.raise_for_status()
+                        fd, p = tempfile.mkstemp()
+                        os.close(fd)
+                        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()
+                        finally:
+                            try: os.remove(p)
+                            except Exception: pass
             except Exception:
-                return empty_img, empty_vid, gr.update(visible=False)
+                pass
+        r = safe_get(src, timeout=timeout)
+        return r.content
+    else:
+        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)
+    return path
+
+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)
+    except Exception: pass
+    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)
+    buf = BytesIO(); 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: 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)
+            if os.path.exists(tmp) and os.path.getsize(tmp) > 0:
+                with open(tmp, "rb") as f: frames.append(f.read())
+        except Exception:
+            pass
+        finally:
+            try: os.remove(tmp)
+            except Exception: pass
+    return frames
+
+def chat_complete(client, model: str, messages, timeout: int = 120) -> str:
+    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)
+            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)
+        first = choices[0]
+        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)
+        return content.strip() if isinstance(content, str) else str(content)
+    except Exception as e:
+        return f"Error during model call: {e}"
 
-    # Remote: check headers for content-type
-    head = safe_head(url)
-    if head:
-        ctype = (head.headers.get("content-type") or "").lower()
-        if ctype.startswith("video/") or any(url.lower().endswith(ext) for ext in VIDEO_EXTS):
-            return empty_img, gr.update(value=url, visible=True), gr.update(visible=True)
+def upload_file_to_mistral(client, path: str, filename: str | None = None, purpose: str = "batch", timeout: int = 120) -> str:
+    fname = filename or os.path.basename(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)
+            fid = getattr(res, "id", None) or (res.get("id") if isinstance(res, dict) else None)
+            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=timeout); r.raise_for_status(); jr = r.json()
+        return jr.get("id") or jr.get("data", [{}])[0].get("id")
 
-    # Try load as image
+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:
-        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, gr.update(visible=False)
+        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\nInstruction: 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:
-        return empty_img, empty_vid, gr.update(visible=False)
+        frames = extract_best_frames_bytes(video_path, sample_count=6)
+        if not frames: return "Error: could not upload video and no frames could be extracted."
+        image_entries = []
+        for i, fb in enumerate(frames, start=1):
+            try:
+                j = convert_to_jpeg_bytes(fb, base_h=720)
+                image_entries.append({"type": "image_url", "image_url": b64_bytes(j, mime="image/jpeg"), "meta": {"frame_index": i}})
+            except Exception:
+                continue
+        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}, {"role": "user", "content": content}]
+        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="Analyzing image")
+        try:
+            return analyze_image_structured(client, raw, prompt)
+        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_path = save_bytes_to_temp(raw, suffix=ext_from_src(src) or ".mp4")
+        try:
+            progress(0.2, desc="Analyzing video")
+            return analyze_video_cohesive(client, tmp_path, prompt)
+        finally:
+            try: os.remove(tmp_path)
+            except Exception: pass
+    try:
+        raw = fetch_bytes(src)
+        progress(0.2, desc="Treating as image")
+        return analyze_image_structured(client, raw, prompt)
+    except Exception as e:
+        return f"Unable to determine media type or fetch file: {e}"
+
+# --- Gradio UI
+css = ".preview_media img, .preview_media video { max-width: 100%; height: auto; border-radius:6px; }"
+
+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(css=css, title="Media Analysis") as demo:
+    with gr.Blocks(title="Flux Multimodal", css=css) as demo:
         with gr.Row():
             with gr.Column(scale=1):
-                preview_image = gr.Image(label="Preview Image", type="pil", visible=False, elem_classes="preview_media")
-                preview_video = gr.Video(label="Preview Video", visible=False, elem_classes="preview_media")
-                pip_button = gr.Button("Open Video in PiP", visible=False, elem_classes="pip_button")
-            with gr.Column(scale=1):
+                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)
+                pip_button = gr.Button("Open Video in PiP", visible=False)
+            with gr.Column(scale=2):
                 url_input = gr.Textbox(label="Image / Video URL or local path", placeholder="https://... or /path/to/file", lines=1)
                 with gr.Accordion("Prompt (optional)", open=False):
                     custom_prompt = gr.Textbox(label="Prompt", lines=4, 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(_btn_label_for_status("idle"), variant="primary")
+                submit_btn = gr.Button(_btn_label_for_status("idle"))
                 clear_btn = gr.Button("Clear")
-                final_md = gr.Markdown(value="", label="Result")
+                output_md = gr.Markdown("")
+                status_state = gr.State("idle")
+
+        pip_html = gr.HTML("""<div id="pip-root" style="display:none"></div>
+<script>
+window.openPiP = (sel) => {
+  try {
+    const v = document.querySelector(sel);
+    if (!v) return "no-video";
+    if (v.requestPictureInPicture) { v.requestPictureInPicture(); return "opened"; }
+    return "unsupported";
+  } catch(e){ return "error:"+e; }
+};
+</script>""")
 
-        status = gr.State("idle")
+        def load_preview(url: str):
+            empty_img = gr.update(value=None, visible=False)
+            empty_vid = gr.update(value=None, visible=False)
+            pip_vis = gr.update(visible=False)
+            if not url: return empty_img, empty_vid, pip_vis
+            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), gr.update(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, pip_vis
+                    except Exception: return empty_img, empty_vid, pip_vis
+            head = safe_head(url)
+            if head:
+                ctype = (head.headers.get("content-type") or "").lower()
+                if ctype.startswith("video/") or any(url.lower().endswith(ext) for ext in VIDEO_EXTS):
+                    return empty_img, gr.update(value=url, visible=True), gr.update(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, pip_vis
+            except Exception:
+                return empty_img, empty_vid, pip_vis
 
         url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video, pip_button])
 
-        def do_clear():
-            return "", gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value="", visible=True)
-        clear_btn.click(fn=do_clear, inputs=[], outputs=[url_input, preview_image, preview_video, final_md])
+        def clear_all():
+            return "", gr.update(value=None, visible=False), gr.update(value=None, visible=False), "idle", gr.update(value=_btn_label_for_status("idle"))
+        clear_btn.click(fn=clear_all, inputs=[], outputs=[url_input, preview_image, preview_video, status_state, submit_btn])
+
+        def pip_click(_):
+            js = "<script>setTimeout(()=>window.openPiP('video.preview_media'),50);</script>"
+            return gr.HTML.update(value=js)
+        pip_button.click(fn=pip_click, inputs=[url_input], outputs=[pip_html])
 
         def start_busy():
-            return "busy", gr.update(value=_btn_label_for_status("busy"), interactive=False)
+            s = "busy"
+            return s, gr.update(value=_btn_label_for_status(s))
+        submit_btn.click(fn=start_busy, inputs=[], outputs=[status_state, submit_btn])
 
         def worker(url: str, prompt: str, key: str, progress=gr.Progress()):
             return process_media(url or "", prompt or "", key or "", progress=progress)
-
-        def finish(result: str):
-            s = "done"
-            if not result:
-                md = "**Error:** no result returned."
-                s = "error"
-            elif isinstance(result, str) and result.lower().startswith("error"):
-                md = f"**Error:** {result}"
-                s = "error"
-            else:
-                md = result
-            return s, gr.update(value=md)
-
-        submit_btn.click(fn=start_busy, inputs=[], outputs=[status, submit_btn])
-
-        submit_btn.click(fn=worker, inputs=[url_input, custom_prompt, api_key], outputs=[final_md], queue=True).then(
-            fn=finish, inputs=[final_md], outputs=[status, final_md]
+        submit_btn.click(fn=worker, inputs=[url_input, custom_prompt, api_key], outputs=[output_md], queue=True).then(
+            fn=lambda res: ("error", "**Error:** no result returned.") if not res else
+                           ("error", f"**Error:** {res}") if isinstance(res, str) and res.lower().startswith("error") else ("done", res),
+            inputs=[output_md],
+            outputs=[status_state, output_md],
         )
 
         def btn_label_for_state(s: str):
-            return gr.update(value=_btn_label_for_status(s), interactive=(s != "busy"))
-        status.change(fn=btn_label_for_state, inputs=[status], outputs=[submit_btn])
-
-        # PiP JS: use a JS function string without Python .format interpolation
-        pip_js = (
-            "const vid = document.querySelector('video');"
-            "if (!vid) return;"
-            "if (document.pictureInPictureElement) { document.exitPictureInPicture().catch(()=>{}); }"
-            "else { vid.requestPictureInPicture().catch(()=>{}); }"
-        )
-        # Attach PiP: client-side JS (no inputs/outputs)
-        pip_button.click(fn=None, _js=pip_js, inputs=[], outputs=[])
+            return _btn_label_for_status(s)
+        status_state.change(fn=btn_label_for_state, inputs=[status_state], outputs=[submit_btn])
 
     return demo
 
-
 if __name__ == "__main__":
-    demo = create_demo()
-    demo.launch(server_name="0.0.0.0", share=False)
+    create_demo().launch()