Update app.py

app.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-
 import os
 import subprocess
 import tempfile
@@ -11,11 +10,10 @@ from PIL import Image, UnidentifiedImageError
 import gradio as gr
 from mistralai import Mistral
 
-# Config
 DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
 DEFAULT_IMAGE_MODEL = "pixtral-12b-2409"
 DEFAULT_VIDEO_MODEL = "voxtral-mini-latest"
-STREAM_THRESHOLD = 20 * 1024 * 1024  # 20 MB
+STREAM_THRESHOLD = 20 * 1024 * 1024
 
 SYSTEM_INSTRUCTION = (
     "You are a clinical visual analyst. Only analyze media actually provided (image data or extracted frames). "
@@ -69,7 +67,6 @@ def fetch_bytes(src: str, stream_threshold=STREAM_THRESHOLD, timeout=60) -> byte
 def convert_to_jpeg_bytes(media_bytes: bytes, base_h=480) -> bytes:
     img = Image.open(BytesIO(media_bytes))
     try:
-        # For animated images, use first frame
         if getattr(img, "is_animated", False):
             img.seek(0)
     except Exception:
@@ -193,7 +190,6 @@ def extract_best_frame_bytes(media_path: str, sample_count: int = 5, timeout_pro
 
 def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
     fname = filename or os.path.basename(path)
-    # Try SDK upload
     try:
         with open(path, "rb") as fh:
             res = client.files.upload(file={"file_name": fname, "content": fh}, purpose=purpose)
@@ -207,7 +203,6 @@ def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
             raise RuntimeError(f"No file id returned: {res}")
         return fid
     except Exception:
-        # Fallback to HTTP upload
         api_key = client.api_key if hasattr(client, "api_key") else os.getenv("MISTRAL_API_KEY", "")
         url = "https://api.mistral.ai/v1/files"
         headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
@@ -224,38 +219,21 @@ 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):
-    """
-    Build messages using structured content per Mistral vision API:
-    - For remote images: include an {"type":"image_url","image_url":...} item
-    - For local bytes: include {"type":"image_base64","image_base64": "..."} (no data: URI prefix)
-    The user content is a list of typed items.
-    """
-    user_content = []
-    user_content.append({"type": "text", "text": prompt})
+    user_content = [{"type": "text", "text": prompt}]
     if image_url:
         user_content.append({"type": "image_url", "image_url": image_url})
     elif b64_jpg:
         user_content.append({"type": "image_base64", "image_base64": b64_jpg})
     else:
         raise ValueError("Either image_url or b64_jpg required")
-    return [
-        {"role": "system", "content": SYSTEM_INSTRUCTION},
-        {"role": "user", "content": user_content},
-    ]
+    return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": user_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}"},
-    ]
+    return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": f"{prompt}\n\n{extra_text}"}]
 
 
 def stream_and_collect(client, model, messages, parts: list):
-    """
-    Use client.chat.stream if available; otherwise use complete.
-    Appends textual pieces to parts list.
-    """
     try:
         stream_gen = None
         try:
@@ -304,16 +282,12 @@ def stream_and_collect(client, model, messages, parts: list):
 
 
 def generate_final_text(src: str, custom_prompt: str, api_key: str):
-    """
-    Main entry for Submit button. Returns final text (string).
-    """
     client = get_client(api_key)
     prompt = (custom_prompt.strip() if custom_prompt and custom_prompt.strip() else "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)
     parts = []
 
-    # Image handling: remote image_url or local image_base64
     if is_image:
         try:
             if is_remote(src):
@@ -321,16 +295,14 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             else:
                 raw = fetch_bytes(src)
                 jpg = convert_to_jpeg_bytes(raw, base_h=480)
-                b64 = b64_jpeg(jpg)  # NOTE: this is plain base64 string (no data: prefix)
+                b64 = b64_jpeg(jpg)
                 msgs = build_messages_for_image(prompt, b64_jpg=b64)
         except Exception as e:
             return f"Error processing image: {e}"
         stream_and_collect(client, DEFAULT_IMAGE_MODEL, msgs, parts)
         return "".join(parts).strip()
 
-    # Video handling (remote/local)
     if is_remote(src):
-        # download remote media, try upload to Mistral Files; fallback to a representative frame
         try:
             media_bytes = fetch_bytes(src, timeout=120)
         except Exception as e:
@@ -341,7 +313,6 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             try:
                 file_id = upload_file_to_mistral(client, tmp_media, filename=os.path.basename(src.split("?")[0]))
             except Exception as e:
-                # fallback to sending representative frame
                 frame_bytes = extract_best_frame_bytes(tmp_media)
                 if not frame_bytes:
                     return f"Error uploading to Mistral and no frame fallback available: {e}"
@@ -354,10 +325,7 @@ 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()
 
-            extra = (
-                f"Remote video uploaded 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."
-            )
+            extra = f"Remote video uploaded to Mistral Files with id: {file_id}\n\nInstruction: 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, DEFAULT_VIDEO_MODEL, msgs, parts)
             return "".join(parts).strip()
@@ -368,7 +336,6 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             except Exception:
                 pass
 
-    # Local video: upload or fallback to frames
     tmp_media = None
     try:
         media_bytes = fetch_bytes(src)
@@ -377,10 +344,7 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         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))
-            extra = (
-                f"Local video uploaded 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."
-            )
+            extra = f"Local video uploaded to Mistral Files with id: {file_id}\n\nInstruction: 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, DEFAULT_VIDEO_MODEL, msgs, parts)
             return "".join(parts).strip()
@@ -401,17 +365,10 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             pass
 
 
-# --- Gradio UI ---
-css = """
-.preview_media img, .preview_media video { max-width: 100%; height: auto; }
-"""
+css = ".preview_media img, .preview_media video { max-width: 100%; height: auto; }"
+
 
 def load_preview(url: str):
-    """
-    Returns: (image_or_None, video_or_None, mime_label)
-    - For images: return PIL.Image, None, "Image"
-    - For videos: return None, url, "Video"
-    """
     if not url:
         return None, None, ""
     try:
@@ -436,25 +393,20 @@ def load_preview(url: str):
 with gr.Blocks(title="Flux", css=css) as demo:
     with gr.Row():
         with gr.Column(scale=1):
-            # Top-left controls
             url_input = gr.Textbox(label="Image or Video URL", placeholder="https://...", 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")
-
-            # Single preview area (either image or video shown)
             preview_image = gr.Image(label="Preview", type="pil", elem_classes="preview_media")
             preview_video = gr.Video(label="Preview", elem_classes="preview_media")
 
         with gr.Column(scale=2):
-            # Right column: plain text output (rendered as Markdown/HTML allowed)
-            final_text = gr.Markdown(value="")  # use Markdown so long text renders nicely
+            final_text = gr.Markdown(value="")
 
-    # Wire up events
     url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video, gr.Textbox(visible=False)])
-    # For submit, use queue to avoid blocking UI
-    submit.click(fn=generate_final_text, inputs=[url_input, custom_prompt, api_key], outputs=[final_text], queue=True)
+    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)), enable_queue=True)
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))