app.py

"""PixelLock Space app: custom UI at `/`, real Gradio mounted at `/gradio`.

The public-facing experience is a hand-rolled HTML/CSS/JS workbench, but the
process still runs a genuine Gradio app for Build Small eligibility. The model
path, grammar-constrained llama.cpp call, and footprint verification match the
validated Gradio prototype.
"""

from __future__ import annotations

import base64
import io
import json
import os
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

import gradio as gr
import httpx
from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
from pydantic import BaseModel


APP_DIR = Path(__file__).resolve().parent
for _path in (APP_DIR, APP_DIR.parent, APP_DIR.parent / "scratch"):
    sys.path.insert(0, str(_path))

import config  # noqa: E402
import validate  # noqa: E402

config.MAX_PALETTE = 64
validate.MAX_PALETTE = 64

import pixel_editor as pe  # noqa: E402


ENDPOINT = os.environ.get(
    "PIXELLOCK_ENDPOINT", "http://127.0.0.1:8080/v1/chat/completions"
)
MODEL = os.environ.get("PIXELLOCK_MODEL", "pixellock")
MAX_DIM = int(os.environ.get("PIXELLOCK_MAX_DIM", "64"))

STATIC_DIR = APP_DIR / "static"
EXAMPLES_DIR = STATIC_DIR / "examples"
ASSETS_DIR = STATIC_DIR / "assets"

THEMES = {
    "🌋  Molten lava": "a molten lava theme — glowing magma core, charred black edges, white-hot highlights, bold orange-red ramp",
    "❄️  Frozen ice": "a frozen ice theme — pale cyan, frosted surface, white frost highlights, cool blue shadows",
    "🪙  Solid gold": "a solid gold royal theme — shimmering gold ramp, warm highlights, deep amber shadows",
    "☢️  Toxic": "a toxic radioactive theme — sickly neon green, glowing hazard spots, dark slime shadows",
    "🍂  Autumn dusk": "an autumn dusk theme — warm amber and crimson with deep purple shadows",
    "🌌  Cosmic galaxy": "a cosmic galaxy theme — deep space-purple with tiny star speckles and glowing cyan accents",
    "🖤  Dark emo": "a dark emo theme — near-black base with glowing magenta and purple accents, moody",
    "🌊  Deep ocean": "a deep ocean theme — teal and aqua, blue-green ramp, soft glow",
}

CHECKER_LIGHT = (235, 235, 235)
CHECKER_DARK = (205, 205, 205)


class EditRequest(BaseModel):
    image: str
    prompt: str = ""
    mode: str = "exact"
    theme: str | None = None


