app/api.py

"""FastAPI service exposing the memorization guardrail.

Runs the canonical MemGuard flow -- ``prompt -> Stable Diffusion ->
memorization score`` -- with the real anisotropy metric. The Stable Diffusion
pipeline is loaded once, lazily, on the first request.

Run locally:  uvicorn app.api:app --reload
Requires the ``api`` + ``diffusers`` extras (the metric needs a real model):
    pip install ".[api,diffusers]"
First scored request downloads SD v1-4 (~4 GB) and is slow on CPU.
"""

from __future__ import annotations

from fastapi import FastAPI
from pydantic import BaseModel

from memguard.detector import MemorizationDetector
from memguard.pipeline import DEFAULT_MODEL_ID, load_sd_pipeline

app = FastAPI(title="MemGuard", version="0.1.0")
detector = MemorizationDetector()

_pipe = None


def _get_pipe():
    """Load Stable Diffusion once (lazily) so startup/health stay cheap."""
    global _pipe
    if _pipe is None:
        _pipe = load_sd_pipeline(DEFAULT_MODEL_ID)
    return _pipe


class ScoreRequest(BaseModel):
    prompt: str
    seed: int | None = None


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "version": "0.1.0", "model": DEFAULT_MODEL_ID}


@app.post("/score")
def score(req: ScoreRequest) -> dict:
    """Score a prompt for training-data memorization on Stable Diffusion v1-4."""
    kwargs = {}
    if req.seed is not None:
        import torch

        kwargs["generator"] = torch.Generator().manual_seed(int(req.seed))
    return detector.check(_get_pipe(), req.prompt, **kwargs)