images.py

"""FLUX.2 [klein] 4B image pipeline. Pluggable GPU backend.

One model, one `image=` parameter does both jobs:
  - initial generation: image=[user_photo]  (a LIST -> multi-reference; puts the user in the scene)
  - edit (next pass):    image=base_image    (a single image -> instruction-guided edit)

Shared art_style + a fixed per-adventure seed give every quest a cohesive look. Always EDIT the
initial image for success/failure states (consistency) - never text-to-image from scratch.

FROGQUEST_BACKEND selects WHERE inference runs (public functions identical either way):
  - "zerogpu" (default): diffusers inside @spaces.GPU on the HF Space's ZeroGPU.
  - "modal": forward to a deployed Modal class (see modal_app.py); the Space runs on CPU-basic.

IMPORTANT: torch / diffusers are imported LAZILY (inside _get_pipe/_gen), and the weight prefetch
is gated to the zerogpu backend, so importing this module on a CPU-basic Space (modal backend)
drags in nothing heavy and downloads nothing.
"""
from __future__ import annotations

import base64
import io
import os

from PIL import Image  # light; used by the CPU-side data-url helpers and type hints

from gpu_shared import (
    GUIDANCE,
    LOW_VRAM_GB,
    MAX_SIDE,
    MODEL_ID,
    STEPS,
    build_edit_prompt,
    build_initial_prompt,
)

BACKEND = os.environ.get("FROGQUEST_BACKEND", "zerogpu").lower()
if BACKEND != "modal":  # the local/ZeroGPU path (default + any unrecognized value) needs the decorator
    import spaces

# Best-effort: pre-fetch the weights at startup so the first @spaces.GPU call doesn't pay the
# multi-GB download out of its (metered, on ZeroGPU) duration. ZeroGPU-only — a CPU-basic Space
# (modal backend) must NOT pull ~23GB. No-op offline / on a fresh local checkout.
if BACKEND != "modal":
    try:
        from huggingface_hub import snapshot_download
        snapshot_download(MODEL_ID)
    except Exception:
        pass

_pipe = None
_offloaded = False  # True when we used CPU offload (small GPU) instead of a full .to("cuda")


def _get_pipe():
    """Construct the pipeline lazily INSIDE the GPU call so we can read the REAL device's caps
    and adapt — bf16 on Blackwell (ZeroGPU), fp16 on Turing, CPU offload when VRAM is tight. torch
    and diffusers are imported here (not at module top) so the modal backend never loads them."""
    global _pipe, _offloaded
    if _pipe is None:
        import torch
        from diffusers import Flux2KleinPipeline

        bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
        dtype = torch.bfloat16 if bf16 else torch.float16  # Turing (T4) has no bf16
        _pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype)
        vram_gb = (torch.cuda.get_device_properties(0).total_memory / 1e9
                   if torch.cuda.is_available() else 0)
        if 0 < vram_gb < LOW_VRAM_GB:
            # Small GPU: stream modules GPU<-CPU per step so the ~23GB of weights fit in ~24GB VRAM.
            _pipe.enable_model_cpu_offload()
            _offloaded = True
    return _pipe


def _gen(prompt: str, image, seed: int) -> Image.Image:
    import torch

    pipe = _get_pipe()
    if not _offloaded:
        pipe.to("cuda")  # full-residency path (big GPU); offload manages its own device moves
    generator = torch.Generator("cuda").manual_seed(int(seed))
    result = pipe(
        prompt=prompt,
        image=image,
        generator=generator,
        num_inference_steps=STEPS,
        guidance_scale=GUIDANCE,
        height=MAX_SIDE,
        width=MAX_SIDE,
    )
    return result.images[0]


# ----------------------------- local (in-Space, ZeroGPU) implementations -----------------------------

def _initial_image_local(user_photo: Image.Image, art_style: str, scene_prompt: str, seed: int) -> Image.Image:
    """Generate the quest's initial scene with the user as the hero (photo as reference)."""
    return _gen(build_initial_prompt(art_style, scene_prompt), image=[user_photo], seed=seed)


def _initial_images_local(user_photo: Image.Image, art_style: str, scene_prompts: list[str],
                          seed: int) -> list[Image.Image]:
    """Batch: ALL of a forge's initial scenes in ONE GPU call. On ZeroGPU each @spaces.GPU call is
    a separate metered reservation (+~30s admission overhead), so looping initial_image would cost
    N reservations; here the pipeline loads once and the gens run back-to-back."""
    return [_gen(build_initial_prompt(art_style, p), image=[user_photo], seed=seed)
            for p in scene_prompts]


def _edit_image_local(base_image: Image.Image, edit_instruction: str, art_style: str, seed: int) -> Image.Image:
    """Edit the existing image into a success/failure state."""
    return _gen(build_edit_prompt(art_style, edit_instruction), image=base_image, seed=seed)


# ----------------------------- modal (off-Space) wrappers -----------------------------

def _initial_image_modal(user_photo: Image.Image, art_style: str, scene_prompt: str, seed: int) -> Image.Image:
    import modal
    flux = modal.Cls.from_name("frogquest", "Flux")()
    return flux.initial.remote(user_photo, art_style, scene_prompt, seed)


def _initial_images_modal(user_photo: Image.Image, art_style: str, scene_prompts: list[str],
                          seed: int) -> list[Image.Image]:
    import modal
    flux = modal.Cls.from_name("frogquest", "Flux")()
    return flux.initials.remote(user_photo, art_style, scene_prompts, seed)


def _edit_image_modal(base_image: Image.Image, edit_instruction: str, art_style: str, seed: int) -> Image.Image:
    import modal
    flux = modal.Cls.from_name("frogquest", "Flux")()
    return flux.edit.remote(base_image, edit_instruction, art_style, seed)


# ----------------------------- bind public names from the backend -----------------------------
if BACKEND == "modal":
    initial_image = _initial_image_modal
    initial_images = _initial_images_modal
    edit_image = _edit_image_modal
else:
    initial_image = spaces.GPU(duration=30)(_initial_image_local)
    initial_images = spaces.GPU(duration=60)(_initial_images_local)  # one reservation, many gens
    edit_image = spaces.GPU(duration=30)(_edit_image_local)


# --- base64 <-> PIL helpers (photo arrives transiently as a data URL; nothing is persisted) ---

def b64_to_pil(data_url_or_b64: str) -> Image.Image:
    s = data_url_or_b64
    if "," in s and s.strip().startswith("data:"):
        s = s.split(",", 1)[1]
    raw = base64.b64decode(s)
    return Image.open(io.BytesIO(raw)).convert("RGB")


def pil_to_data_url(img: Image.Image, fmt: str = "JPEG", quality: int = 82) -> str:
    """Encode a PIL image as a data URL. Defaults to JPEG (~80-200KB for a 768px scene) so the
    quest-image cache fits in localStorage; pass fmt="PNG" for lossless when size doesn't matter.
    """
    fmt = (fmt or "JPEG").upper()
    if fmt in ("JPG", "JPEG"):
        fmt = "JPEG"
        img = img.convert("RGB")  # JPEG has no alpha channel
        save_kwargs = {"quality": quality}
        mime = "jpeg"
    else:
        save_kwargs = {}
        mime = fmt.lower()
    buf = io.BytesIO()
    img.save(buf, format=fmt, **save_kwargs)
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/{mime};base64,{b64}"