"""FLUX.2 [dev] 4-bit studio — text-to-image and multi-reference editing, self-contained.

Runs the pre-quantized `diffusers/FLUX.2-dev-bnb-4bit` checkpoint: full dev quality at
about a fifth of the footprint, with the text encoder loaded INSIDE this Space. That
last part matters — the public dev demo delegates its text encoder to a separate Space,
so it breaks for good whenever that other Space goes away.

The /infer signature mirrors the common FLUX.2 Space API, so existing clients only need
to change the space id.
"""
import os
import random

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import Flux2Pipeline
from PIL import Image

MODEL_ID = os.environ.get("FLUX2_MODEL", "diffusers/FLUX.2-dev-bnb-4bit")
MAX_SEED = np.iinfo(np.int32).max
MAX_SIDE = 1536

pipe = Flux2Pipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map=None,               # let the offload hook place things, not the loader
)
pipe.enable_model_cpu_offload()    # 4-bit weights still want room to breathe


def _duration(prompt, image_list, seed, width, height, num_inference_steps, guidance_scale):
    """ZeroGPU reserves the slot before the call and passes it the same arguments, so
    this signature must mirror _run's exactly."""
    n = 1 + 0.6 * len(image_list or [])
    return int(min(300, max(90, int(num_inference_steps) * 2.4 * n + 40)))


@spaces.GPU(duration=_duration)
def _run(prompt, image_list, seed, width, height, num_inference_steps, guidance_scale):
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    kwargs = dict(prompt=prompt, width=int(width), height=int(height),
                  num_inference_steps=int(num_inference_steps),
                  guidance_scale=float(guidance_scale), generator=generator)
    if image_list:
        kwargs["image"] = image_list
    return pipe(**kwargs).images[0]


def infer(prompt, input_images, seed=0, randomize_seed=False, width=1024, height=1024,
          num_inference_steps=28, guidance_scale=4.0, prompt_upsampling=False,
          progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    images = []
    for item in (input_images or []):
        path = item[0] if isinstance(item, (list, tuple)) else item
        if isinstance(path, dict):
            path = path.get("image") or path.get("path") or path.get("name")
        if path:
            im = Image.open(path).convert("RGB")
            im.thumbnail((MAX_SIDE, MAX_SIDE))
            images.append(im)
    width = min(int(width), MAX_SIDE)
    height = min(int(height), MAX_SIDE)
    out = _run(prompt, images, seed, width, height, num_inference_steps, guidance_scale)
    return out, seed


with gr.Blocks(title="FLUX.2 dev 4-bit studio") as demo:
    gr.Markdown("## FLUX.2 [dev] 4-bit studio\nText-to-image and multi-reference editing. "
                "Self-contained: the text encoder lives in this Space.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=4)
            input_images = gr.Gallery(label="Reference images (optional)", type="filepath",
                                      columns=4, height=200)
            run = gr.Button("Generate", variant="primary")
            with gr.Accordion("Settings", open=False):
                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                randomize_seed = gr.Checkbox(True, label="Randomize seed")
                width = gr.Slider(256, MAX_SIDE, value=1024, step=32, label="Width")
                height = gr.Slider(256, MAX_SIDE, value=1024, step=32, label="Height")
                num_inference_steps = gr.Slider(4, 50, value=28, step=1, label="Steps")
                guidance_scale = gr.Slider(1.0, 10.0, value=4.0, step=0.1, label="Guidance")
                prompt_upsampling = gr.Checkbox(False, label="Prompt upsampling (unused)")
        with gr.Column():
            result = gr.Image(label="Result", type="pil")
            used_seed = gr.Number(label="Seed used")

    run.click(infer,
              inputs=[prompt, input_images, seed, randomize_seed, width, height,
                      num_inference_steps, guidance_scale, prompt_upsampling],
              outputs=[result, used_seed], api_name="infer")

demo.queue().launch(show_error=True)