import glob
import os
import random
import subprocess
import sys
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
RAW_REPO = "krea/Krea-2-Raw"
TURBO_REPO = "krea/Krea-2-Turbo"
MAX_SEED = 2**31 - 1

# Both models are loaded at global scope. They share the architecture, so Turbo
# reuses Raw's text encoder / tokenizer / VAE and only loads its own transformer.
pipe_raw = Krea2Pipeline.from_pretrained(RAW_REPO, torch_dtype=DTYPE)
pipe_turbo = Krea2Pipeline.from_pretrained(
    TURBO_REPO,
    text_encoder=pipe_raw.text_encoder,
    tokenizer=pipe_raw.tokenizer,
    vae=pipe_raw.vae,
    torch_dtype=DTYPE,
)
pipe_raw.to("cuda")
pipe_turbo.to("cuda")



def _load_aoti():
    # One compiled Krea2TransformerBlock (kernels only, weights stay live) serves
    # both pipelines since their transformer configs are identical. Mirrors
    # spaces.aoti_blocks_load, but downloads from the private artifact dataset
    # with an explicit write token instead of ambient model-repo auth.
    from huggingface_hub import hf_hub_download
    from spaces.zero.torch.aoti import LazyAOTIModel

    pt2 = hf_hub_download(
        repo_id="multimodalart/Krea-2-aoti",
        filename="Krea2TransformerBlock/package.pt2",
        repo_type="dataset",
        token=os.environ.get("HF_WRITE_TOKEN"),
    )
    aoti_model = LazyAOTIModel(pt2)
    for pipe in (pipe_raw, pipe_turbo):
        for block in pipe.transformer.modules():
            if block.__class__.__name__ == "Krea2TransformerBlock":
                spaces.aoti_patch(block, aoti_model)


try:
    _load_aoti()
    print("AoTI blocks loaded.")
except Exception as e:
    print(f"AoTI load skipped ({e}); running eager.")

PIPES = {"Raw": pipe_raw, "Turbo": pipe_turbo}
DEFAULTS = {
    "Raw": {"steps": 28, "guidance": 4.5},
    "Turbo": {"steps": 8, "guidance": 0.0},
}

# Resolution presets. The model renders up to 2K, but the compiled transformer
# block can exceed this Space's GPU memory above 1024, so 1024 is the default
# and larger sizes are opt-in (see the OOM guard in generate).
RESOLUTIONS = {
    "Square · 1024": (1024, 1024),
    "Portrait · 1024": (832, 1216),
    "Landscape · 1024": (1216, 832),
    "Square · 2K": (2048, 2048),
}

PROMPT_TIPS = """\
Krea 2 is tuned for natural language. Describe the image the way you would describe it to a person.

- Write in full sentences or rich phrases. Longer, more specific prompts give the best results, but short prompts work too.
- Name the things that matter: subject, setting, lighting, color, framing, medium, and mood.
- To render text in the image, wrap the words in quotes, for example: a storefront window with a neon sign that reads "open late".
- The model can render up to 2K, but very high resolutions may run out of GPU memory on this Space. 1024 is the reliable default.

Want help writing longer prompts? An `expansion.txt` system prompt is provided in the [model repo](https://huggingface.co/krea/Krea-2-Turbo) for use with any LLM.
"""

# Drawn from the official Krea 2 prompt guide. These demonstrate the
# detailed, natural-language style the model rewards.
EXAMPLE_PROMPTS = [
    ["immense rocket launch exhaust as seen from extremely close up"],
    [
        "3D rendered matte black designer toy figure, stylized round anthropomorphic shape, "
        "backward black baseball cap, oversized gold-rimmed aviator sunglasses, white traditional "
        "line-art tattoos of tiger and bird on torso, black studded belt with gold buckle, smooth "
        "vinyl texture, studio lighting, solid vibrant blue background, high contrast minimal composition"
    ],
    [
        "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 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."
    ],
    [
        "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 close-up portrait of a young East Asian woman with straight black hair, loose strands sweeping "
        "across her fair skin, and an intense gaze. She wears a light grey collared shirt with a black tie. "
        "A vibrant bouquet of pink and orange lilies with lush green leaves sits in the blurred right "
        "foreground. The background is a solid, striking crimson red. Soft, directional studio lighting "
        "highlights her facial features, creating a high-contrast composition with a shallow depth of field."
    ],
]

# Official sample renders for the prompts above, in the same order. Drop the
# matching PNGs into assets/samples/ and the gallery shows them; clicking a
# thumbnail loads its prompt. Filenames follow the Krea 2 prompt guide. If the
# files are absent, the UI falls back to the text example prompts below.
SAMPLE_DIR = os.path.join(os.path.dirname(__file__), "assets", "samples")
SAMPLE_FILES = ["takeoff.png", "3d.png", "mouse.png", "red.png", "beach.png", "future.png", "jester.png", "flowers.png"]
SAMPLE_LABELS = [
    "Rocket exhaust",
    "Designer toy",
    "Harvest mouse",
    "Editorial portrait",
    "Ligne claire beach",
    "Liquid chrome",
    "Jester",
    "Crimson portrait",
]