def _checkerboard(width: int, height: int) -> Image.Image:
    bg = Image.new("RGB", (width, height), CHECKER_LIGHT)
    px = bg.load()
    for y in range(height):
        for x in range(width):
            if ((x // 16) + (y // 16)) % 2:
                px[x, y] = CHECKER_DARK
    return bg


def _file_b64(path: Path) -> str:
    return "data:image/png;base64," + base64.b64encode(path.read_bytes()).decode()


def _image_b64(image: Image.Image) -> str:
    out = io.BytesIO()
    image.save(out, format="PNG")
    return "data:image/png;base64," + base64.b64encode(out.getvalue()).decode()


def _decode_image_b64(image_b64: str) -> tuple[Path, str | None]:
    raw = image_b64.split(",", 1)[1] if "," in image_b64 else image_b64
    image = Image.open(io.BytesIO(base64.b64decode(raw))).convert("RGBA")
    original_w, original_h = image.size
    note = None
    if max(original_w, original_h) > MAX_DIM:
        scale = MAX_DIM / max(original_w, original_h)
        new_w = max(1, round(original_w * scale))
        new_h = max(1, round(original_h * scale))
        image = image.resize((new_w, new_h), Image.Resampling.NEAREST)
        note = "downscaled {}x{} -> {}x{}".format(
            original_w, original_h, new_w, new_h
        )
    tmp_in = Path(tempfile.gettempdir()) / "pixellock_in.png"
    image.save(tmp_in)
    return tmp_in, note


def _render_crisp(sprite: validate.Sprite, target: int = 512) -> Image.Image:
    scale = max(1, target // max(sprite.width, sprite.height))
    out_w, out_h = sprite.width * scale, sprite.height * scale
    sprite_image = Image.new("RGBA", (sprite.width, sprite.height))
    sprite_image.putdata(
        [
            (0, 0, 0, 0)
            if sprite.palette[ch] is None
            else (*sprite.palette[ch], 255)
            for row in sprite.rows
            for ch in row
        ]
    )
    sprite_image = sprite_image.resize((out_w, out_h), Image.Resampling.NEAREST)
    bg = _checkerboard(out_w, out_h).convert("RGBA")
    bg.alpha_composite(sprite_image)
    return bg.convert("RGB")


def _build_user_message(wire: str, width: int, height: int, instruction: str, upscale: bool) -> tuple[str, int]:
    if upscale:
        contract = (
            "Redraw this texture at {}x{} (2x). Every input pixel becomes a 2x2 "
            "block: transparent stays transparent, colored stays colored. Add "
            "finer shading and detail within that constraint."
        ).format(width * 2, height * 2)
        out_cells = (width * 2) * (height * 2)
    else:
        contract = (
            "Edit this texture. The grid stays {}x{} and every transparent cell "
            "stays transparent; change only the colors of non-transparent cells."
        ).format(width, height)
        out_cells = width * height

    return (
        "{}\n\nInstruction: {}\n\nHere is the input texture:\n{}".format(
            contract, instruction, wire
        ),
        out_cells,
    )


def _run_engine(image_path: Path, instruction: str, upscale: bool, note: str | None) -> dict[str, Any]:
    wire, width, height = pe.png_to_wire(image_path, spaced=True)
    input_sprite, parse_error = validate.parse_sprite(wire)
    if input_sprite is None:
        raise ValueError("Could not read sprite: {}".format(parse_error))

    opaque_key_count = len([key for key in input_sprite.palette if key != "."])
    grammar = pe.build_grammar(input_sprite.rows, opaque_key_count, upscale, spaced=True)
    user_msg, out_cells = _build_user_message(wire, width, height, instruction, upscale)
    payload = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": pe.APP_SYSTEM},
            {"role": "user", "content": user_msg},
        ],
        "max_tokens": min(int(out_cells * 1.8) + 800, 40000),
        "temperature": 0.7,
        "chat_template_kwargs": {"enable_thinking": False},
        "grammar": grammar,
    }

    started = time.perf_counter()
    try:
        response = httpx.post(ENDPOINT, json=payload, timeout=600.0)
        response.raise_for_status()
    except Exception as exc:
        raise RuntimeError(
            "Model backend is still waking or unreachable: {}".format(exc)
        ) from exc
    latency = time.perf_counter() - started

    body = response.json()
    text = body["choices"][0]["message"]["content"] or ""
    output_sprite, output_error = validate.parse_sprite(text)
    if output_sprite is None:
        raise ValueError("Model output failed to parse: {}".format(output_error))

    input_footprint = validate.footprint(input_sprite)
    if upscale:
        input_footprint = {
            (2 * x + dx, 2 * y + dy)
            for (x, y) in input_footprint
            for dx in (0, 1)
            for dy in (0, 1)
        }
    footprint_ok = input_footprint == validate.footprint(output_sprite)
    color_count = len([key for key in output_sprite.palette if key != "."])

    true_png = Path(tempfile.gettempdir()) / "pixellock_out.png"
    pe.wire_to_png(output_sprite, true_png)

    status_bits = [
        "{}x{}".format(output_sprite.width, output_sprite.height),
        "{} colors".format(color_count),
        "{:.1f}s".format(latency),
    ]
    if note:
        status_bits.insert(0, note)

    return {
        "ok": True,
        "status": "  ·  ".join(status_bits),
        "footprint_ok": footprint_ok,
        "footprint_perfect": footprint_ok,
        "width": output_sprite.width,
        "height": output_sprite.height,
        "colors": color_count,
        "latency": round(latency, 1),
        "image": _file_b64(true_png),
        "input": _file_b64(image_path),
        "preview": _image_b64(_render_crisp(output_sprite)),
        "wire": text.strip(),
        "tokens": (body.get("usage") or {}).get("completion_tokens"),
    }


def _asset_rows() -> list[dict[str, str]]:
    rows = []
    for path in sorted(ASSETS_DIR.glob("*.png")):
        rows.append(
            {
                "id": path.stem,
                "title": path.stem.replace("gen_", "").replace("_", " ").title(),
                "input": _file_b64(path),
                "prompt": "",
                "mode": "exact",
            }
        )
    return rows


def _example_rows() -> list[dict[str, str]]:
    manifest = EXAMPLES_DIR / "examples.json"
    if not manifest.exists():
        return []
    rows = []
    for item in json.loads(manifest.read_text(encoding="utf-8")):
        input_path = EXAMPLES_DIR / item["input"]
        output_path = EXAMPLES_DIR / item["output"]
        if not input_path.exists() or not output_path.exists():
            continue
        rows.append(
            {
                "id": str(item.get("id", input_path.stem)),
                "title": item["title"],
                "theme": item.get("theme", ""),
                "prompt": item["prompt"],
                "mode": item.get("mode", "exact"),
                "input": _file_b64(input_path),
                "output": _file_b64(output_path),
            }
        )
    return rows


api = FastAPI(title="PixelLock")


@api.get("/", response_class=HTMLResponse)
async def index() -> str:
    return (STATIC_DIR / "index.html").read_text(encoding="utf-8")


@api.get("/api/health")
async def api_health() -> JSONResponse:
    base = ENDPOINT.rsplit("/v1/", 1)[0] if "/v1/" in ENDPOINT else ENDPOINT
    model_online = False
    detail = ""
    try:
        with httpx.Client(timeout=2.0) as client:
            resp = client.get(base.rstrip("/") + "/v1/models")
            model_online = resp.status_code < 500
    except Exception as exc:
        detail = str(exc)
    return JSONResponse(
        {
            "ok": True,
            "model_online": model_online,
            "endpoint": ENDPOINT,
            "model": MODEL,
            "detail": detail,
        }
    )


@api.get("/api/themes")
async def api_themes() -> JSONResponse:
    return JSONResponse(
        {"themes": [{"key": key, "prompt": prompt} for key, prompt in THEMES.items()]}
    )


@api.get("/api/assets")
async def api_assets() -> JSONResponse:
    return JSONResponse(_asset_rows())


@api.get("/api/examples")
async def api_examples() -> JSONResponse:
    return JSONResponse(_example_rows())


@api.post("/api/edit")
async def api_edit(req: EditRequest) -> JSONResponse:
    instruction = (req.prompt or "").strip()
    if not instruction and req.theme:
        instruction = THEMES.get(req.theme, "").strip()
    if not instruction:
        return JSONResponse(
            {"ok": False, "status": "Pick a theme or type a prompt."},
            status_code=400,
        )

    try:
        image_path, note = _decode_image_b64(req.image)
        result = _run_engine(
            image_path,
            instruction,
            req.mode == "upscale2x" or req.mode.lower().startswith("upscale"),
            note,
        )
    except ValueError as exc:
        return JSONResponse({"ok": False, "status": str(exc)}, status_code=400)
    except RuntimeError as exc:
        return JSONResponse({"ok": False, "status": str(exc)}, status_code=503)
    return JSONResponse(result)


api.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


GRADIO_CSS = """
footer { display: none !important; }
.gradio-container { max-width: 860px !important; }
"""

with gr.Blocks(
    title="PixelLock Gradio Runtime",
    theme=gr.themes.Base(primary_hue="purple", neutral_hue="slate"),
    css=GRADIO_CSS,
) as gradio_demo:
    gr.Markdown(
        """
        # PixelLock Gradio Runtime

        The custom PixelLock UI is served at `/`. This mounted Gradio surface is
        kept live for Space eligibility and API introspection. It uses the same
        backend model, grammar-constrained decoding, and footprint checks.
        """
    )
    gr.JSON(
        value={
            "custom_ui": "/",
            "edit_api": "/api/edit",
            "model": MODEL,
            "grammar_locked": True,
            "thinking_disabled": True,
        },
        label="Runtime contract",
    )


app = gr.mount_gradio_app(api, gradio_demo, path="/gradio")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))