server.py

"""FastAPI server — DiffuMT project page + 2AFC game + DDIM generation."""
from __future__ import annotations

import asyncio
import base64
import io
import json
import os
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Optional

import numpy as np
import torch
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.responses import Response, StreamingResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
from pydantic import BaseModel

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IMAGE_SIZE = 256
IMAGES_DIR = Path(__file__).parent / "images"
ADMIN_TOKEN = os.environ.get("ADMIN_TOKEN", "")

def _resolve_data_file() -> Path:
    """
    Prefer /data/results.json (HF persistent storage) if /data is a writable mount.
    Fall back to /app/results.json (committed to repo, survives restarts within the
    same build, resets on the next Docker rebuild).
    Seeds /data from the repo copy on first boot if persistent storage is now attached.
    """
    explicit = os.environ.get("DATA_PATH")
    if explicit:
        return Path(explicit)
    persistent = Path("/data/results.json")
    bundled    = Path(__file__).parent / "results.json"
    # /data is available if the directory exists and is writable
    if persistent.parent.exists() and os.access(persistent.parent, os.W_OK):
        if not persistent.exists() and bundled.exists():
            import shutil
            shutil.copy(bundled, persistent)
        return persistent
    return bundled

DATA_FILE = _resolve_data_file()
MODEL_CACHE = Path("/data/model") if Path("/data").exists() and os.access("/data", os.W_OK) \
              else Path(__file__).parent / ".model_cache"

model_state: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    print("lifespan: importing diffusers/huggingface_hub", flush=True)
    from diffusers import DDIMScheduler, UNet2DModel
    from huggingface_hub import snapshot_download
    print("lifespan: imports done, downloading model snapshot", flush=True)

    MODEL_CACHE.mkdir(parents=True, exist_ok=True)
    model_dir = snapshot_download(
        "HTW-KI-Werkstatt/DiffuMT",
        repo_type="model",
        local_dir=str(MODEL_CACHE),
    )
    print(f"lifespan: snapshot at {model_dir}, loading UNet", flush=True)
    unet = (
        UNet2DModel.from_pretrained(f"{model_dir}/unet", use_safetensors=True)
        .to(DEVICE)
        .eval()
    )
    print("lifespan: UNet loaded, loading scheduler", flush=True)
    scheduler = DDIMScheduler.from_pretrained(f"{model_dir}/scheduler")
    model_state["unet"] = unet
    model_state["scheduler"] = scheduler
    print(f"Model loaded on {DEVICE}", flush=True)
    yield


# ── Data store ───────────────────────────────────────────────────────────────

DATA_FILE.parent.mkdir(parents=True, exist_ok=True)
if not DATA_FILE.exists():
    DATA_FILE.write_text(json.dumps({"sessions": [], "trials": []}, indent=2))


def _read() -> dict:
    return json.loads(DATA_FILE.read_text())


def _write(data: dict) -> None:
    DATA_FILE.write_text(json.dumps(data, indent=2))


# ── App ──────────────────────────────────────────────────────────────────────

app = FastAPI(lifespan=lifespan)

IMAGE_EXT = {".jpg", ".jpeg", ".png", ".tiff", ".tif", ".webp", ".bmp", ".gif"}


# ── Game API ─────────────────────────────────────────────────────────────────

@app.get("/api/images")
def list_images():
    def ls(subdir: str):
        d = IMAGES_DIR / subdir
        if not d.exists():
            return []
        return [
            {"id": f.name, "src": f"/images/{subdir}/{f.name}", "type": subdir}
            for f in sorted(d.iterdir())
            if f.suffix.lower() in IMAGE_EXT
        ]
    return {"real": ls("real"), "fake": ls("fake")}


class TrialPayload(BaseModel):
    sessionId: str
    trial: dict


@app.post("/api/trial")
def save_trial(payload: TrialPayload):
    data = _read()
    data["trials"].append({"session_id": payload.sessionId, "player_name": None, **payload.trial})
    _write(data)
    return {"ok": True}


class SubmitPayload(BaseModel):
    sessionId: str
    playerName: str


@app.post("/api/submit")
def submit_score(payload: SubmitPayload):
    data = _read()
    trials = [t for t in data["trials"] if t["session_id"] == payload.sessionId]
    if not trials:
        raise HTTPException(status_code=404, detail="Session not found")
    data["trials"] = [
        {**t, "player_name": payload.playerName} if t["session_id"] == payload.sessionId else t
        for t in data["trials"]
    ]
    score = sum(1 for t in trials if t.get("correct"))
    total = len(trials)
    data["sessions"].append({
        "session_id": payload.sessionId,
        "player_name": payload.playerName,
        "timestamp": trials[0].get("timestamp", ""),
        "score": score,
        "total": total,
        "percentage": round((score / total) * 100) if total else 0,
    })
    _write(data)
    return {"sessionId": payload.sessionId, "score": score, "total": total}


@app.get("/api/leaderboard")
def leaderboard():
    sessions = _read()["sessions"]
    return sorted(sessions, key=lambda s: (-s["percentage"], s["timestamp"]))[:50]


def _check_admin(x_admin_token: str, authorization: str):
    if not ADMIN_TOKEN:
        raise HTTPException(status_code=503, detail="Admin access disabled: ADMIN_TOKEN not set")
    bearer = authorization.removeprefix("Bearer ").strip()
    if bearer != ADMIN_TOKEN and x_admin_token != ADMIN_TOKEN:
        raise HTTPException(status_code=401, detail="Unauthorized")


