Handle unsupported/corrupt image uploads with a clear error; add HEIC support

app.py

@@ -1,8 +1,10 @@
 import os
+import tempfile
 import spaces
 
 import torch
 import gradio as gr
+from PIL import Image
 from threading import Thread
 from transformers import (AutoProcessor, AutoModelForImageTextToText,
                           BitsAndBytesConfig, TextIteratorStreamer)
@@ -33,6 +35,55 @@ model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, **load_kwargs)
 model.eval()
 
 
+# transformers' image loader delegates to torchvision.decode_image, which accepts only
+# jpeg/png/webp/gif and raises a bare RuntimeError on anything else -- the user saw a
+# 30-line traceback for an ordinary phone photo. Normalising through PIL first means
+# BMP/TIFF/ICO (and HEIC, when pillow-heif is installed) work instead of failing, and
+# anything genuinely unreadable produces a message that says what to do about it.
+try:
+    import pillow_heif
+    pillow_heif.register_heif_opener()
+    _HEIF = True
+except Exception:
+    _HEIF = False
+
+NATIVE_FORMATS = {"JPEG", "PNG", "WEBP", "GIF"}
+SUPPORTED_HINT = ("JPEG, PNG, WebP, GIF, BMP, TIFF"
+                  + (", HEIC/HEIF" if _HEIF else ""))
+_image_cache = {}
+
+
+def prepare_image(path):
+    """Return a path the processor can definitely decode, or raise a clear gr.Error."""
+    cached = _image_cache.get(path)
+    if cached and os.path.exists(cached):
+        return cached
+
+    name = os.path.basename(path or "file")
+    try:
+        with Image.open(path) as probe:
+            probe.verify()                      # catches truncated/corrupt files
+        with Image.open(path) as img:
+            fmt = (img.format or "").upper()
+            if fmt in NATIVE_FORMATS and img.mode in ("RGB", "L"):
+                _image_cache[path] = path
+                return path
+            converted = img.convert("RGB")
+        handle, out = tempfile.mkstemp(suffix=".png")
+        os.close(handle)
+        converted.save(out, format="PNG")
+    except gr.Error:
+        raise
+    except Exception as exc:
+        raise gr.Error(
+            f"Couldn't read the image '{name}'. Supported formats: {SUPPORTED_HINT}. "
+            f"If it's an AVIF or HEIC photo, re-save it as JPEG or PNG and try again."
+        ) from exc
+
+    _image_cache[path] = out
+    return out
+
+
 def to_messages(message, history):
     messages = []
     for turn in history:
@@ -42,13 +93,14 @@ def to_messages(message, history):
             continue
         if isinstance(content, (tuple, list)):
             messages.append({"role": role,
-                             "content": [{"type": "image", "image": content[0]}]})
+                             "content": [{"type": "image",
+                                          "image": prepare_image(content[0])}]})
         else:
             messages.append({"role": role, "content": [{"type": "text", "text": content}]})
 
     content = []
     for path in message.get("files") or []:
-        content.append({"type": "image", "image": path})
+        content.append({"type": "image", "image": prepare_image(path)})
     if message.get("text"):
         content.append({"type": "text", "text": message["text"]})
     messages.append({"role": "user", "content": content})
@@ -70,10 +122,6 @@ def split_thinking(text):
     return [gr.ChatMessage(role="assistant", content=text)]
 
 
-# size='xlarge': this is a 27B at bf16 (~54.7 GB packed), which does not fit the default
-# ZeroGPU slice -- the symptom was "GPU task aborted" with nothing in the Space logs.
-# duration is kept at 120s, the usual ZeroGPU ceiling; 150 was over it.
-@spaces.GPU(duration=120, size="xlarge")
 def respond(message: dict, history: list, temperature: float, top_p: float, top_k: int):
     """Chat with Qwen3.8-27B-Uncensored. Accepts text and images, streams the reply.
 
@@ -84,7 +132,18 @@ def respond(message: dict, history: list, temperature: float, top_p: float, top_
         top_p: Nucleus sampling cutoff.
         top_k: Top-k sampling cutoff.
     """
+    # Build (and validate) the turn outside the GPU worker, so an unreadable upload
+    # fails immediately with a readable message instead of consuming a ZeroGPU slot
+    # and surfacing as a traceback from inside the worker.
     messages = to_messages(message, history)
+    yield from generate(messages, temperature, top_p, top_k)
+
+
+# size='xlarge': this is a 27B at bf16 (~54.7 GB packed), which does not fit the default
+# ZeroGPU slice -- the symptom was "GPU task aborted" with nothing in the Space logs.
+# duration is kept at 120s, the usual ZeroGPU ceiling; 150 was over it.
+@spaces.GPU(duration=120, size="xlarge")
+def generate(messages: list, temperature: float, top_p: float, top_k: int):
     inputs = processor.apply_chat_template(
         messages,
         add_generation_prompt=True,

requirements.txt

@@ -5,3 +5,6 @@ transformers==5.15.0
 accelerate==1.14.0
 bitsandbytes==0.50.1
 torchvision
+# HEIC/HEIF support: torchvision's decoder rejects them outright, so phone photos
+# fail without this. Optional at runtime -- app.py degrades to a clear error.
+pillow-heif