app.py

"""Hugging Face Space entry point for the trusted HTML shell."""

from __future__ import annotations

import base64
import binascii
import threading
import time
import warnings
from io import BytesIO
from pathlib import Path
from typing import Any, Callable

from fastapi import HTTPException, Request
from fastapi.responses import HTMLResponse
from fastapi.responses import JSONResponse
from PIL import Image, UnidentifiedImageError

from snap2sim.backend import InferenceClient, Settings
from snap2sim.schema import normalize_confidence_threshold

try:
    from gradio import Server
except ImportError:
    from fastapi import FastAPI
    import uvicorn

    class Server(FastAPI):  # type: ignore[no-redef]
        """Local compatibility shim for environments older than Gradio Server."""

        def api(self, name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
            def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
                return func

            return decorator

        def launch(self, **kwargs: Any) -> None:
            uvicorn.run(
                self,
                host=kwargs.get("server_name", "0.0.0.0"),
                port=kwargs.get("server_port", 7860),
            )


app = Server()
INDEX_PATH = Path(__file__).with_name("index.html")
MAX_IMAGE_BASE64_CHARS = 12 * 1024 * 1024
MAX_IMAGE_BYTES = 9 * 1024 * 1024
MAX_IMAGE_PIXELS = 12_000_000
RATE_LIMIT_WINDOW_SECONDS = 60
RATE_LIMIT_PER_CLIENT = 12
RATE_LIMIT_GLOBAL = 72
RATE_LIMIT_PATHS = {"/analyze_image", "/generate_scene"}

Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS

_rate_lock = threading.Lock()
_client_hits: dict[str, list[float]] = {}
_global_hits: list[float] = []


@app.middleware("http")
async def rate_limit_api(request: Request, call_next: Callable[..., Any]) -> Any:
    if request.method == "POST" and request.url.path in RATE_LIMIT_PATHS:
        allowed, retry_after = _record_request(_client_id(request))
        if not allowed:
            return JSONResponse(
                {"detail": f"Rate limit exceeded. Retry after {retry_after} seconds."},
                status_code=429,
                headers={"Retry-After": str(retry_after)},
            )
    return await call_next(request)


@app.get("/", response_class=HTMLResponse)
async def homepage() -> str:
    return INDEX_PATH.read_text(encoding="utf-8")


@app.get("/manifest.json")
async def manifest() -> dict[str, Any]:
    return {
        "name": "Snap2Sim Inside the Machine",
        "short_name": "Snap2Sim",
        "start_url": "/",
        "display": "standalone",
        "background_color": "#0F1318",
        "theme_color": "#E8A33D",
    }


@app.api(name="analyze_image")
def analyze_image_api(image_base64: str) -> dict[str, Any]:
    return _analyze_image(image_base64)


@app.post("/analyze_image")
def analyze_image_http(payload: dict[str, Any]) -> dict[str, Any]:
    return _analyze_image(str(payload.get("image_base64", "")))


@app.api(name="generate_scene")
def generate_scene_api(
    analysis: dict[str, Any],
    confidence_threshold: float | None = None,
) -> dict[str, Any]:
    return _generate_scene(analysis, confidence_threshold)


@app.post("/generate_scene")
def generate_scene_http(payload: dict[str, Any]) -> dict[str, Any]:
    return _generate_scene(
        payload.get("analysis") or {},
        payload.get("confidence_threshold"),
    )


def _analyze_image(image_base64: str) -> dict[str, Any]:
    image = _decode_image(image_base64) if image_base64 else None
    return InferenceClient(Settings()).analyze_image(image)


def _generate_scene(analysis: dict[str, Any], threshold: Any = None) -> dict[str, Any]:
    return InferenceClient(Settings()).generate_scene(
        analysis,
        normalize_confidence_threshold(threshold),
    )


def _decode_image(image_base64: str) -> Image.Image:
    if "," in image_base64 and image_base64.lstrip().startswith("data:"):
        image_base64 = image_base64.split(",", 1)[1]
    if len(image_base64) > MAX_IMAGE_BASE64_CHARS:
        raise HTTPException(status_code=413, detail="Image upload is too large.")
    try:
        raw = base64.b64decode(image_base64, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise HTTPException(status_code=400, detail="Image payload is not valid base64.") from exc
    if len(raw) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=413, detail="Image upload is too large.")

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("error", Image.DecompressionBombWarning)
            image = Image.open(BytesIO(raw))
            image.load()
    except Image.DecompressionBombWarning as exc:
        raise HTTPException(status_code=413, detail="Image dimensions are too large.") from exc
    except Image.DecompressionBombError as exc:
        raise HTTPException(status_code=413, detail="Image dimensions are too large.") from exc
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise HTTPException(status_code=400, detail="Upload a valid image file.") from exc

    if image.width * image.height > MAX_IMAGE_PIXELS:
        raise HTTPException(status_code=413, detail="Image dimensions are too large.")
    return image.convert("RGB")


def _client_id(request: Request) -> str:
    forwarded_for = request.headers.get("x-forwarded-for", "")
    if forwarded_for:
        return forwarded_for.split(",", 1)[0].strip()
    return request.client.host if request.client else "unknown"


def _record_request(client_id: str) -> tuple[bool, int]:
    now = time.monotonic()
    cutoff = now - RATE_LIMIT_WINDOW_SECONDS
    with _rate_lock:
        _global_hits[:] = [hit for hit in _global_hits if hit >= cutoff]
        hits = [hit for hit in _client_hits.get(client_id, []) if hit >= cutoff]
        if len(hits) >= RATE_LIMIT_PER_CLIENT or len(_global_hits) >= RATE_LIMIT_GLOBAL:
            oldest = min(hits[0] if hits else now, _global_hits[0] if _global_hits else now)
            retry_after = max(1, int(RATE_LIMIT_WINDOW_SECONDS - (now - oldest)))
            _client_hits[client_id] = hits
            return False, retry_after
        hits.append(now)
        _global_hits.append(now)
        _client_hits[client_id] = hits
    return True, 0


if __name__ == "__main__":
    app.launch()