modal_workers/klein_image.py

"""Modal worker for black-forest-labs/FLUX.2-klein-4B image generation.

Deploy:
    modal deploy modal_workers/klein_image.py

The deployed ASGI app exposes one base URL with:
    POST /generate - accepts {"prompt": str, "seed": int | None}
    GET  /health - reports worker readiness
    GET  /media/{filename} - serves generated PNG files

Tiny Narrator's app.py calls these routes through KLEIN_MODAL_ENDPOINT.
"""

import io
import os
from pathlib import Path
from uuid import uuid4

import modal


APP_NAME = "tiny-narrator-klein"
CACHE_DIR = Path("/cache")
MEDIA_DIR = Path("/outputs")
IMAGE_MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"

app = modal.App(APP_NAME)
model_cache = modal.Volume.from_name("tiny-narrator-klein-cache", create_if_missing=True)
output_volume = modal.Volume.from_name("tiny-narrator-klein-outputs", create_if_missing=True)

klein_image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("git")
    .pip_install(
        "torch==2.9.1",
        index_url="https://download.pytorch.org/whl/cu128",
    )
    .pip_install(
        "accelerate==1.12.0",
        "diffusers @ git+https://github.com/huggingface/diffusers.git",
        "fastapi[standard]==0.136.3",
        "huggingface-hub[hf-transfer]==0.36.0",
        "pillow==12.2.0",
        "safetensors>=0.8.0rc0",
        "sentencepiece==0.2.1",
        "transformers==4.57.3",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)


def _secret_names() -> list[modal.Secret]:
    """Return the fixed secrets used by Modal deployments.

    Modal requires dependencies to be identical between local deploy-time import
    and remote container import. Keep this list static; do not derive it from
    environment variables.
    """

    return [modal.Secret.from_name("tiny-narrator-klein-token")]


async def _media_file_response(filename: str):
    from fastapi import HTTPException
    from fastapi.responses import Response

    path = MEDIA_DIR / filename
    if not path.exists() or not path.is_file():
        try:
            await output_volume.reload.aio()
        except RuntimeError as exc:
            if not path.exists() or not path.is_file():
                raise HTTPException(status_code=503, detail=f"Media volume reload failed: {exc}") from exc

    if not path.exists() or not path.is_file():
        raise HTTPException(status_code=404, detail=f"Media file not found: {filename}")

    content_type = "image/png"
    if filename.endswith(".webp"):
        content_type = "image/webp"
    elif filename.endswith((".jpg", ".jpeg")):
        content_type = "image/jpeg"
    return Response(content=path.read_bytes(), media_type=content_type)


@app.cls(
    image=klein_image,
    gpu=os.getenv("KLEIN_MODAL_GPU", "A10G"),
    volumes={str(CACHE_DIR): model_cache, str(MEDIA_DIR): output_volume},
    secrets=_secret_names(),
    timeout=900,
    scaledown_window=300,
    max_containers=1,
)
class KleinImageModel:
    @modal.enter()
    def load(self) -> None:
        import torch
        from diffusers import Flux2KleinPipeline

        self.model_id = os.environ.get("KLEIN_IMAGE_MODEL_ID", IMAGE_MODEL_ID)
        token = os.environ.get("HF_TOKEN")
        self.pipe = Flux2KleinPipeline.from_pretrained(
            self.model_id,
            cache_dir=str(CACHE_DIR / "huggingface"),
            torch_dtype=torch.bfloat16,
            token=token,
        )
        self.pipe.to("cuda")

    @modal.method()
    def generate(self, prompt: str, seed: int | None = None) -> dict:
        import torch

        MEDIA_DIR.mkdir(parents=True, exist_ok=True)
        generator = None
        if seed is not None:
            generator = torch.Generator(device="cuda").manual_seed(seed)

        result = self.pipe(
            prompt=prompt,
            num_inference_steps=int(os.environ.get("KLEIN_MODAL_STEPS", "4")),
            guidance_scale=float(os.environ.get("KLEIN_MODAL_GUIDANCE", "1.0")),
            generator=generator,
        )
        image = result.images[0]

        filename = f"klein-{uuid4().hex}.png"
        output_path = MEDIA_DIR / filename
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        output_path.write_bytes(buffer.getvalue())
        output_volume.commit()

        return {
            "filename": filename,
            "model": self.model_id,
        }


@app.function(
    image=klein_image,
    volumes={str(MEDIA_DIR): output_volume},
    secrets=_secret_names(),
    timeout=300,
)
@modal.concurrent(max_inputs=20)
@modal.asgi_app(label="tiny-narrator-klein")
def klein_api():
    from fastapi import FastAPI, HTTPException, Request

    import time

    api = FastAPI(title="Tiny Narrator Klein Worker")
    model = KleinImageModel()

    def _check_token(request: Request) -> None:
        """Reject the request if a token is configured but not provided or mismatched."""
        expected = os.getenv("KLEIN_MODAL_TOKEN", "")
        if not expected:
            return
        auth_header = request.headers.get("authorization", "")
        token_header = request.headers.get("x-tiny-narrator-token", "")
        provided = ""
        if auth_header.startswith("Bearer "):
            provided = auth_header[len("Bearer "):]
        elif token_header:
            provided = token_header
        if provided != expected:
            raise HTTPException(status_code=401, detail="Unauthorized")

    @api.get("/health")
    async def health(request: Request) -> dict:
        _check_token(request)
        return {
            "ok": True,
            "model": IMAGE_MODEL_ID,
            "runtime": "modal-klein",
        }

    @api.get("/media/{filename}")
    async def media(filename: str):
        return await _media_file_response(filename)

    @api.post("/generate")
    async def generate(request: Request) -> dict:
        _check_token(request)
        start = time.perf_counter()
        body = await request.json()
        prompt = str(body.get("prompt") or "").strip()
        if not prompt:
            raise HTTPException(status_code=400, detail="prompt is required")

        seed = body.get("seed")
        if seed is not None:
            try:
                seed = int(seed)
            except (TypeError, ValueError) as exc:
                raise HTTPException(status_code=400, detail="seed must be an integer") from exc

        result = await model.generate.remote.aio(prompt, seed)
        model_id = result.get("model") or IMAGE_MODEL_ID
        if model_id != IMAGE_MODEL_ID:
            raise HTTPException(status_code=500, detail=f"unexpected model: {model_id}")

        filename = result["filename"]
        image_url = str(request.url_for("media", filename=filename))
        return {
            "ok": True,
            "runtime": "modal-klein",
            "model": IMAGE_MODEL_ID,
            "image_url": image_url,
            "prompt": prompt,
            "seed": seed,
            "elapsed_ms": round((time.perf_counter() - start) * 1000),
        }

    return api