import os

os.environ.setdefault("VF_HF_ATTN_IMPL", "sdpa")

import spaces
import torch
import gradio as gr
import random

from mage_flow import MageFlowPipeline
from mage_flow.models.modules._attn_backend import set_attn_backend

MODEL_ID = "microsoft/Mage-Flow-Base"

pipe = MageFlowPipeline.from_pretrained(MODEL_ID, device="cuda")

# Override the attention backend to SDPA after model load — the model's
# __init__ resets it to "flash2" (the default attn_type), but flash_attn
# is not available on ZeroGPU. SDPA is torch-native and works everywhere.
set_attn_backend("sdpa")


def _estimate_duration(prompt, negative_prompt, height, width, steps, cfg, seed, *args, **kwargs):
    """Estimate GPU time needed based on steps and resolution.
    Measured: ~15s for 30 steps at 1024x1024 with SDPA backend + content screening.
    """
    pixels = height * width
    base = 20  # content screening + text encoding overhead
    step_time = 0.5 if pixels <= 1024 * 1024 else 1.0
    return min(300, base + int(steps * step_time))


@spaces.GPU(duration=_estimate_duration)
def generate(
    prompt: str,
    negative_prompt: str = "",
    height: int = 1024,
    width: int = 1024,
    steps: int = 30,
    cfg: float = 5.0,
    seed: int = 42,
    randomize_seed: bool = False,
):
    """Generate an image from a text prompt using Mage-Flow-Base.

    Args:
        prompt: Text description of the image to generate.
        negative_prompt: What to avoid in the generated image.
        height: Output image height (512-2048, multiple of 16).
        width: Output image width (512-2048, multiple of 16).
        steps: Number of denoising steps (more = higher quality, slower).
        cfg: Classifier-free guidance scale (higher = more prompt adherence).
        seed: Random seed for reproducibility.
        randomize_seed: If True, use a random seed each time.
    """
    if randomize_seed:
        seed = random.randint(0, 2**32 - 1)
    seed = int(seed)
    height = max(512, min(2048, int(height) // 16 * 16))
    width = max(512, min(2048, int(width) // 16 * 16))

    neg = negative_prompt.strip() if negative_prompt else " "

    result = pipe.generate(
        [prompt],
        neg_prompts=[neg],
        seeds=[seed],
        steps=int(steps),
        cfg=float(cfg),
        heights=[height],
        widths=[width],
    )
    return result[0], seed


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

with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Mage-Flow-Base Text-to-Image")
    gr.Markdown(
        "Generate images with [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base), "
        "a 4.1B parameter native-resolution MMDiT model. Supports any resolution from 512 to 2048."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to generate...",
                show_label=False,
                container=False,
                scale=4,
            )
            run_btn = gr.Button("Generate", variant="primary", scale=1)

        output = gr.Image(label="Generated Image", show_label=False, height=512)

        with gr.Accordion("Advanced settings", open=False):
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                placeholder="What to avoid...",
                value="",
            )
            with gr.Row():
                height = gr.Slider(
                    label="Height", minimum=512, maximum=2048, step=16, value=1024
                )
                width = gr.Slider(
                    label="Width", minimum=512, maximum=2048, step=16, value=1024
                )
            with gr.Row():
                steps = gr.Slider(
                    label="Steps", minimum=1, maximum=100, step=1, value=30
                )
                cfg = gr.Slider(
                    label="CFG scale", minimum=1.0, maximum=20.0, step=0.1, value=5.0
                )
            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        gr.Examples(
            examples=[
                ["A close-up portrait of an elderly African man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic."],
                ["An immersive close-up of a steaming bowl of Sichuan mapo tofu over jasmine rice served on a hand-thrown ceramic plate, finished with a wedge of citrus. Surface oils catch a tiny specular highlight. Shot with a Hasselblad H6D-100c, ambient window light, the kind of image that makes the viewer hungry."],
                ["the Salar de Uyuni mirror surface captured at high noon, with intimate stillness permeating the air. dew beads on every blade of grass. National Geographic editorial, cinematic depth, fine-grained natural texture."],
                ["A serene Japanese garden with a red wooden bridge over a koi pond, cherry blossoms floating on still water, morning mist, photorealistic."],
            ],
            inputs=[prompt],
            outputs=[output, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, height, width, steps, cfg, seed, randomize_seed],
        outputs=[output, seed],
        api_name="generate",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)