from __future__ import annotations

import os
import random
import threading
import time

try:
    import spaces  # type: ignore

    GPU = spaces.GPU
except Exception:  # pragma: no cover — local/box runs without the ZeroGPU shim
    def GPU(*dargs, **dkwargs):  # noqa: N802
        def wrap(fn):
            return fn
        if len(dargs) == 1 and callable(dargs[0]) and not dkwargs:
            return dargs[0]
        return wrap

import gradio as gr
import torch
from diffusers import Flux2KleinPipeline
from PIL import Image

# Env-overridable so the same file serves ZeroGPU (distilled, 4-step) and the
# local box (e.g. KLEIN_MODEL=black-forest-labs/FLUX.2-klein-base-4B KLEIN_STEPS=50).
MODEL_ID = os.environ.get("KLEIN_MODEL", "black-forest-labs/FLUX.2-klein-4B")
STEPS = int(os.environ.get("KLEIN_STEPS", "4"))
GUIDANCE = float(os.environ.get("KLEIN_GUIDANCE", "1.0"))
MAX_SEED = 2_147_483_647

# Optional character/style LoRAs, loadable per-request via the `lora` argument.
LORA_REPO = os.environ.get("KLEIN_LORA_REPO", "polats/weiner-klein-lora")
LORAS = {
    "weiner": "weiner_klein_4b_v1.safetensors",          # step 1000 (final)
    "weiner750": "weiner_klein_4b_v1_000000750.safetensors",
    "weiner500": "weiner_klein_4b_v1_000000500.safetensors",
}

print(f"Loading {MODEL_ID} on CPU")
_t0 = time.time()
pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
print(f"Loaded in {time.time() - _t0:.1f}s")

_lora_active: str | None = None
# The pipe + its loaded LoRA are process-global and shared by EVERY caller of this Space.
# Serialize LoRA-state changes with the generation so concurrent requests can't clobber
# each other's adapter (e.g. a no-LoRA call rendering while another request has weiner
# attached). Held across _ensure_lora + pipe.to() + pipe() for the whole critical section.
_gen_lock = threading.Lock()


def _ensure_lora(name: str) -> None:
    global _lora_active
    name = (name or "").strip()
    if name == (_lora_active or ""):
        return
    if _lora_active:
        pipe.unload_lora_weights()
        _lora_active = None
    if name:
        if name not in LORAS:
            raise gr.Error(f"unknown lora '{name}' (have: {', '.join(LORAS)})")
        pipe.load_lora_weights(
            LORA_REPO,
            weight_name=LORAS[name],
            token=os.environ.get("HF_TOKEN") or None,
        )
        _lora_active = name


@GPU(duration=60)
def generate(prompt: str, seed: int = 42, lora: str = "", ref_image=None):
    if not prompt or not prompt.strip():
        raise gr.Error("prompt required")
    if seed is None or int(seed) < 0:
        seed = random.randint(0, MAX_SEED)
    # Exclusive ownership of the shared pipe for the whole LoRA-set + generate, so another
    # caller's adapter can't leak into this render (or vice-versa).
    with _gen_lock:
        _ensure_lora(lora)
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        pipe.to(dev)
        kwargs = dict(
            prompt=prompt.strip(),
            width=1024,
            height=1024,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE,
            generator=torch.Generator(device=dev).manual_seed(int(seed)),
        )
        if ref_image is not None:
            # FLUX.2 native text+image->image: the reference (e.g. a pose-control
            # skeleton rendered from exact mocap joints) steers the composition.
            img = ref_image if isinstance(ref_image, Image.Image) else Image.open(ref_image)
            kwargs["image"] = img.convert("RGB").resize((1024, 1024))
        img = pipe(**kwargs).images[0]
        pipe.to("cpu")
        if dev == "cuda":
            torch.cuda.empty_cache()
    return img


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", lines=4),
        gr.Number(label="Seed", value=42, precision=0),
        gr.Textbox(label="LoRA", value="", placeholder="blank = none; weiner / weiner750 / weiner500"),
        gr.Image(label="Pose reference (optional)", type="pil"),
    ],
    outputs=gr.Image(type="pil", label="Portrait"),
    api_name="generate",
    title="Tiny Army Klein ZeroGPU",
)

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