import spaces  # must be imported before torch

import os
import random

import gradio as gr
import torch

from pipeline import generate_image, load_models

HERE = os.path.dirname(os.path.abspath(__file__))
MAX_SEED = 2**31 - 1

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

MODELS = load_models(device="cuda")


def _duration(prompt, seed=0, randomize_seed=True, num_steps=50, *args, **kwargs):
    # measured on ZeroGPU: ~7.9 s cold / ~2.9 s warm for 50 steps (incl. VAE decode)
    return int(7 + 0.07 * int(num_steps))


@spaces.GPU(duration=_duration)
def generate(
    prompt: str,
    seed: int = 0,
    randomize_seed: bool = True,
    num_steps: int = 50,
    cfg_scale: float = 7.0,
    guidance_low: float = 0.35,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a 512x512 image from a text prompt with the GenFirst EiT MMDiT-XL/2 prior (variant A).

    Args:
        prompt: Text description of the image.
        seed: Random seed (ignored when randomize_seed is True).
        randomize_seed: Pick a random seed for every call.
        num_steps: Number of Euler sampling steps.
        cfg_scale: Classifier-free guidance scale.
        guidance_low: CFG is only applied for t in [guidance_low, 1.0] (t=1 is noise).

    Returns:
        The generated image and the seed that was used.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    num_steps = int(num_steps)
    progress(0, desc="Sampling")
    image = generate_image(
        MODELS, prompt.strip(), seed, num_steps, float(cfg_scale), float(guidance_low), 1.0,
        device="cuda", autocast=True,
        callback=lambda i: progress((i + 1) / num_steps, desc=f"Step {i + 1}/{num_steps}"),
    )
    return image, seed


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["a photograph of an astronaut riding a horse"],
    ["a cute corgi puppy in a field of flowers, bokeh"],
    ["an oil painting of a lighthouse during a storm"],
    ["a bowl of ramen with an egg, food photography"],
    ["a photo of a purple wine glass and a black apple"],
    ["a photo of a dog right of a teddy bear"],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """# GenFirst: Text-to-Image (EiT MMDiT-XL/2)
512x512 text-to-image from **GenFirst** (*Generation Before Reconstruction*): an MMDiT-XL/2 flow-matching prior with a CLIP-L text encoder, trained on an end-to-end GenFirst FLUX VAE (variant A, GenEval 0.90).
[Paper](https://arxiv.org/abs/2608.29335) · [Code](https://github.com/ByteDance-Seed/GenFirst) · [Weights](https://huggingface.co/ByteDance-Seed/GenFirst) · [Project page](https://tom-zgt.github.io/GenFirst/) · weights are CC BY-NC 4.0 (non-commercial)."""
        )
        with gr.Row():
            prompt = gr.Textbox(
                show_label=False, placeholder="Describe the image you want…", container=False, scale=4, max_lines=3
            )
            run = gr.Button("Generate", variant="primary", scale=1)
        result = gr.Image(label="Result", type="numpy", format="png", height=512)
        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():
                num_steps = gr.Slider(label="Sampling steps", minimum=10, maximum=100, step=1, value=50)
                cfg_scale = gr.Slider(label="Guidance scale (CFG)", minimum=1.0, maximum=12.0, step=0.5, value=7.0)
                guidance_low = gr.Slider(
                    label="CFG interval start (t_low)", minimum=0.0, maximum=1.0, step=0.05, value=0.35,
                    info="CFG is applied on t ∈ [t_low, 1]; the paper's GenEval setting is 0.35",
                )
        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt],
            outputs=[result, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    gr.on(
        triggers=[run.click, prompt.submit],
        fn=generate,
        inputs=[prompt, seed, randomize_seed, num_steps, cfg_scale, guidance_low],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)