better image handling

app.py

@@ -1,24 +1,10 @@
 import base64
-from pathlib import Path
-from typing import Any
+import io
 
 import gradio as gr
 from huggingface_hub import InferenceClient
 from PIL import Image
 
-
-def _image_to_data_url(image_path: str, max_side: int = 1120, quality: int = 85) -> str:
-    """Resize and encode a local image file as a base64 JPEG data URL."""
-    with Image.open(image_path) as img:
-        img = img.convert("RGB")
-        w, h = img.size
-        if max(w, h) > max_side:
-            scale = max_side / max(w, h)
-            img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
-        buf = io.BytesIO()
-        img.save(buf, format="JPEG", quality=quality)
-    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
-
 from agent import run_agent
 
 MODEL = "Qwen/Qwen3-VL-30B-A3B-Thinking"
@@ -33,72 +19,51 @@ DEFAULT_SYSTEM = (
 
 
 def _image_to_data_url(image_path: str, max_side: int = 1120, quality: int = 85) -> str:
-    """
-    Convert a local image file to a base64 data URL, resizing it first so the
-    payload stays within the HF router's request-size limit.
-
-    Images are downscaled so their longest side is at most `max_side` pixels,
-    then saved as JPEG at `quality` to keep the base64 size small.
-    """
-    from PIL import Image
-    import io
-
+    """Resize and base64-encode a local image so the HF payload stays under the limit."""
     with Image.open(image_path) as img:
         img = img.convert("RGB")
         w, h = img.size
         if max(w, h) > max_side:
             scale = max_side / max(w, h)
             img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
-
         buf = io.BytesIO()
         img.save(buf, format="JPEG", quality=quality)
-        data = base64.b64encode(buf.getvalue()).decode("utf-8")
-
-    return f"data:image/jpeg;base64,{data}"
-
-
-def _build_user_content(text: str, image_path: str | None) -> Any:
-    """
-    Build the `content` field for a user message.
-
-    Returns a plain string when there is no image, or a list of content parts
-    (text + image_url) when an image is present.
-    """
-    if not image_path:
-        return text or ""
-
-    parts: list[dict] = []
-    if text:
-        parts.append({"type": "text", "text": text})
-    parts.append(
-        {
-            "type": "image_url",
-            "image_url": {"url": _image_to_data_url(image_path)},
-        }
-    )
-    return parts
+    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
 
 
 def respond(
-    message: str,
+    message,
     history: list[dict],
-    image: str | None,
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
     hf_token: gr.OAuthToken,
 ):
     client = InferenceClient(token=hf_token.token, model=MODEL)
 
-    # Build the full message list: system + history + new user turn
-    messages: list[dict] = [{"role": "system", "content": system_message}]
+    # multimodal=True sends {"text": str, "files": [path, ...]}
+    # Guard against plain strings in case of edge-case history replay
+    if isinstance(message, dict):
+        text = message.get("text", "")
+        files = message.get("files", [])
+    else:
+        text = message or ""
+        files = []
+
+    if files:
+        content = []
+        if text:
+            content.append({"type": "text", "text": text})
+        for f in files:
+            content.append({"type": "image_url", "image_url": {"url": _image_to_data_url(f)}})
+    else:
+        content = text
+
+    messages = [{"role": "system", "content": system_message}]
     messages.extend(history)
+    messages.append({"role": "user", "content": content})
 
-    user_content = _build_user_content(message, image)
-    messages.append({"role": "user", "content": user_content})
-
-    # Run the agentic loop and get back the final answer string
     answer = run_agent(
         messages=messages,
         client=client,
@@ -108,36 +73,28 @@ def respond(
         top_p=top_p,
     )
 
-    # Yield in chunks to preserve Gradio's streaming UX
-    chunk_size = 8
     partial = ""
-    for i in range(0, len(answer), chunk_size):
-        partial += answer[i : i + chunk_size]
+    for i in range(0, len(answer), 8):
+        partial += answer[i : i + 8]
         yield partial
 
 
+chatbot = gr.ChatInterface(
+    respond,
+    multimodal=True,
+    additional_inputs=[
+        gr.Textbox(value=DEFAULT_SYSTEM, label="System message"),
+        gr.Slider(minimum=1, maximum=16384, value=8192, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+    ],
+)
+
 with gr.Blocks() as demo:
     with gr.Sidebar():
         gr.LoginButton()
-        image_input = gr.Image(label="Upload image", type="filepath")
-
-    gr.ChatInterface(
-        respond,
-        additional_inputs=[
-            image_input,
-            gr.Textbox(value=DEFAULT_SYSTEM, label="System message"),
-            gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-            gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.05,
-                label="Top-p (nucleus sampling)",
-            ),
-        ],
-    )
+    chatbot.render()
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(ssr_mode=False)