Fix gallery image outputs

app.py

@@ -1,6 +1,7 @@
 import base64
 import io
 import re
+import tempfile
 import time
 from dataclasses import dataclass
 from html import escape
@@ -8,7 +9,7 @@ from urllib.parse import quote
 
 import gradio as gr
 import requests
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 
 
 POLLINATIONS_URL = (
@@ -118,25 +119,45 @@ def parse_assets(raw_roles: str, style_hint: str) -> list[AssetSpec]:
     return specs
 
 
-def image_to_data_uri(content: bytes, width: int, height: int) -> str:
+def image_to_png_bytes(content: bytes, width: int, height: int) -> bytes:
     image = Image.open(io.BytesIO(content)).convert("RGBA")
     image = image.resize((width, height), Image.LANCZOS)
     out = io.BytesIO()
     image.save(out, format="PNG")
-    return "data:image/png;base64," + base64.b64encode(out.getvalue()).decode("ascii")
+    return out.getvalue()
 
 
-def placeholder_data_uri(role: str, width: int, height: int) -> str:
+def png_bytes_to_data_uri(content: bytes) -> str:
+    return "data:image/png;base64," + base64.b64encode(content).decode("ascii")
+
+
+def write_gallery_image(content: bytes, role: str) -> str:
+    handle = tempfile.NamedTemporaryFile(
+        prefix=f"{slugify(role)}_",
+        suffix=".png",
+        delete=False,
+    )
+    handle.write(content)
+    handle.close()
+    return handle.name
+
+
+def placeholder_png_bytes(role: str, width: int, height: int) -> bytes:
     label = escape(role[:18])
-    svg = f"""<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}">
-<rect width="100%" height="100%" fill="#222"/>
-<rect x="4" y="4" width="{width - 8}" height="{height - 8}" rx="12" fill="#3b82f6"/>
-<text x="50%" y="50%" text-anchor="middle" dominant-baseline="middle" fill="white" font-family="Arial" font-size="18">{label}</text>
-</svg>"""
-    return "data:image/svg+xml;base64," + base64.b64encode(svg.encode("utf-8")).decode("ascii")
+    image = Image.new("RGBA", (width, height), "#222222")
+    draw = ImageDraw.Draw(image)
+    draw.rounded_rectangle((4, 4, width - 4, height - 4), radius=12, fill="#3b82f6")
+    font = ImageFont.load_default()
+    bbox = draw.textbbox((0, 0), label, font=font)
+    x = (width - (bbox[2] - bbox[0])) / 2
+    y = (height - (bbox[3] - bbox[1])) / 2
+    draw.text((x, y), label, fill="white", font=font)
+    out = io.BytesIO()
+    image.save(out, format="PNG")
+    return out.getvalue()
 
 
-def generate_asset(spec: AssetSpec, index: int) -> tuple[str, str | None]:
+def generate_asset(spec: AssetSpec, index: int) -> tuple[str, str, str | None]:
     seed = abs(hash((spec.role, spec.prompt, index))) % 999999
     url = POLLINATIONS_URL.format(
         prompt=quote(spec.prompt),
@@ -149,13 +170,28 @@ def generate_asset(spec: AssetSpec, index: int) -> tuple[str, str | None]:
         try:
             response = requests.get(url, timeout=120)
             response.raise_for_status()
-            return image_to_data_uri(response.content, spec.width, spec.height), None
+            png_content = image_to_png_bytes(response.content, spec.width, spec.height)
+            return (
+                png_bytes_to_data_uri(png_content),
+                write_gallery_image(png_content, spec.role),
+                None,
+            )
         except Exception as exc:
             if attempt < 2:
                 time.sleep(8)
             else:
-                return placeholder_data_uri(spec.role, spec.width, spec.height), str(exc)
-    return placeholder_data_uri(spec.role, spec.width, spec.height), "Unknown generation error"
+                png_content = placeholder_png_bytes(spec.role, spec.width, spec.height)
+                return (
+                    png_bytes_to_data_uri(png_content),
+                    write_gallery_image(png_content, spec.role),
+                    str(exc),
+                )
+    png_content = placeholder_png_bytes(spec.role, spec.width, spec.height)
+    return (
+        png_bytes_to_data_uri(png_content),
+        write_gallery_image(png_content, spec.role),
+        "Unknown generation error",
+    )
 
 
 def replacement_names(spec: AssetSpec) -> set[str]:
@@ -235,9 +271,9 @@ def generate_images_and_game(html_code: str, roles: str, style_hint: str):
     errors = []
 
     for index, spec in enumerate(specs):
-        data_uri, error = generate_asset(spec, index)
+        data_uri, gallery_path, error = generate_asset(spec, index)
         assets[spec.role] = data_uri
-        gallery.append((data_uri, f"{spec.role} -> {spec.filename}"))
+        gallery.append((gallery_path, f"{spec.role} -> {spec.filename}"))
         if error:
             errors.append(f"{spec.role}: fallback used ({error})")