import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import torch
import gradio as gr
import numpy as np
import random

from diffusers import DiffusionPipeline, AutoencoderKL, AutoencoderTiny
from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images

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

MODEL_ID = "black-forest-labs/FLUX.1-dev"

# Tiny VAE for fast live previews during the denoising loop.
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
# Full VAE for the crisp final decode.
good_vae = AutoencoderKL.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=dtype).to(device)

pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype, vae=taef1).to(device)
torch.cuda.empty_cache()

pipe.flux_pipe_call_that_returns_an_iterable_of_images = (
    flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048


def _estimate_duration(prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, *args, **kwargs):
    pixels = (width * height) / (1024 * 1024)
    return int(min(240, 30 + num_inference_steps * pixels * 2.2))


@spaces.GPU(duration=_estimate_duration)
def generate(
    prompt: str,
    seed: int = 0,
    randomize_seed: bool = True,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3.5,
    num_inference_steps: int = 28,
):
    """Generate an image from a text prompt with FLUX.1-dev. Streams live previews while denoising."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        output_type="pil",
        good_vae=good_vae,
    ):
        yield img, seed


examples = [
    "a tiny astronaut hatching from an egg on the moon",
    "a cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
    "a futuristic cyborg chef cooking in a neon-lit kitchen",
    "a serene mountain lake at golden hour, ultra detailed, photorealistic",
]

css = """
#header { text-align: center; margin-bottom: 0.5rem; }
#subtitle { text-align: center; opacity: 0.75; }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # FLUX.1 [dev] ⚡
        Text-to-image with the 12B-param [FLUX.1 dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
        rectified flow transformer. Live previews stream while your image is being denoised.
        """,
        elem_id="header",
    )

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            placeholder="Describe the image you want to generate…",
            scale=4,
            autofocus=True,
        )
        run_button = gr.Button("Generate", variant="primary", scale=1)

    result = gr.Image(label="Result", interactive=False, height=520)

    with gr.Accordion("Advanced settings", open=False):
        with gr.Row():
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row():
            width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
        with gr.Row():
            guidance_scale = gr.Slider(label="Guidance scale", minimum=1, maximum=15, step=0.1, value=3.5)
            num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=28)

    gr.Examples(
        examples=examples,
        fn=generate,
        inputs=[prompt],
        outputs=[result, seed],
        cache_examples=True,
        cache_mode="lazy",
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=generate,
        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, seed],
        api_name="generate",
    )

    gr.Markdown(
        """
        *Model: [FLUX.1 [dev]](https://huggingface.co/black-forest-labs/FLUX.1-dev) ·
        [non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md) ·
        [FLUX.1 dev GitHub](https://github.com/black-forest-labs/flux) ·
        Built with [diffusers](https://github.com/huggingface/diffusers) on ZeroGPU*
        """,
        elem_id="subtitle",
    )

demo.launch(
    mcp_server=True,
    theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
    css=css,
)