src/apps/p5_memory_quilt/flux_backend.py

from __future__ import annotations

"""Strict local FLUX backend.

This module is deliberately import-safe: it should not require torch/diffusers unless
its functions are called. The backend refuses to synthesize or fetch fallback tiles.
To enable real FLUX tile generation, point P5_FLUX_MODEL_DIR at a local diffusers
checkpoint directory and install the optional deps (torch, diffusers, transformers).
"""

from dataclasses import dataclass
from pathlib import Path
import os
import time


@dataclass(frozen=True)
class FluxConfig:
    model_id: str = "black-forest-labs/FLUX.2-klein-base-4B"
    model_dir: str = ""
    steps: int = 4
    guidance_scale: float | None = None
    seed: int = 0
    height: int = 512
    width: int = 512


def _load_config(size: int = 512) -> FluxConfig:
    model_id = os.environ.get("P5_FLUX_MODEL_ID", "black-forest-labs/FLUX.2-klein-base-4B").strip()
    model_dir = os.environ.get("P5_FLUX_MODEL_DIR", "").strip()
    steps = int(os.environ.get("P5_FLUX_STEPS", "4"))
    guidance_raw = os.environ.get("P5_FLUX_GUIDANCE", "").strip()
    guidance_scale = float(guidance_raw) if guidance_raw else None
    seed = int(os.environ.get("P5_FLUX_SEED", "0"))
    return FluxConfig(
        model_id=model_id,
        model_dir=model_dir,
        steps=steps,
        guidance_scale=guidance_scale,
        seed=seed,
        height=size,
        width=size,
    )


def available() -> bool:
    """Heuristic: do we have a local directory pointer for FLUX weights?"""
    model_dir = os.environ.get("P5_FLUX_MODEL_DIR", "").strip()
    return bool(model_dir and Path(model_dir).exists())


def try_generate(prompt: str, size: int = 512):
    """Generate a tile with locally available FLUX weights.

    Returns:
        (image, meta) where image is a PIL Image, meta includes backend details.

    Raises:
        RuntimeError if optional deps are missing or the local model isn't usable.
    """

    cfg = _load_config(size=size)
    started_at = time.perf_counter()

    # Import lazily so unit tests and lightweight installs don't pay torch/diffusers.
    try:
        import torch  # type: ignore
    except Exception as exc:
        raise RuntimeError("torch is not installed; cannot run FLUX backend") from exc

    try:
        import diffusers  # type: ignore
    except Exception as exc:
        raise RuntimeError("diffusers is not installed; cannot run FLUX backend") from exc

    if not cfg.model_dir:
        raise RuntimeError("P5_FLUX_MODEL_DIR is not set; real FLUX generation requires a local checkpoint directory")

    model_ref: str = cfg.model_dir

    # Prefer the explicit Flux2KleinPipeline if present; otherwise fall back to FluxPipeline or DiffusionPipeline.
    PipelineCls = getattr(diffusers, "Flux2KleinPipeline", None)
    if PipelineCls is None:
        PipelineCls = getattr(diffusers, "FluxPipeline", None)
    if PipelineCls is None:
        from diffusers import DiffusionPipeline as PipelineCls  # type: ignore

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    # Force offline behavior.
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

    pipe = PipelineCls.from_pretrained(
        model_ref,
        torch_dtype=dtype,
        local_files_only=True,
    )
    pipe = pipe.to(device)

    generator = torch.Generator(device=device)
    generator.manual_seed(cfg.seed)

    kwargs = {
        "prompt": prompt,
        "num_inference_steps": cfg.steps,
        "height": cfg.height,
        "width": cfg.width,
        "generator": generator,
    }
    if cfg.guidance_scale is not None:
        kwargs["guidance_scale"] = cfg.guidance_scale

    result = pipe(**kwargs)

    images = getattr(result, "images", None)
    if not images:
        raise RuntimeError("FLUX pipeline returned no images")

    image = images[0]
    elapsed_ms = round((time.perf_counter() - started_at) * 1000.0, 2)
    generation_stats = {
        "steps": cfg.steps,
        "guidance_scale": cfg.guidance_scale,
        "seed": cfg.seed,
        "height": cfg.height,
        "width": cfg.width,
        "elapsed_ms": elapsed_ms,
        "device": device,
        "dtype": str(dtype),
    }
    meta = {
        "adapter_name": "flux-diffusers",
        "backend": "flux-diffusers",
        "device": device,
        "dtype": str(dtype),
        "model_dir": cfg.model_dir,
        "model_id": cfg.model_id,
        "generation_stats": generation_stats,
    }
    return image, meta