@app.get("/api/admin/sessions")
def admin_sessions():
    sessions = _read()["sessions"]
    return sorted(sessions, key=lambda s: s["timestamp"], reverse=True)


@app.get("/api/export")
def export_data(x_admin_token: str = Header(default=""), authorization: str = Header(default="")):
    _check_admin(x_admin_token, authorization)
    return Response(
        content=DATA_FILE.read_text(),
        media_type="application/json",
        headers={"Content-Disposition": 'attachment; filename="2afc_data.json"'},
    )


@app.post("/api/admin/clear")
def clear_all(x_admin_token: str = Header(default=""), authorization: str = Header(default="")):
    _check_admin(x_admin_token, authorization)
    _write({"sessions": [], "trials": []})
    return {"ok": True}


class ClearUserPayload(BaseModel):
    playerName: str


@app.post("/api/admin/clear-user")
def clear_user(payload: ClearUserPayload, x_admin_token: str = Header(default=""), authorization: str = Header(default="")):
    _check_admin(x_admin_token, authorization)
    data = _read()
    bs, bt = len(data["sessions"]), len(data["trials"])
    data["sessions"] = [s for s in data["sessions"] if s["player_name"] != payload.playerName]
    data["trials"]   = [t for t in data["trials"]   if t.get("player_name") != payload.playerName]
    _write(data)
    return {"ok": True, "removedSessions": bs - len(data["sessions"]), "removedTrials": bt - len(data["trials"])}


class ClearSessionPayload(BaseModel):
    sessionId: str
    deleteTrials: bool = True


@app.post("/api/admin/clear-session")
def clear_session(payload: ClearSessionPayload, x_admin_token: str = Header(default=""), authorization: str = Header(default="")):
    _check_admin(x_admin_token, authorization)
    data = _read()
    bs, bt = len(data["sessions"]), len(data["trials"])
    data["sessions"] = [s for s in data["sessions"] if s["session_id"] != payload.sessionId]
    if payload.deleteTrials:
        data["trials"] = [t for t in data["trials"] if t.get("session_id") != payload.sessionId]
    _write(data)
    return {"ok": True, "removedSessions": bs - len(data["sessions"]), "removedTrials": bt - len(data["trials"])}


@app.post("/api/admin/import")
async def import_data(request: Request, x_admin_token: str = Header(default=""), authorization: str = Header(default="")):
    _check_admin(x_admin_token, authorization)
    body = await request.json()
    if "sessions" not in body or "trials" not in body:
        raise HTTPException(status_code=400, detail="Body must have 'sessions' and 'trials' keys")
    _write(body)
    return {"ok": True, "sessions": len(body["sessions"]), "trials": len(body["trials"])}


# ── DDIM generation (SSE) ─────────────────────────────────────────────────────

class GeneratePayload(BaseModel):
    mask: str       # base64 PNG from canvas.toDataURL()
    steps: int = 30
    seed: int = -1


def _prep_mask(b64: str) -> torch.Tensor:
    raw = base64.b64decode(b64)
    img = Image.open(io.BytesIO(raw)).convert("L")
    img = img.resize((IMAGE_SIZE, IMAGE_SIZE), Image.NEAREST)
    arr = (np.asarray(img) > 0).astype(np.float32) / 255.0
    return torch.from_numpy(arr)[None, None].to(DEVICE)


def _tensor_to_b64(x: torch.Tensor) -> str:
    arr = (x / 2 + 0.5).clamp(0, 1)
    arr = (arr[0].cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    buf = io.BytesIO()
    Image.fromarray(arr).save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


@app.post("/api/generate")
async def generate(payload: GeneratePayload):
    unet = model_state.get("unet")
    scheduler = model_state.get("scheduler")
    if unet is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet — please wait a moment")

    async def stream():
        try:
            mask = _prep_mask(payload.mask)
            if mask.sum() == 0:
                yield f"data: {json.dumps({'error': 'Empty mask — draw some strokes first'})}\n\n"
                return

            gen = None
            if payload.seed >= 0:
                gen = torch.Generator(device=DEVICE).manual_seed(payload.seed)

            x = torch.randn((1, 3, IMAGE_SIZE, IMAGE_SIZE), generator=gen, device=DEVICE)
            scheduler.set_timesteps(payload.steps)
            timesteps = scheduler.timesteps
            n = len(timesteps)
            snap = set(np.linspace(0, n - 1, min(8, n)).round().astype(int).tolist())

            with torch.no_grad():
                for i, t in enumerate(timesteps):
                    x_in = torch.cat([x, mask], dim=1)
                    eps = unet(x_in, t).sample
                    x = scheduler.step(eps, t, x, generator=gen).prev_sample

                    if i in snap or i == n - 1:
                        msg = json.dumps({
                            "step": i + 1,
                            "total": n,
                            "image": _tensor_to_b64(x),
                            "done": i == n - 1,
                        })
                        yield f"data: {msg}\n\n"
                        await asyncio.sleep(0)

        except Exception as exc:
            yield f"data: {json.dumps({'error': str(exc)})}\n\n"

    return StreamingResponse(
        stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


# ── Static files (must be last) ───────────────────────────────────────────────

app.mount("/images", StaticFiles(directory=str(IMAGES_DIR)), name="images")
app.mount("/", StaticFiles(directory=str(Path(__file__).parent / "public"), html=True), name="static")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))