Update app.py

app.py

@@ -10,10 +10,9 @@ from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
 from mistralai import Mistral
 
-# Config
 DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
-PIXTRAL_MODEL = "pixtral-12b-2409"     # image-capable multimodal model
-VIDEO_MODEL = "voxtral-mini-latest"    # replace with your preferred video model
+PIXTRAL_MODEL = "pixtral-12b-2409"
+VIDEO_MODEL = "voxtral-mini-latest"
 STREAM_THRESHOLD = 20 * 1024 * 1024
 FFMPEG_BIN = shutil.which("ffmpeg")
 
@@ -30,7 +29,7 @@ Image.MAX_IMAGE_PIXELS = 10000 * 10000
 IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
 VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".flv")
 
-def get_client(key: str = None):
+def get_client(key: str = None) -> Mistral:
     api_key = (key or "").strip() or DEFAULT_KEY
     return Mistral(api_key=api_key)
 
@@ -63,14 +62,14 @@ def fetch_bytes(src: str, stream_threshold=STREAM_THRESHOLD, timeout=60) -> byte
     with open(src, "rb") as f:
         return f.read()
 
-def save_bytes_to_temp(b: bytes, suffix: str):
+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=480) -> bytes:
+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):
@@ -89,21 +88,14 @@ def convert_to_jpeg_bytes(img_bytes: bytes, base_h=480) -> bytes:
 def b64_jpeg(img_bytes: bytes) -> str:
     return base64.b64encode(img_bytes).decode("utf-8")
 
-def extract_best_frames_bytes(media_path: str, sample_count: int = 5, timeout_probe: int = 10, timeout_extract: int = 15):
+def extract_best_frames_bytes(media_path: str, sample_count: int = 5, timeout_probe: int = 10, timeout_extract: int = 15) -> list:
     if not FFMPEG_BIN or not os.path.exists(media_path):
         return []
-    frames = []
     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)
-    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
+        out = subprocess.check_output(probe_cmd, timeout=timeout_probe).strip()
+        duration = float(out) if out else None
     except Exception:
         duration = None
 
@@ -112,42 +104,29 @@ def extract_best_frames_bytes(media_path: str, sample_count: int = 5, timeout_pr
     else:
         timestamps = [0.5, 1.0, 2.0][:sample_count]
 
+    frames = []
     for i, t in enumerate(timestamps):
-        fd, tmp_frame = tempfile.mkstemp(suffix=f"_{i}.jpg")
+        fd, tmp = tempfile.mkstemp(suffix=f"_{i}.jpg")
         os.close(fd)
-        cmd = [
-            FFMPEG_BIN, "-nostdin", "-y", "-i", media_path,
-            "-ss", str(t), "-frames:v", "1", "-q:v", "2", tmp_frame
-        ]
-        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        cmd = [FFMPEG_BIN, "-nostdin", "-y", "-i", media_path, "-ss", str(t), "-frames:v", "1", "-q:v", "2", tmp]
         try:
-            proc.communicate(timeout=timeout_extract)
-        except subprocess.TimeoutExpired:
-            try: proc.kill()
+            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())
+        finally:
+            try: os.remove(tmp)
             except Exception: pass
-            proc.communicate()
-        if proc.returncode == 0 and os.path.exists(tmp_frame) and os.path.getsize(tmp_frame) > 0:
-            with open(tmp_frame, "rb") as f:
-                frames.append(f.read())
-        try:
-            if os.path.exists(tmp_frame): os.remove(tmp_frame)
-        except Exception:
-            pass
     return frames
 
-def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
+def upload_file_to_mistral(client: Mistral, path: str, filename: str | None = None, purpose: str = "batch") -> str:
     fname = filename or os.path.basename(path)
     try:
         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:
-            try:
-                fid = res["data"][0]["id"]
-            except Exception:
-                pass
-        if not fid:
-            raise RuntimeError(f"No file id returned: {res}")
+            fid = res["data"][0]["id"]
         return fid
     except Exception:
         api_key = client.api_key if hasattr(client, "api_key") else os.getenv("MISTRAL_API_KEY", "")
@@ -159,269 +138,163 @@ def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
             r = requests.post(url, headers=headers, files=files, data=data, timeout=120)
             r.raise_for_status()
             jr = r.json()
-            fid = jr.get("id") or jr.get("data", [{}])[0].get("id")
-            if not fid:
-                raise RuntimeError(f"Upload failed to return id: {jr}")
-            return fid
+            return jr.get("id") or jr.get("data", [{}])[0].get("id")
 
-def build_messages_for_image(prompt: str, b64_jpg: str = None, image_url: str = None):
-    if image_url:
-        content = f"{prompt}\n\nImage: {image_url}"
-    elif b64_jpg:
-        content = f"{prompt}\n\nImage (base64): data:image/jpeg;base64,{b64_jpg}"
-    else:
-        raise ValueError("Either image_url or b64_jpg required")
+def build_messages_for_image(prompt: str, b64_jpg: str) -> list:
+    content = f"{prompt}\n\nImage (base64): data:image/jpeg;base64,{b64_jpg}"
     return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": content}]
 
