app.py

"""
PixelForge AI — generative upscale engine (Hugging Face Space app).

Exposes:
  POST /upscale   — accepts raw image bytes (or multipart "image" file part),
                    returns upscaled image bytes. Query params: scale=2|4.
  GET  /health    — liveness probe; reports the selected engine (no secrets).

Engine selection (env, optional):
  UPSCALE_ENGINE  = "supir" (default, generative restoration) | "sdxl_tile"
                    (lighter generative fallback for small-GPU hardware)
  SPACE_ACCESS_TOKEN = if set, every /upscale request must send
                    `Authorization: Bearer <same value>`. The PixelForge site
                    provider sends its HF_TOKEN here — set this to the same
                    value to lock the endpoint down.

Designed for Hugging Face Spaces on the free ZeroGPU hardware tier (Docker
SDK). On ZeroGPU the GPU is attached per-request; this app wraps GPU work in
`gpu_context()` (the `spaces` package), so it also runs unchanged on a
dedicated GPU Space or a local GPU box.

NOTE: GPU inference cannot be exercised in the team sandbox (no GPU, ~4 GB
RAM). The model paths are written against the verified public APIs of SUPIR
(create_SUPIR_model / batchify_sample from the official codebase) and
diffusers (StableDiffusionXLControlNetPipeline); the README lists exactly what
still needs an HF account + first GPU request to confirm.
"""

from __future__ import annotations

import io
import logging
import os
import time
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request, UploadFile
from fastapi.responses import Response
from PIL import Image

from engines import UpscaleEngine, gpu_context
from engines.supir_engine import SupirEngine
from engines.sdxl_tile_engine import SdxlTileEngine

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
log = logging.getLogger("upscale.app")

ENGINE = os.environ.get("UPSCALE_ENGINE", "supir").strip().lower()
ACCESS_TOKEN = os.environ.get("SPACE_ACCESS_TOKEN", "").strip()
MAX_UPLOAD_BYTES = int(os.environ.get("MAX_UPLOAD_BYTES", str(25 * 1024 * 1024)))

_ENGINES: dict[str, UpscaleEngine] = {
    "supir": SupirEngine(),
    "sdxl_tile": SdxlTileEngine(),
}


def get_engine() -> UpscaleEngine:
    engine = _ENGINES.get(ENGINE)
    if engine is None:
        raise HTTPException(
            status_code=500,
            detail=f"unknown UPSCALE_ENGINE '{ENGINE}' (expected supir or sdxl_tile)",
        )
    return engine


def _check_auth(request: Request) -> None:
    """Optional bearer-token gate — used by the PixelForge site provider."""
    if not ACCESS_TOKEN:
        return
    header = request.headers.get("authorization", "")
    provided = header[7:].strip() if header.lower().startswith("bearer ") else ""
    import hmac

    if not provided or not hmac.compare_digest(provided, ACCESS_TOKEN):
        raise HTTPException(status_code=401, detail="unauthorized")


def _load_image(body: bytes) -> Image.Image:
    try:
        return Image.open(io.BytesIO(body))
    except Exception as exc:  # noqa: BLE001 — surface as 400, not 500
        raise HTTPException(status_code=400, detail=f"not a decodable image: {exc}") from exc


@asynccontextmanager
async def lifespan(_: FastAPI):
    log.info("PixelForge upscale Space starting (engine=%s, gpu=%s)", ENGINE, _gpu_available())
    # Warm-up is optional; SUPIR weights (~9.5 GB) download on first request
    # so the Space itself boots fast on ZeroGPU.
    if os.environ.get("PRELOAD", "0").strip().lower() in ("1", "true", "yes"):
        log.info("PRELOAD=1 — loading %s engine at startup", ENGINE)
        get_engine().load()
    yield
    log.info("PixelForge upscale Space stopped")


app = FastAPI(title="PixelForge AI — generative upscale engine", lifespan=lifespan)


def _gpu_available() -> bool:
    try:
        import torch

        return bool(torch.cuda.is_available())
    except Exception:
        return False


@app.get("/health")
def health():
    return {"status": "ok", "engine": ENGINE, "gpu": _gpu_available()}


@app.post("/upscale")
async def upscale(request: Request, scale: int = 4, file: UploadFile | None = None):
    _check_auth(request)
    if scale not in (2, 4):
        raise HTTPException(status_code=400, detail="scale must be 2 or 4")

    # Accept either a multipart file part named "image" or a raw image body.
    if file is not None:
        body = await file.read()
    else:
        content_type = request.headers.get("content-type", "").lower()
        if not (content_type.startswith("image/") or content_type.startswith("application/octet-stream")):
            raise HTTPException(
                status_code=415,
                detail="send raw image bytes (image/*) or multipart with an 'image' file part",
            )
        body = await request.body()

    if not body:
        raise HTTPException(status_code=400, detail="empty body")
    if len(body) > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail=f"image too large (max {MAX_UPLOAD_BYTES} bytes)")

    engine = get_engine()
    img = _load_image(body)

    t0 = time.time()
    log.info("upscale start: engine=%s scale=%s input=%s", engine.name(), scale, img.size)
    with gpu_context():
        out = engine.upscale(img, scale)
    log.info("upscale done: %s -> %s in %.1fs", img.size, out.size, time.time() - t0)

    buf = io.BytesIO()
    out.save(buf, format="PNG")
    return Response(
        content=buf.getvalue(),
        media_type="image/png",
        headers={"X-Upscale-Engine": engine.name()},
    )

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