import os
import random
import spaces
import torch
import gradio as gr
from huggingface_hub import login

if os.environ.get("HF_TOKEN"):
    login(token=os.environ["HF_TOKEN"])

from diffusers import Krea2Pipeline

DTYPE = torch.bfloat16
TURBO_REPO = "krea/Krea-2-Turbo"
MAX_SEED = 2**31 - 1

pipe = Krea2Pipeline.from_pretrained(TURBO_REPO, torch_dtype=DTYPE)
pipe.to("cuda")

RESOLUTIONS = {
    "Square 1024": (1024, 1024),
    "Portrait 1024": (832, 1216),
    "Landscape 1024": (1216, 832),
}

EXAMPLE_PROMPTS = [
    ["immense rocket launch exhaust as seen from extremely close up"],
    [
        "A tiny, russet-brown harvest mouse clings to a slender diagonal branch amid vibrant green "
        "lobed leaves and small round buds. The mouse has soft textured fur, glossy black eyes, a pink "
        "nose, fine whiskers, and delicate pink paws firmly gripping the wood. In this macro photograph, "
        "an extremely shallow depth of field sharply focuses on the animal's face. The deep green "
        "background dissolves into a smooth, creamy bokeh, illuminated by soft, diffused natural lighting "
        "that highlights the intricate details of the fur and foliage."
    ],
    [
        "high-fashion editorial portrait of a young East Asian woman, short choppy platinum blonde bob "
        "with heavy bangs, looking over her bare shoulder to the right, lips playfully pursed, wearing a "
        "structured black top with an architectural protruding bust detail and thin straps, delicate gold "
        "hoop earrings, arm bent with hand resting on hip, warm skin tones, solid striking crimson red "
        "background, soft directional studio lighting, cinematic color palette, medium close-up shot"
    ],
    [
        "A surreal retro-futuristic space scene features liquid chrome forming an abstract face merging "
        "with a glowing planetary horizon. The foreground is dominated by swirling, highly reflective "
        "metallic fluid that distorts into a stylized, melting facial profile with deep shadows and bright "
        "silver highlights. This undulating chrome form rests against the curved, atmospheric edge of a "
        "massive planet bathed in a soft electric blue and purple glow. Set against a deep black starfield, "
        "the artwork employs a vintage 1980s airbrush aesthetic with smooth gradients, ethereal lighting, "
        "and high-contrast metallic rendering."
    ],
    [
        "Stylized digital painting of a menacing jester figure rendered with bold, expressive brushstrokes "
        "and a vibrant, almost psychedelic color palette against a pitch-black background. Dynamic low-angle "
        "perspective forces a dramatic, imposing composition as the character leans forward, one leg raised "
        "high. The jester wears a classic multi-pointed hat with bells, a ruffled collar, and striped tights "
        "in alternating shades of purple, blue, and chartreuse. The figure's face is a smooth, faceless, pale "
        "mauve mask with a single glowing white point of light at the center, and it grips a massive ornate "
        "sword with a glowing ethereal white blade. Theatrical lighting, dark fantasy concept-art aesthetic."
    ],
    [
        "A minimalist flat-color illustration of a person wading through expansive shallow ocean waves "
        "beneath a pale peach sky. The dark-skinned figure, wearing an orange swim cap, light blue top, and "
        "bright green shorts, steps carefully through knee-deep water. The ocean is rendered in muted mint "
        "green with delicate, thin black linework detailing the continuous ripples and gentle whitecaps. "
        "Soft pinkish-peach reflections echo the sky on the water's surface. The high-angle wide perspective "
        "emphasizes the vast negative space of the water, utilizing a clean ligne claire drawing aesthetic "
        "with a subtle paper texture."
    ],
]

PLACEHOLDER = (
    "Describe your image in natural language. e.g. a russet harvest mouse clinging to a "
    'branch, macro photograph, shallow depth of field, creamy green bokeh, soft natural light. '
    'Wrap words in "quotes" to render them as text.'
)


def _duration(prompt, steps, guidance, width, height, seed, randomize, *args, **kwargs):
    megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024))
    return int(int(steps) * 2 * megapixels + 25)


