Update app.py

app.py

@@ -6,7 +6,7 @@ import shutil
 from io import BytesIO
 import base64
 import requests
-from PIL import Image, UnidentifiedImageError
+from PIL import Image, UnidentifiedImageError, ImageFile
 import gradio as gr
 from mistralai import Mistral
 
@@ -14,6 +14,7 @@ DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
 DEFAULT_IMAGE_MODEL = "pixtral-12b-2409"
 DEFAULT_VIDEO_MODEL = "voxtral-mini-latest"
 STREAM_THRESHOLD = 20 * 1024 * 1024
+FFMPEG_BIN = shutil.which("ffmpeg")
 
 SYSTEM_INSTRUCTION = (
     "You are a clinical visual analyst. Only analyze media actually provided (image data or extracted frames). "
@@ -25,6 +26,9 @@ SYSTEM_INSTRUCTION = (
 IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
 VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".flv")
 
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+Image.MAX_IMAGE_PIXELS = 10000 * 10000
+
 
 def get_client(key: str = None):
     api_key = (key or "").strip() or DEFAULT_KEY
@@ -48,16 +52,18 @@ def fetch_bytes(src: str, stream_threshold=STREAM_THRESHOLD, timeout=60) -> byte
             if cl and int(cl) > stream_threshold:
                 fd, path = tempfile.mkstemp()
                 os.close(fd)
-                with open(path, "wb") as f:
-                    for chunk in r.iter_content(8192):
-                        if chunk:
-                            f.write(chunk)
-                with open(path, "rb") as f:
-                    data = f.read()
                 try:
-                    os.remove(path)
-                except Exception:
-                    pass
+                    with open(path, "wb") as f:
+                        for chunk in r.iter_content(8192):
+                            if chunk:
+                                f.write(chunk)
+                    with open(path, "rb") as f:
+                        data = f.read()
+                finally:
+                    try:
+                        os.remove(path)
+                    except Exception:
+                        pass
                 return data
             return r.content
     with open(src, "rb") as f:
@@ -93,46 +99,19 @@ def save_bytes_to_temp(b: bytes, suffix: str):
     return path
 
 
-def extract_delta(chunk):
-    if not chunk:
-        return None
-    data = getattr(chunk, "data", None) or getattr(chunk, "response", None) or getattr(chunk, "delta", None)
-    if not data:
-        return None
-    try:
-        content = data.choices[0].delta.content
-        if content is None:
-            return None
-        return str(content)
-    except Exception:
-        pass
-    try:
-        msg = data.choices[0].message
-        if isinstance(msg, dict):
-            content = msg.get("content")
-        else:
-            content = getattr(msg, "content", None)
-        if content is None:
-            return None
-        return str(content)
-    except Exception:
-        pass
-    try:
-        return str(data)
-    except Exception:
-        return None
-
-
-def extract_best_frame_bytes(media_path: str, sample_count: int = 5, timeout_probe: int = 10, timeout_extract: int = 15):
-    ffmpeg = shutil.which("ffmpeg")
-    if not ffmpeg or not os.path.exists(media_path):
-        return None
+def extract_best_frames_bytes(media_path: str, sample_count: int = 5, timeout_probe: int = 10, timeout_extract: int = 15):
+    if not FFMPEG_BIN or not os.path.exists(media_path):
+        return []
     tmp_frames = []
     try:
-        probe_cmd = [ffmpeg, "-v", "error", "-show_entries", "format=duration",
+        probe_cmd = [FFMPEG_BIN, "-v", "error", "-show_entries", "format=duration",
                      "-of", "default=noprint_wrappers=1:nokey=1", media_path]
         proc = subprocess.Popen(probe_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        out, _ = proc.communicate(timeout=timeout_probe)
+        try:
+            out, _ = proc.communicate(timeout=timeout_probe)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            out, _ = proc.communicate()
         duration = None
         try:
             duration = float(out.strip().split(b"\n")[0]) if out else None
@@ -148,7 +127,7 @@ def extract_best_frame_bytes(media_path: str, sample_count: int = 5, timeout_pro
             fd, tmp_frame = tempfile.mkstemp(suffix=f"_{i}.jpg")
             os.close(fd)
             cmd = [
-                ffmpeg, "-nostdin", "-y", "-i", media_path,
+                FFMPEG_BIN, "-nostdin", "-y", "-i", media_path,
                 "-ss", str(t),
                 "-frames:v", "1",
                 "-q:v", "2",
@@ -164,29 +143,18 @@ def extract_best_frame_bytes(media_path: str, sample_count: int = 5, timeout_pro
                     pass
                 proc.communicate()
             if proc.returncode == 0 and os.path.exists(tmp_frame) and os.path.getsize(tmp_frame) > 0:
-                tmp_frames.append(tmp_frame)
-            else:
-                try:
-                    if os.path.exists(tmp_frame):
-                        os.remove(tmp_frame)
-                except Exception:
-                    pass
-
-        if not tmp_frames:
-            return None
-
-        chosen = max(tmp_frames, key=lambda p: os.path.getsize(p) if os.path.exists(p) else 0)
-        with open(chosen, "rb") as f:
-            data = f.read()
-        return data
-    finally:
-        for fpath in tmp_frames:
+                with open(tmp_frame, "rb") as f:
+                    tmp_frames.append(f.read())
             try:
-                if os.path.exists(fpath):
-                    os.remove(fpath)
+                if os.path.exists(tmp_frame):
+                    os.remove(tmp_frame)
             except Exception:
                 pass
 
+        return tmp_frames
+    finally:
+        pass
+
 
 def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
     fname = filename or os.path.basename(path)
@@ -219,11 +187,9 @@ def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
 
 
 def build_messages_for_image(prompt: str, b64_jpg: str = None, image_url: str = None):
-    # Keep user-visible prompt and include a single string content containing a data URL or image URL.
     if image_url:
         content = f"{prompt}\n\nImage: {image_url}"
     elif b64_jpg:
-        # Use explicit data URL so Mistral chat content is a single string (fixes Pydantic validation)
         content = f"{prompt}\n\nImage (base64): data:image/jpeg;base64,{b64_jpg}"
     else:
         raise ValueError("Either image_url or b64_jpg required")
@@ -234,6 +200,36 @@ def build_messages_for_text(prompt: str, extra_text: str):
     return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": f"{prompt}\n\n{extra_text}"}]
 
 
+def extract_delta(chunk):
+    if not chunk:
+        return None
+    data = getattr(chunk, "data", None) or getattr(chunk, "response", None) or getattr(chunk, "delta", None)
+    if not data:
+        return None
+    try:
+        content = data.choices[0].delta.content
+        if content is None:
+            return None
+        return str(content)
+    except Exception:
+        pass
+    try:
+        msg = data.choices[0].message
+        if isinstance(msg, dict):
+            content = msg.get("content")
+        else:
+            content = getattr(msg, "content", None)
+        if content is None:
+            return None
+        return str(content)
+    except Exception:
+        pass
+    try:
+        return str(data)
+    except Exception:
+        return None
+
+
 def extract_text_from_response(res, parts: list):
     try:
         choices = getattr(res, "choices", None) or res.get("choices", [])
@@ -265,60 +261,87 @@ def extract_text_from_response(res, parts: list):
 
 
 def stream_and_collect(client, model, messages, parts: list):
-    try:
-        # We assume messages are plain strings already for image/text functions above.
-        # For safety, if content is a list/dict (old codepath), coerce minimally here:
-        norm_msgs = []
-        for m in messages:
-            if not isinstance(m, dict):
-                norm_msgs.append(m)
-                continue
-            c = m.get("content")
-            if isinstance(c, list):
-                # prefer image_url or image_base64 if present, else join text chunks
-                picked = []
-                for item in c:
-                    if isinstance(item, dict):
-                        if item.get("type") == "image_url" and item.get("image_url"):
-                            picked.append(item["image_url"])
-                        elif item.get("type") == "image_base64" and item.get("image_base64"):
-                            picked.append("data:image/jpeg;base64," + item["image_base64"])
-                        elif item.get("type") == "text" and item.get("text"):
-                            picked.append(item["text"])
-                    elif isinstance(item, str):
-                        picked.append(item)
-                newc = "\n\n".join(p for p in picked if p).strip()
+    norm_msgs = []
+    for m in messages:
+        if not isinstance(m, dict):
+            norm_msgs.append(m)
+            continue
+        c = m.get("content")
+        if isinstance(c, list):
+            picked = []
+            for item in c:
+                if isinstance(item, dict):
+                    if item.get("type") == "image_url" and item.get("image_url"):
+                        picked.append(item["image_url"])
+                    elif item.get("type") == "image_base64" and item.get("image_base64"):
+                        picked.append("data:image/jpeg;base64," + item["image_base64"])
+                    elif item.get("type") == "text" and item.get("text"):
+                        picked.append(item["text"])
+                elif isinstance(item, str):
+                    picked.append(item)
+            newc = "\n\n".join(p for p in picked if p).strip()
+            nm = m.copy()
+            nm["content"] = newc
+            norm_msgs.append(nm)
+        else:
+            if not isinstance(c, str):
                 nm = m.copy()
-                nm["content"] = newc
+                nm["content"] = str(c or "")
                 norm_msgs.append(nm)
             else:
-                if not isinstance(c, str):
-                    nm = m.copy()
-                    nm["content"] = str(c or "")
-                    norm_msgs.append(nm)
-                else:
-                    norm_msgs.append(m)
+                norm_msgs.append(m)
 
-        # Try streaming first (client may expose .chat.stream); fall back to non-streaming complete
+    stream_gen = None
+    try:
+        stream_gen = client.chat.stream(model=model, messages=norm_msgs)
+    except Exception:
         stream_gen = None
-        try:
-            stream_gen = client.chat.stream(model=model, messages=norm_msgs)
-        except Exception:
-            stream_gen = None
-        if stream_gen:
-            for chunk in stream_gen:
-                d = extract_delta(chunk)
-                if d is None:
-                    continue
-                if d.strip() == "" and parts:
-                    continue
-                parts.append(d)
-            return
-
-        res = client.chat.complete(model=model, messages=norm_msgs, stream=False)
-        extract_text_from_response(res, parts)
-    except Exception as e:
-        parts.append(f"[Model error: {e}]")
+    if stream_gen:
+        for chunk in stream_gen:
+            d = extract_delta(chunk)
+            if d is None:
+                continue
+            if d.strip() == "" and parts:
+                continue
+            parts.append(d)
+        return
+
+    res = client.chat.complete(model=model, messages=norm_msgs, stream=False)
+    extract_text_from_response(res, parts)
+
+
+def analyze_image_bytes(client, img_bytes: bytes, prompt: str, model=DEFAULT_IMAGE_MODEL):
+    jpg = convert_to_jpeg_bytes(img_bytes, base_h=480)
+    b64 = b64_jpeg(jpg)
+    msgs = build_messages_for_image(prompt, b64_jpg=b64)
+    parts = []
+    stream_and_collect(client, model, msgs, parts)
+    return "".join(parts).strip()
+
+
+def analyze_multiple_frames(client, frames_bytes_list, prompt: str, model=DEFAULT_IMAGE_MODEL):
+    results = []
+    for i, fb in enumerate(frames_bytes_list):
+        res = analyze_image_bytes(client, fb, f"{prompt}\n\nFrame index: {i+1}", model=model)
+        results.append((i, res))
+
+    merged = []
+    for i, text in results:
+        merged.append(f"Frame {i+1} analysis:\n{text}")
+
+    consolidation_prompt = (
+        prompt
+        + "\n\nConsolidate the key consistent observations across the provided frame analyses below. "
+        "List consistent findings first, then note any differences between frames."
+        + "\n\n" + "\n\n".join(f"Frame {i+1}:\n{text}" for i, text in results)
+    )
+    parts = []
+    msgs = build_messages_for_text(consolidation_prompt, "")
+    stream_and_collect(client, DEFAULT_IMAGE_MODEL, msgs, parts)
+    consolidated = "".join(parts).strip()
+    if consolidated:
+        merged.append("Consolidated summary:\n" + consolidated)
+    return "\n\n".join(merged)
 
 
 def generate_final_text(src: str, custom_prompt: str, api_key: str):
@@ -328,31 +351,25 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
     is_image = ext in IMAGE_EXTS or (not is_remote(src) and os.path.isfile(src) and ext in IMAGE_EXTS)
     parts = []
 
-    # IMAGE path: keep exactly the previous behavior but send content as a single string data URL
     if is_image:
         try:
             if is_remote(src):
-                msgs = build_messages_for_image(prompt, image_url=src)
+                raw = fetch_bytes(src)
+                return analyze_image_bytes(client, raw, prompt, model=DEFAULT_IMAGE_MODEL)
             else:
                 raw = fetch_bytes(src)
-                jpg = convert_to_jpeg_bytes(raw, base_h=480)
-                b64 = b64_jpeg(jpg)
-                msgs = build_messages_for_image(prompt, b64_jpg=b64)
+                return analyze_image_bytes(client, raw, prompt, model=DEFAULT_IMAGE_MODEL)
+        except UnidentifiedImageError:
+            return "Error: provided file is not a valid image."
         except Exception as e:
             return f"Error processing image: {e}"
-        stream_and_collect(client, DEFAULT_IMAGE_MODEL, msgs, parts)
-        return "".join(parts).strip()
 
-    # REMOTE VIDEO path
     if is_remote(src):
+        tmp_media = None
         try:
             media_bytes = fetch_bytes(src, timeout=120)
-        except Exception as e:
-            return f"Error downloading remote media: {e}"
-        ext = ext_from_src(src) or ".mp4"
-        tmp_media = save_bytes_to_temp(media_bytes, suffix=ext)
-        try:
-            # Try uploading full video to Mistral files first
+            ext = ext_from_src(src) or ".mp4"
+            tmp_media = save_bytes_to_temp(media_bytes, suffix=ext)
             try:
                 file_id = upload_file_to_mistral(client, tmp_media, filename=os.path.basename(src.split("?")[0]))
                 extra = (
@@ -362,19 +379,11 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
                 msgs = build_messages_for_text(prompt, extra)
                 stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
                 return "".join(parts).strip()
-            except Exception as e_upload:
-                # Fallback: extract a representative frame and analyze as image
-                frame_bytes = extract_best_frame_bytes(tmp_media)
-                if not frame_bytes:
-                    return f"Error uploading to Mistral and no frame fallback available: {e_upload}"
-                try:
-                    jpg = convert_to_jpeg_bytes(frame_bytes, base_h=480)
-                except UnidentifiedImageError:
-                    jpg = frame_bytes
-                b64 = b64_jpeg(jpg)
-                msgs = build_messages_for_image(prompt, b64_jpg=b64)
-                stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
-                return "".join(parts).strip()
+            except Exception:
+                frames = extract_best_frames_bytes(tmp_media, sample_count=5)
+                if not frames:
+                    return "Error: could not upload remote video and no frames extracted."
+                return analyze_multiple_frames(client, frames, prompt, model=DEFAULT_IMAGE_MODEL)
         finally:
             try:
                 if tmp_media and os.path.exists(tmp_media):
@@ -382,7 +391,6 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             except Exception:
                 pass
 
-    # LOCAL VIDEO path
     tmp_media = None
     try:
         media_bytes = fetch_bytes(src)
@@ -390,7 +398,6 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         ext = ext or ".mp4"
         tmp_media = save_bytes_to_temp(media_bytes, suffix=ext)
         try:
-            # Upload local video file to Mistral files and analyze by file id
             file_id = upload_file_to_mistral(client, tmp_media, filename=os.path.basename(src))
             extra = (
                 f"Local video uploaded to Mistral Files with id: {file_id}\n\n"
@@ -400,15 +407,10 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
             return "".join(parts).strip()
         except Exception:
-            # fallback to frame extraction + image analysis
-            frame_bytes = extract_best_frame_bytes(tmp_media)
-            if not frame_bytes:
+            frames = extract_best_frames_bytes(tmp_media, sample_count=5)
+            if not frames:
                 return "Unable to process the provided file. Provide a direct image/frame URL or a remote video URL."
-            jpg = convert_to_jpeg_bytes(frame_bytes, base_h=480)
-            b64 = b64_jpeg(jpg)
-            msgs = build_messages_for_image(prompt, b64_jpg=b64)
-            stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
-            return "".join(parts).strip()
+            return analyze_multiple_frames(client, frames, prompt, model=DEFAULT_IMAGE_MODEL)
     finally:
         try:
             if tmp_media and os.path.exists(tmp_media):
@@ -419,16 +421,13 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
 
 css = ".preview_media img, .preview_media video { max-width: 100%; height: auto; }"
 
-
 def load_preview(url: str):
-    # Return (image_or_None, video_or_None, status_text)
     if not url:
         return None, None, ""
     try:
         r = requests.get(url, timeout=30, stream=True)
         r.raise_for_status()
         ctype = (r.headers.get("content-type") or "").lower()
-        # If content-type indicates video or URL ends with a video ext, return video preview
         if (ctype and ctype.startswith("video/")) or any(url.lower().split("?")[0].endswith(ext) for ext in VIDEO_EXTS):
             return None, url, "Video"
         data = r.content
@@ -451,17 +450,16 @@ with gr.Blocks(title="Flux", css=css) as demo:
             with gr.Accordion("Mistral API Key (optional)", open=False):
                 api_key = gr.Textbox(label="API Key", type="password", max_lines=1)
             submit = gr.Button("Submit")
-            preview_image = gr.Image(label="Preview", type="pil", elem_classes="preview_media")
-            preview_video = gr.Video(label="Preview", elem_classes="preview_media")
+            preview_image = gr.Image(label="Preview", type="pil", elem_classes="preview_media", visible=False)
+            preview_video = gr.Video(label="Preview", elem_classes="preview_media", visible=False)
 
         with gr.Column(scale=2):
             final_text = gr.Markdown(value="")
 
-    # Ensure preview outputs get None when not applicable so Gradio hides them
     url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video, gr.Textbox(visible=False)])
     submit.click(fn=generate_final_text, inputs=[url_input, custom_prompt, api_key], outputs=[final_text])
     demo.queue()
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
-            
+