_gallery = [
    (os.path.join(SAMPLE_DIR, fname), label, prompt[0])
    for fname, label, prompt in zip(SAMPLE_FILES, SAMPLE_LABELS, EXAMPLE_PROMPTS)
    if os.path.exists(os.path.join(SAMPLE_DIR, fname))
]
GALLERY_ITEMS = [(path, label) for path, label, _ in _gallery]
GALLERY_PROMPTS = [prompt for _, _, prompt in _gallery]

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, negative_prompt, model, steps, guidance, width, height, seed, randomize, progress=None):
    # Scale the GPU reservation by step count and pixel area so larger renders
    # are not killed before they finish.
    megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024))
    return int(int(steps) * 2 * megapixels + 25)


@spaces.GPU(duration=_duration, size="xlarge")
def generate(
    prompt,
    negative_prompt="",
    model="Turbo",
    steps=8,
    guidance=0,
    width=1024,
    height=1024,
    seed=42,
    randomize=False,
    progress=gr.Progress(track_tqdm=True),
):
    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)
    pipe = PIPES[model]
    try:
        image = pipe(
            prompt=prompt,
            negative_prompt=(negative_prompt or None) if guidance > 0 else None,
            height=int(height),
            width=int(width),
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=generator,
        ).images[0]
    except RuntimeError as exc:
        # At high resolution the compiled transformer block can exhaust GPU
        # memory, which surfaces as a CUDA allocation / AOTI runtime error.
        # Recover the worker and tell the user how to fix it.
        torch.cuda.empty_cache()
        raise gr.Error(
            f"Generation failed at {int(width)}x{int(height)}. This is usually the GPU running "
            "out of memory at high resolution. Try 1024x1024 or a smaller size."
        ) from exc
    return image, seed


def on_model_change(model):
    d = DEFAULTS[model]
    return (
        gr.update(value=d["steps"]),
        gr.update(value=d["guidance"]),
        gr.update(interactive=d["guidance"] > 0),
    )


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


# Krea brand identity: neutral grayscale foundation with a single blue action
# accent (krea.ai/press). Dark surfaces, mono utility type, accent reserved for
# the primary action and focus states.
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 .meta {
  margin-top: 18px;
  display: flex;
  justify-content: space-between;
  align-items: center;
  flex-wrap: wrap;
  gap: 12px;
}
#krea-header .badges { 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;
}
#krea-header .links { display: flex; gap: 16px; }
#krea-header .links a {
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 11px;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  color: #737373;
  text-decoration: none;
  transition: color 0.15s ease;
}
#krea-header .links a:hover { color: #f5f5f5; }

#generate-btn { font-weight: 600; letter-spacing: 0.01em; }

#result-image { min-height: 420px; border-radius: 10px; overflow: hidden; }

/* Inline code chips legible on the dark theme (e.g. the expansion.txt mention in tips). */
.gradio-container code,
.gradio-container .prose code {
  background: #171717 !important;
  color: #d4d4d5 !important;
  border: 1px solid #262626 !important;
  border-radius: 5px !important;
  padding: 2px 7px !important;
  font-family: 'JetBrains Mono', ui-monospace, monospace !important;
  font-size: 0.85em !important;
}

footer { display: none !important; }
.gradio-container .prose a { color: """ + KREA_ACCENT + """; }
"""

with gr.Blocks(title="Krea 2") 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</h1>
              <p class="subtitle">Generate images from natural language. Pick Raw for CFG-guided control or Turbo for fast, few-step results.</p>
              <div class="meta">
                <div class="badges">
                  <span class="badge">Raw · CFG</span>
                  <span class="badge">Turbo · few-step</span>
                </div>
                <div class="links">
                  <a href="https://www.krea.ai/blog/krea-2-technical-report" target="_blank" rel="noopener">Technical report ↗</a>
                  <a href="https://github.com/krea-ai/krea-2" target="_blank" rel="noopener">GitHub ↗</a>
                </div>
              </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,
                )
                model = gr.Radio(["Turbo", "Raw"], value="Turbo", label="Model")
                run = gr.Button("Generate", variant="primary", elem_id="generate-btn")

                with gr.Accordion("Prompting tips", open=False):
                    gr.Markdown(PROMPT_TIPS)

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

                with gr.Accordion("Advanced", open=False):
                    negative_prompt = gr.Textbox(
                        label="Negative prompt",
                        lines=1,
                        interactive=False,
                        info="Available with Raw, where guidance is above 0.",
                    )
                    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.Slider(512, 2048, value=1024, step=16, label="Width")
                        height = gr.Slider(512, 2048, value=1024, step=16, label="Height")
                    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")

        if GALLERY_ITEMS:
            gallery = gr.Gallery(
                value=GALLERY_ITEMS,
                label="Example prompts",
                columns=4,
                height="auto",
                object_fit="cover",
                allow_preview=False,
                elem_id="examples-gallery",
            )

            def use_example(evt: gr.SelectData):
                # Clicking a sample loads its prompt into the box, ready to run.
                return GALLERY_PROMPTS[evt.index]

            gallery.select(use_example, None, prompt)
        else:
            # No bundled sample images present; show the prompts as text.
            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,
            )

    model.change(on_model_change, model, [steps, guidance, negative_prompt])
    resolution.change(on_resolution_change, resolution, [width, height])

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

demo.launch(theme=theme, css=CSS)