@spaces.GPU(duration=_duration)
def generate(
    prompt,
    steps=8,
    guidance=0.0,
    width=1024,
    height=1024,
    seed=42,
    randomize=False,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt using Krea 2 Turbo."""
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt to generate an image.")
    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    generator = torch.Generator("cuda").manual_seed(seed)
    try:
        image = pipe(
            prompt=prompt,
            height=int(height),
            width=int(width),
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=generator,
        ).images[0]
    except RuntimeError as exc:
        torch.cuda.empty_cache()
        raise gr.Error(
            f"Generation failed at {int(width)}x{int(height)}. Try 1024x1024 or a smaller size."
        ) from exc
    return image, seed


def on_resolution_change(label):
    w, h = RESOLUTIONS[label]
    return gr.update(value=w), gr.update(value=h)


KREA_ACCENT = "#2b5cff"

theme = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    neutral_hue=gr.themes.colors.neutral,
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
    body_background_fill="#000000",
    body_background_fill_dark="#000000",
    body_text_color="#f5f5f5",
    background_fill_primary="#0d0d0d",
    background_fill_secondary="#0d0d0d",
    block_background_fill="#0d0d0d",
    block_border_color="#262626",
    block_border_width="1px",
    block_label_background_fill="#0d0d0d",
    block_label_text_color="#737373",
    block_title_text_color="#d4d4d5",
    border_color_primary="#262626",
    input_background_fill="#000000",
    input_border_color="#262626",
    input_border_color_focus=KREA_ACCENT,
    button_primary_background_fill=KREA_ACCENT,
    button_primary_background_fill_hover="#1f4fff",
    button_primary_text_color="#ffffff",
    button_primary_border_color=KREA_ACCENT,
    button_secondary_background_fill="#171717",
    button_secondary_background_fill_hover="#262626",
    button_secondary_text_color="#f5f5f5",
    button_secondary_border_color="#262626",
    slider_color=KREA_ACCENT,
)

CSS = """
.gradio-container { background: #000 !important; }
#page { max-width: 1120px; margin: 0 auto; padding: 4px 8px 32px; }

#krea-header {
  padding: 32px 6px 22px;
  border-bottom: 1px solid #1a1a1a;
  margin-bottom: 22px;
}
#krea-header .eyebrow {
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 11px;
  letter-spacing: 0.24em;
  text-transform: uppercase;
  color: #737373;
}
#krea-header h1 {
  font-size: 42px;
  font-weight: 600;
  letter-spacing: -0.025em;
  line-height: 1.05;
  margin: 10px 0 6px;
  color: #fff;
}
#krea-header .subtitle {
  font-size: 15px;
  line-height: 1.5;
  color: #a3a3a3;
  margin: 0;
  max-width: 60ch;
}
#krea-header .badges {
  margin-top: 18px;
  display: flex;
  gap: 8px;
}
#krea-header .badge {
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 10px;
  letter-spacing: 0.12em;
  text-transform: uppercase;
  color: #d4d4d5;
  border: 1px solid #262626;
  border-radius: 999px;
  padding: 4px 10px;
}

#generate-btn { font-weight: 600; letter-spacing: 0.01em; }
#result-image { min-height: 420px; border-radius: 10px; overflow: hidden; }

footer { display: none !important; }
"""

with gr.Blocks(title="Krea 2 Turbo Demo", theme=theme, css=CSS) as demo:
    with gr.Column(elem_id="page"):
        gr.HTML(
            """
            <header id="krea-header">
              <div class="eyebrow">KREA · TEXT-TO-IMAGE</div>
              <h1>Krea 2 Turbo</h1>
              <p class="subtitle">Fast, few-step image generation from natural language. 8 steps, zero guidance.</p>
              <div class="badges">
                <span class="badge">13B params</span>
                <span class="badge">8 steps</span>
                <span class="badge">Turbo distilled</span>
              </div>
            </header>
            """
        )

        with gr.Row(equal_height=False):
            with gr.Column(scale=5, elem_classes="panel"):
                prompt = gr.Textbox(
                    label="Prompt",
                    lines=4,
                    placeholder=PLACEHOLDER,
                    show_label=True,
                    autofocus=True,
                )
                run = gr.Button("Generate", variant="primary", elem_id="generate-btn")

                resolution = gr.Radio(
                    list(RESOLUTIONS.keys()),
                    value="Square 1024",
                    label="Resolution",
                )

                with gr.Accordion("Advanced", open=False):
                    steps = gr.Slider(1, 50, value=8, step=1, label="Steps")
                    guidance = gr.Slider(0.0, 10.0, value=0.0, step=0.1, label="Guidance scale")
                    with gr.Row():
                        width = gr.Number(value=1024, label="Width", precision=0, interactive=False)
                        height = gr.Number(value=1024, label="Height", precision=0, interactive=False)
                    with gr.Row():
                        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                        randomize = gr.Checkbox(value=True, label="Randomize seed")

            with gr.Column(scale=6, elem_classes="panel"):
                output = gr.Image(label="Result", format="png", elem_id="result-image")

        gr.Examples(
            fn=generate,
            examples=EXAMPLE_PROMPTS,
            inputs=[prompt],
            outputs=[output, seed],
            cache_examples=True,
            cache_mode="lazy",
            label="Example prompts",
            examples_per_page=4,
        )

    resolution.change(on_resolution_change, resolution, [width, height])

    inputs = [prompt, steps, guidance, width, height, seed, randomize]
    run.click(generate, inputs, [output, seed])
    prompt.submit(generate, inputs, [output, seed])

demo.launch(mcp_server=True)