Use free local Diffusers image model

app.py

@@ -103,9 +103,15 @@ class StylePlan:
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
+FREE_IMAGE_MODEL = os.environ.get("FREE_IMAGE_MODEL", "segmind/tiny-sd")
+FREE_IMAGE_STEPS = int(os.environ.get("FREE_IMAGE_STEPS", "5"))
+USE_HF_PROMPT_PROVIDER = os.environ.get("USE_HF_PROMPT_PROVIDER", "0") == "1"
+USE_HF_IMAGE_PROVIDER = os.environ.get("USE_HF_IMAGE_PROVIDER", "0") == "1"
 HF_IMAGE_MODEL = os.environ.get("HF_IMAGE_MODEL", "black-forest-labs/FLUX.1-schnell")
 HF_PROMPT_MODEL = os.environ.get("HF_PROMPT_MODEL", "Qwen/Qwen2.5-Coder-7B-Instruct:fastest")
 HF_PROMPT_ENDPOINT = os.environ.get("HF_PROMPT_ENDPOINT", "https://router.huggingface.co/v1/chat/completions")
+FREE_DIFFUSION_PIPE = None
+FREE_DIFFUSION_ERROR = None
 
 
 def slugify(value: str) -> str:
@@ -263,6 +269,8 @@ def short_error(exc: Exception) -> str:
 
 
 def hf_prompt_json(html_code: str, role_lines: list[tuple[str, str]], style_hint: str) -> tuple[dict[str, str] | None, str | None]:
+    if not USE_HF_PROMPT_PROVIDER:
+        return None, None
     if not HF_TOKEN:
         return None, "HF_TOKEN is not visible to the Space runtime"
 
@@ -978,7 +986,78 @@ def placeholder_png_bytes(role: str, width: int, height: int) -> bytes:
     return out.getvalue()
 
 
+def diffusion_dimensions(spec: AssetSpec) -> tuple[int, int]:
+    if is_background_spec(spec):
+        return 384, 216
+    return 256, 256
+
+
+def polish_diffusion_asset(image: Image.Image, spec: AssetSpec) -> bytes:
+    image = image.convert("RGBA")
+    if not is_background_spec(spec):
+        # Text-to-image models do not produce real transparency. This makes sprites
+        # usable by fading out colors similar to the generated corner background.
+        small = image.resize((128, 128), Image.LANCZOS)
+        corners = [
+            small.getpixel((0, 0)),
+            small.getpixel((127, 0)),
+            small.getpixel((0, 127)),
+            small.getpixel((127, 127)),
+        ]
+        bg = tuple(sum(pixel[i] for pixel in corners) // len(corners) for i in range(3))
+        pixels = small.load()
+        for y in range(small.height):
+            for x in range(small.width):
+                r, g, b, a = pixels[x, y]
+                dist = abs(r - bg[0]) + abs(g - bg[1]) + abs(b - bg[2])
+                edge = min(x, y, small.width - 1 - x, small.height - 1 - y)
+                if dist < 64 or edge < 3:
+                    a = 0
+                elif dist < 125:
+                    a = max(0, min(a, (dist - 64) * 4))
+                pixels[x, y] = (r, g, b, a)
+        image = small.resize((spec.width, spec.height), Image.LANCZOS)
+    else:
+        image = image.resize((spec.width, spec.height), Image.LANCZOS)
+    out = io.BytesIO()
+    image.save(out, format="PNG")
+    return out.getvalue()
+
+
+def free_diffusion_png(spec: AssetSpec, index: int, run_id: int) -> tuple[bytes | None, str | None]:
+    global FREE_DIFFUSION_PIPE, FREE_DIFFUSION_ERROR
+    if FREE_DIFFUSION_ERROR:
+        return None, FREE_DIFFUSION_ERROR
+    try:
+        import torch
+        from diffusers import DiffusionPipeline
+
+        if FREE_DIFFUSION_PIPE is None:
+            FREE_DIFFUSION_PIPE = DiffusionPipeline.from_pretrained(FREE_IMAGE_MODEL)
+            FREE_DIFFUSION_PIPE = FREE_DIFFUSION_PIPE.to("cpu")
+            if hasattr(FREE_DIFFUSION_PIPE, "enable_attention_slicing"):
+                FREE_DIFFUSION_PIPE.enable_attention_slicing()
+
+        width, height = diffusion_dimensions(spec)
+        seed = abs(hash(f"{spec.role}|{spec.prompt}|{index}|{run_id}")) % 2147483647
+        generator = torch.Generator(device="cpu").manual_seed(seed)
+        image = FREE_DIFFUSION_PIPE(
+            spec.prompt,
+            width=width,
+            height=height,
+            num_inference_steps=FREE_IMAGE_STEPS,
+            guidance_scale=6.0,
+            generator=generator,
+        ).images[0]
+        return polish_diffusion_asset(image, spec), None
+    except Exception as exc:
+        FREE_DIFFUSION_ERROR = short_error(exc)
+        return None, FREE_DIFFUSION_ERROR
+
+
 def hf_image_png(spec: AssetSpec, index: int, run_id: int) -> tuple[bytes | None, str | None]:
+    if not USE_HF_IMAGE_PROVIDER:
+        return None, None
     if not HF_TOKEN:
         return None, "HF_TOKEN is not visible to the Space runtime"
     if InferenceClient is None:
@@ -1003,15 +1082,20 @@ def hf_image_png(spec: AssetSpec, index: int, run_id: int) -> tuple[bytes | None
 
 
 def generate_asset(spec: AssetSpec, index: int, run_id: int) -> tuple[str, str, str | None, str]:
-    png_content, hf_error = hf_image_png(spec, index, run_id)
-    source = HF_IMAGE_MODEL
+    png_content, free_error = free_diffusion_png(spec, index, run_id)
+    source = FREE_IMAGE_MODEL
+    error = free_error
+    if png_content is None:
+        png_content, hf_error = hf_image_png(spec, index, run_id)
+        source = HF_IMAGE_MODEL
+        error = hf_error or free_error
     if png_content is None:
         png_content = local_asset_png(spec, index, run_id)
-        source = "local style fallback"
+        source = "local procedural fallback"
     return (
         png_bytes_to_data_uri(png_content),
         write_gallery_image(png_content, spec.role),
-        hf_error if source == "local style fallback" else None,
+        error if source == "local procedural fallback" else None,
         source,
     )
 
@@ -1191,7 +1275,7 @@ def generate_images_and_game(html_code: str, roles: str, style_hint: str):
         gallery.append((gallery_path, f"{spec.role} -> {spec.filename}"))
         model_rows.append((spec.role, prompt_model, image_model))
         if error:
-            errors.append(f"{spec.role}: image model failed ({error}); used local style fallback")
+            errors.append(f"{spec.role}: image model failed ({error}); used local procedural fallback")
 
     rewritten = embed_assets(html_code, assets, specs)
     status = (