server.py

import base64
import io
import json
from pathlib import Path
from urllib.parse import quote

from fastapi import FastAPI
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from pipeline import (
    generate, NUM_STEPS, MAX_STEPS, GUIDANCE_SCALE, GUIDANCE_RESCALE,
    DEFAULT_SPACING,
)

SPACINGS = ("trailing", "quadratic")

BASE       = Path(__file__).parent
SPRITE_DIR = BASE / "sprites"
STATIC_DIR = BASE / "static"

app = FastAPI(title="Senior LDM")


class CachedStaticFiles(StaticFiles):
    """StaticFiles with long-lived cache headers. The sprite animation swaps the
    <img> src ~10×/sec; without caching the browser re-downloads every frame
    forever. Sprites/gallery are immutable assets, so cache them aggressively."""

    async def get_response(self, path, scope):
        response = await super().get_response(path, scope)
        response.headers["Cache-Control"] = "public, max-age=31536000, immutable"
        return response


app.mount("/sprites", CachedStaticFiles(directory=str(SPRITE_DIR)), name="sprites")
app.mount("/static",  CachedStaticFiles(directory=str(STATIC_DIR)), name="static")

# ── sprite state machine → frame files ──────────────────────────────────────
SPRITE_MAP = {
    "idle": ["Screenshot 2026-05-31 020023.png", "Screenshot 2026-05-31 020056.png"],
    "run":  ["Screenshot 2026-05-31 020138.png", "Screenshot 2026-05-31 020141.png",
             "Screenshot 2026-05-31 020145.png", "Screenshot 2026-05-31 020148.png"],
    "turn": ["Screenshot 2026-05-31 020210.png", "Screenshot 2026-05-31 020215.png"],
    "done": ["Screenshot 2026-05-31 020023.png"],
}

# ── prompt ideas — art / LAION-aesthetics styles, no humans ─────────────────
PROMPT_IDEAS = [
    "a serene mountain lake at sunset, oil painting",
    "venice at night, shimmering canals, painterly",
    "a classical still life of fruit and wine",
    "a fantasy castle on a cliff, hd wallpaper",
    "a brown mountain under a starry night sky",
    "a forest at sunset with dramatic clouds",
    "a winter cabin at the water's edge, wildlife art",
    "a lighthouse on a rocky coast during a storm",
    "a field of lavender under a stormy sky",
    "an autumn forest path of golden leaves",
    "a watercolor painting of a snowy village",
    "a futuristic city skyline glowing at night",
    "a tranquil japanese garden with a koi pond",
    "an abstract swirling galaxy of vibrant color",
    "a cozy cabin interior with warm firelight",
    "a snowy forest path in soft morning light",
]

# ── gallery showcase — example outputs served from static/gallery ───────────
def _gallery():
    gdir = STATIC_DIR / "gallery"
    items = []
    for f in sorted(gdir.glob("*.png")):
        # filename → readable label: drop the source/index prefix, de-underscore.
        parts = f.stem.split("_")
        label = " ".join(parts[2:] if len(parts) > 2 else parts)
        items.append({"url": f"/static/gallery/{quote(f.name)}", "label": label})
    return items

GALLERY = _gallery()


class GenRequest(BaseModel):
    prompt: str
    seed: int = 42
    steps: int = NUM_STEPS
    guidance: float = GUIDANCE_SCALE
    guidance_rescale: float = GUIDANCE_RESCALE
    spacing: str = DEFAULT_SPACING


@app.get("/")
def index():
    return FileResponse(str(STATIC_DIR / "index.html"))


@app.get("/api/config")
def config():
    return {
        "sprites":   {state: [f"/sprites/{quote(f)}" for f in files]
                      for state, files in SPRITE_MAP.items()},
        "prompts":   PROMPT_IDEAS,
        "gallery":   GALLERY,
        "num_steps": NUM_STEPS,
        "max_steps": MAX_STEPS,
        "guidance":  GUIDANCE_SCALE,
        "guidance_rescale": GUIDANCE_RESCALE,
        "spacing":   DEFAULT_SPACING,
        "spacings":  list(SPACINGS),
    }


@app.post("/api/generate")
def api_generate(req: GenRequest):
    prompt    = req.prompt.strip()
    steps     = max(1, min(MAX_STEPS, req.steps))
    guidance  = max(0.0, min(20.0, req.guidance))
    g_rescale = max(0.0, min(1.0, req.guidance_rescale))
    spacing   = req.spacing if req.spacing in SPACINGS else DEFAULT_SPACING

    def stream():
        if not prompt:
            yield json.dumps({"type": "error", "message": "empty prompt"}) + "\n"
            return
        for step, total, elapsed, image in generate(
            prompt, seed=req.seed, steps=steps, guidance=guidance,
            spacing=spacing, guidance_rescale=g_rescale,
        ):
            if image is None:
                yield json.dumps({
                    "type": "step", "step": step, "total": total,
                    "elapsed": round(elapsed, 2),
                }) + "\n"
            else:
                buf = io.BytesIO()
                image.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode()
                yield json.dumps({
                    "type": "done", "step": step, "total": total,
                    "elapsed": round(elapsed, 2),
                    "image": "data:image/png;base64," + b64,
                }) + "\n"

    return StreamingResponse(stream(), media_type="application/x-ndjson")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)