-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 build_messages_for_text(prompt: str, extra: str) -> list:
+    return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": f"{prompt}\n\n{extra}"}]
 
-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
+def extract_text_from_response(res, parts: list):
     try:
-        content = data.choices[0].delta.content
-        if content is None:
-            return None
-        return str(content)
+        choices = getattr(res, "choices", None) or res.get("choices", [])
     except Exception:
-        pass
+        choices = []
+    if not choices:
+        parts.append(str(res))
+        return
     try:
-        msg = data.choices[0].message
+        msg = 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", [])
-    except Exception:
-        choices = []
-    if choices:
-        try:
-            msg = choices[0].message
-            if isinstance(msg, dict):
-                content = msg.get("content")
-            else:
-                content = getattr(msg, "content", None)
-            if content:
-                if isinstance(content, str):
-                    parts.append(content)
-                else:
-                    if isinstance(content, list):
-                        for c in content:
-                            if isinstance(c, dict) and c.get("type") == "text":
-                                parts.append(c.get("text", ""))
-                    elif isinstance(content, dict):
-                        text = content.get("text") or content.get("content")
-                        if text:
-                            parts.append(text)
-        except Exception:
-            parts.append(str(res))
-    else:
-        parts.append(str(res))
-
-def stream_and_collect(client, model, messages, parts: list):
-    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)
+        if isinstance(content, str):
+            parts.append(content)
         else:
-            if not isinstance(c, str):
-                nm = m.copy()
-                nm["content"] = str(c or "")
-                norm_msgs.append(nm)
-            else:
-                norm_msgs.append(m)
-
-    stream_gen = None
-    try:
-        stream_gen = client.chat.stream(model=model, messages=norm_msgs)
+            parts.append(str(content))
     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)
+        parts.append(str(res))
 
-def analyze_image_bytes(client, img_bytes: bytes, prompt: str, model=PIXTRAL_MODEL):
-    jpg = convert_to_jpeg_bytes(img_bytes, base_h=480)
-    b64 = b64_jpeg(jpg)
-    msgs = build_messages_for_image(prompt, b64_jpg=b64)
+def chat_complete(client: Mistral, model: str, messages: list) -> str:
     parts = []
-    stream_and_collect(client, model, msgs, parts)
+    res = client.chat.complete(model=model, messages=messages, stream=False)
+    extract_text_from_response(res, parts)
     return "".join(parts).strip()
 
-def analyze_multiple_frames(client, frames_bytes_list, prompt: str, model=PIXTRAL_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)
+def analyze_image(client: Mistral, img_bytes: bytes, prompt: str) -> str:
+    jpeg = convert_to_jpeg_bytes(img_bytes, base_h=480)
+    b64 = b64_jpeg(jpeg)
+    msgs = build_messages_for_image(prompt, b64)
+    return chat_complete(client, PIXTRAL_MODEL, msgs)
+
+def analyze_frames_and_consolidate(client: Mistral, frames: list, prompt: str) -> str:
+    per_frame = []
+    for i, fb in enumerate(frames):
+        txt = analyze_image(client, fb, f"{prompt}\n\nFrame index: {i + 1}")
+        per_frame.append(f"Frame {i + 1} analysis:\n{txt}")
+    consolidation = (
+        f"{prompt}\n\n"
+        "Consolidate 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(per_frame)
     )
-    parts = []
-    msgs = build_messages_for_text(consolidation_prompt, "")
-    stream_and_collect(client, PIXTRAL_MODEL, msgs, parts)
-    consolidated = "".join(parts).strip()
-    if consolidated:
-        merged.append("Consolidated summary:\n" + consolidated)
-    return "\n\n".join(merged)
+    msgs = build_messages_for_text(consolidation, "")
+    summary = chat_complete(client, PIXTRAL_MODEL, msgs)
+    return "\n\n".join(per_frame + [f"Consolidated summary:\n{summary}"])
 
-def generate_final_text(src: str, custom_prompt: str, api_key: str):
+def process_media(src: str, custom_prompt: str, api_key: str) -> str:
     client = get_client(api_key)
-    prompt = (custom_prompt.strip() if custom_prompt and custom_prompt.strip() else "Please provide a detailed visual review.")
-    parts = []
+    prompt = custom_prompt.strip() or "Please provide a detailed visual review."
     ext = ext_from_src(src)
-    is_image = ext in IMAGE_EXTS or (not is_remote(src) and os.path.isfile(src) and ext in IMAGE_EXTS)
-    is_video = ext in VIDEO_EXTS or (not is_remote(src) and os.path.isfile(src) and ext in VIDEO_EXTS)
-
-    # If remote and content-type suggests video, treat as video
+    is_image = ext in IMAGE_EXTS
+    is_video = ext in VIDEO_EXTS
     if is_remote(src):
         try:
-            r = requests.head(src, timeout=10, allow_redirects=True)
-            ctype = (r.headers.get("content-type") or "").lower()
+            h = requests.head(src, timeout=8, allow_redirects=True)
+            ctype = (h.headers.get("content-type") or "").lower()
             if ctype.startswith("video/"):
-                is_video = True
+                is_video = True; is_image = False
             elif ctype.startswith("image/"):
-                is_image = True
+                is_image = True; is_video = False
         except Exception:
             pass
-
     if is_image:
         try:
             raw = fetch_bytes(src)
         except Exception as e:
             return f"Error fetching image: {e}"
         try:
-            return analyze_image_bytes(client, raw, prompt, model=PIXTRAL_MODEL)
+            return analyze_image(client, raw, prompt)
         except UnidentifiedImageError:
             return "Error: provided file is not a valid image."
         except Exception as e:
             return f"Error processing image: {e}"
-
     if is_video:
-        tmp_media = None
+        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 or ".mp4")
         try:
             try:
-                media_bytes = fetch_bytes(src, timeout=120)
-            except Exception as e:
-                return f"Error fetching video: {e}"
-            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]))
+                file_id = upload_file_to_mistral(client, tmp_path, filename=os.path.basename(src.split("?")[0]))
                 extra = (
                     f"Uploaded video to Mistral Files with id: {file_id}\n\n"
                     "Instruction: Analyze the video contents using the uploaded file id. Do not invent frames not present."
                 )
                 msgs = build_messages_for_text(prompt, extra)
-                stream_and_collect(client, VIDEO_MODEL, msgs, parts)
-                return "".join(parts).strip()
+                return chat_complete(client, VIDEO_MODEL, msgs)
             except Exception:
-                frames = extract_best_frames_bytes(tmp_media, sample_count=5)
+                frames = extract_best_frames_bytes(tmp_path, 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=PIXTRAL_MODEL)
+                return analyze_frames_and_consolidate(client, frames, prompt)
         finally:
-            try:
-                if tmp_media and os.path.exists(tmp_media):
-                    os.remove(tmp_media)
-            except Exception:
-                pass
-
+            try: os.remove(tmp_path)
+            except Exception: pass
     return "Unable to determine media type from the provided URL or file extension."
 
-# UI helpers
 css = ".preview_media img, .preview_media video { max-width: 100%; height: auto; }"
+
 def load_preview(url: str):
     if not url:
         return None, None, ""
+    if not is_remote(url) and os.path.exists(url):
+        ext = ext_from_src(url)
+        if ext in VIDEO_EXTS:
+            return None, os.path.abspath(url), "Video"
+        if ext in IMAGE_EXTS:
+            try:
+                img = Image.open(url)
+                return img.convert("RGB"), None, "Image"
+            except Exception:
+                return None, None, "Preview failed"
     try:
-        r = requests.get(url, timeout=30, stream=True)
-        r.raise_for_status()
-        ctype = (r.headers.get("content-type") or "").lower()
-        if (ctype and ctype.startswith("video/")) or any(url.lower().split("?")[0].endswith(ext) for ext in VIDEO_EXTS):
+        h = requests.head(url, timeout=8, allow_redirects=True)
+        ctype = (h.headers.get("content-type") or "").lower()
+        if ctype.startswith("video/") or any(url.lower().split("?")[0].endswith(ext) for ext in VIDEO_EXTS):
             return None, url, "Video"
-        data = r.content
-        try:
-            img = Image.open(BytesIO(data))
-            if getattr(img, "is_animated", False):
-                img.seek(0)
-            img = img.convert("RGB")
-        except UnidentifiedImageError:
-            return None, None, "Preview failed"
-        return img, None, "Image"
+        r = requests.get(url, timeout=20)
+        r.raise_for_status()
+        img = Image.open(BytesIO(r.content))
+        if getattr(img, "is_animated", False):
+            img.seek(0)
+        return img.convert("RGB"), None, "Image"
     except Exception:
         return None, None, "Preview failed"
 
-# Gradio app
-with gr.Blocks(title="Flux Multimodal", css=css) as demo:
+with gr.Blocks(title="Flux Multimodal (fixed)", css=css) as demo:
     with gr.Row():
         with gr.Column(scale=1):
-            url_input = gr.Textbox(label="Image or Video URL or local path", placeholder="https://... or /path/to/file", lines=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 = gr.Button("Submit")
+            submit_btn = gr.Button("Submit")
             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_text = gr.Markdown(value="")
+            final_md = gr.Markdown(value="")
 
-    def _preview_wrapper(url):
+    def preview_update(url):
         img, vid, label = load_preview(url)
-        return img, vid, label
+        return (img if label == "Image" else None, vid if label == "Video" else None, label == "Image", label == "Video")
 
-    url_input.change(fn=_preview_wrapper, 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()
+    url_input.change(fn=preview_update, inputs=[url_input], outputs=[preview_image, preview_video, preview_image, preview_video])
+    submit_btn.click(fn=process_media, inputs=[url_input, custom_prompt, api_key], outputs=[final_md])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    demo.queue().launch()