import random

import gradio as gr
import spaces
import torch

from inference import generate as nf2_generate
from inference import load
from prompt_check import check_prompt
from prompt_edit import edit_prompt

MAX_SEED = 2**31 - 1

pipe = load()
pipe.to("cuda")

RESOLUTIONS = {
    "Square · 1:1": (1024, 1024),
    "Portrait · 2:3": (832, 1216),
    "Landscape · 3:2": (1216, 832),
    "Cinematic · 16:9": (1360, 768),
    "Square · 2K": (2048, 2048),
    "Portrait · 2K": (1344, 2048),
    "Landscape · 2K": (2048, 1344),
    "Cinematic · 2K": (2048, 1152),
}

EXAMPLE_PROMPTS = [
    [
        "Subject:\n"
        "A hyper-detailed close-up portrait of a young red-haired woman with fair freckled skin, long loose curls.\n"
        "Scene:\n"
        "Her face half-veiled by overlapping cherry blossom branches; petals brushing her temple and casting delicate, "
        "fluttering shadows over her eyes. Blossoms framing the face. Shallow DOF.\n"
        "Lighting:\n"
        "Ethereal backlighting with soft god rays filtering through pink blooms; subtle fill to keep facial detail and "
        "gentle rim to separate hair from background.\n"
        "Style:\n"
        "Natural rendering\u2014fine eyelash strands, faint capillaries on eyelids, satin lip balm highlight, micro skin "
        "texture, petals with fine veins and a soft velvety sheen, pastel bokeh background with filmic grain."
    ],
    [
        "Candid snapshot of an ordinary middle-aged French man with a long face, three-day stubble, messy "
        "light-brown hair and a slightly crooked nose, sitting alone at a small wooden table in a dimly lit Paris "
        "bistro, glass of red wine in front of him, shot with a compact point-and-shoot film camera, gentle flash "
        "falloff, documentary realism, visible grain."
    ],
    ["immense rocket launch exhaust as seen from extremely close up"],
    [
        "Subject:\n"
        "A close-up of a German Shepherd's head, black-and-tan coat, upright ears, dark almond-shaped eyes and a moist "
        "nose; only its head and shoulders are in frame. Its head is turned fully to the right of the frame, gaze fixed "
        "on a distant point far to the right, not looking at the camera, with a focused expression.\n"
        "Action:\n"
        "It sits upright among the wheat ears, ears pointed toward where it is looking, gazing motionless to the right.\n"
        "Scene:\n"
        "An endless golden wheat field with great bright cumulus clouds on the horizon; extremely shallow depth of "
        "field, the surroundings present only as blurred patches of colour and light.\n"
        "Color:\n"
        "Golden yellow and bright cloud white.\n"
        "Composition:\n"
        "Vertical close-up, the head filling two-thirds of the frame height, subject on the left third line, generous "
        "empty space on the right where it is looking; background fully out of focus.\n"
        "Medium:\n"
        "Cinematic film still.\n"
        "Medium details:\n"
        "Mamiya 7 II, 80mm f/4, Kodak Portra 160; open bright daylight breaking through the clouds; clear, luminous, "
        "airy image; soft highlights that never clip, light and airy shadows; shallow depth of field with focus on the "
        "eyes; fine film grain; nothing but the golden wheat field in frame, no buildings, vehicles or man-made objects."
    ],
    [
        "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 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."
    ],
    [
        "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 side view of a heavily modified, vintage American muscle car performing a burnout. The car is a 1968-1970 "
        "Dodge Charger, but in a state of disrepair - showing significant rust, faded paint (a mix of teal/blue and "
        "white on the roof), and missing trim. The hood is open, revealing a large, powerful engine with multiple "
        "carburetors. Thick white tire smoke is billowing from the rear tires, obscuring the lower portion of the car."
    ],
    [
        "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."
    ],
    ["film still, astronaut in a dense tropical jungle"],
    [
        "A tiny, vibrant blue cardinal perches delicately on a slender diagonal branch amid lush green lobed "
        "leaves and small round buds. The bird has soft, finely textured sapphire-blue plumage, a distinctive pointed "
        "crest, deep indigo wings, glossy black eyes, a small coral-pink conical beak, and delicate dark feet firmly "
        "gripping the wood. Its head is slightly tilted, giving it a curious and gentle expression. In this macro "
        "wildlife photograph, an extremely shallow depth of field sharply focuses on the bird's face, capturing the "
        "intricate details of its delicate feathers. The deep emerald-green background dissolves into a smooth, creamy "
        "bokeh, illuminated by soft, diffused natural lighting that enhances the rich blue tones of the plumage and the "
        "vibrant surrounding foliage."
    ],
    [
        "Subject:\n"
        "A birch wood; in the foreground two birch trees lean diagonally into the frame from the left, their bark "
        "peeling in papery black-and-white layers; beneath the trees, moss-covered boulders and a pool of still water.\n"
        "Morning mist hangs between the trees, the distant trunks fading one by one into the fog.\n"
        "A late-winter dawn, no wind, the water as still as a mirror, reflecting the trunks.\n"
        "Color:\n"
        "The yellow-green of the moss and the cold grey-white of the mist.\n"
        "Composition:\n"
        "The near birch trunks crossing the picture, the still water in the lower right.\n"
        "Medium:\n"
        "Landscape photograph.\n"
        "Soft diffused light; every layer of birch bark and every fleck of lichen sharp, the fuzz of the moss visible; "
        "the mist lowers the contrast layer by layer into the distance."
    ],
]

EXAMPLE_LABELS = [
    "Cherry blossom portrait",
    "Paris bistro",
    "Rocket exhaust",
    "Shepherd in wheat",
    "Designer toy",
    "Ligne claire beach",
    "Editorial portrait",
    "Muscle car burnout",
    "Liquid chrome",
    "Astronaut in jungle",
    "Blue Cardinal",
    "Misty birch wood",
]

S, P, L, C = "Square · 1:1", "Portrait · 2:3", "Landscape · 3:2", "Cinematic · 16:9"
EXAMPLE_FORMATS = [P, L, P, P, S, C, S, L, S, S, S, L]


def gpu_duration(prompt, steps=2, width=1024, height=1024, *args, **kwargs):
    # measured on ZeroGPU: about 1.3 s + 0.6 s per step for every megapixel, plus ~4 s on a cold start;
    # reserve about twice the cold worst case
    megapixels = int(width) * int(height) / 2**20
    return int(6 + megapixels * (3 + 1.5 * int(steps)))


@spaces.GPU(duration=gpu_duration)
def generate(prompt, steps=2, width=1024, height=1024, seed=2026, randomize=True, 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)
    if check_prompt(pipe.text_encoder, pipe.tokenizer, prompt) != "allow":
        raise gr.Error("This prompt can't be generated in this demo. Please change it and try again.")
    try:
        image = nf2_generate(pipe, prompt, int(steps), int(height), int(width), seed)
    except RuntimeError as exc:
        torch.cuda.empty_cache()
        raise gr.Error(
            f"Generation failed at {int(width)}×{int(height)}, most likely out of GPU memory. Try a 1024 format."
        ) from exc
    return image, seed


def edit_gpu_duration(prompt, command):
    # measured on the 4090: 1.1 s for a 1250-char rewrite (prompt-lookup decoding); reserve ~3x plus a cold start
    return int(8 + len(prompt) / 400)


@spaces.GPU(duration=edit_gpu_duration)
def rewrite_prompt(prompt, command):
    return edit_prompt(pipe.text_encoder, pipe.tokenizer, prompt, command)


def apply_command(prompt, command):
    """Returns the rewritten prompt and the cleared command."""
    if not prompt or not prompt.strip():
        raise gr.Error("Write a prompt first, then give a command to change it.")
    # no letters in any script ("?", "...", "123") is not an instruction: the LLM would invent an edit for it
    if not command or not any(ch.isalpha() for ch in command):
        raise gr.Error('Type a command first, e.g. "make it a film still".')
    rewritten = rewrite_prompt(prompt, command)
    # raising (not a toast + skip) keeps the chained Generate from running on a failed edit
    if not rewritten or rewritten == prompt.strip():
        raise gr.Error("The command did not change the prompt. Try putting it differently.")
    return rewritten, ""


EXAMPLE_STEPS, EXAMPLE_SEED = 2, 2026


@spaces.GPU(duration=30)
def render_example(text, width, height):
    return nf2_generate(pipe, text, EXAMPLE_STEPS, height, width, EXAMPLE_SEED)


# every example at every format, keyed by (prompt, width, height)
example_cache = {
    (text, *size): render_example(text, *size) for (text,) in EXAMPLE_PROMPTS for size in RESOLUTIONS.values()
}


def cached_example(prompt, steps, width, height, seed):
    if seed is None or int(steps) != EXAMPLE_STEPS or int(seed) != EXAMPLE_SEED:
        return None
    return example_cache.get((prompt, int(width), int(height)))


def show_example(i):
    text, fmt = EXAMPLE_PROMPTS[i][0], EXAMPLE_FORMATS[i]
    return text, fmt, EXAMPLE_STEPS, *RESOLUTIONS[fmt], EXAMPLE_SEED, example_cache[(text, *RESOLUTIONS[fmt])]


def on_resolution_change(label, prompt, steps, seed):
    w, h = RESOLUTIONS[label]
    image = cached_example(prompt, steps, w, h, seed)
    return gr.update(value=w), gr.update(value=h), gr.update() if image is None else image


STAGE_HTML = """
<div class="stage-empty">
  <span class="corner tl"></span><span class="corner tr"></span><span class="corner bl"></span><span class="corner br"></span>
  <p class="stage-title">Your image will appear here</p>
  <p class="stage-hint">Write a prompt, or pick an example below.</p>
</div>
"""


INK_1, INK_2, INK_3, INK_4 = "#F2EFE9", "#DEDAD2", "#B3AFA7", "#7A766F"
LINE = "rgba(242,239,233,0.14)"
GOLD = "#C9A86A"

# Spectral (with italics) loads here; the black background holds the page dark before Gradio's own CSS arrives
HEAD = (
    '<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Spectral:ital,wght@0,300;0,400;1,400&display=swap">'
    "<style>html, body { background: #000; }</style>"
)

theme = gr.themes.Base(
    primary_hue=gr.themes.colors.stone,
    neutral_hue=gr.themes.colors.stone,
    font=[gr.themes.GoogleFont("Work Sans", weights=(400, 500, 600)), "Helvetica Neue", "Arial", "sans-serif"],
    font_mono=["ui-monospace", "SF Mono", "Menlo", "monospace"],
    radius_size=gr.themes.sizes.radius_sm,
).set(
    body_background_fill="#000000",
    body_background_fill_dark="#000000",
    body_text_color=INK_1,
    body_text_color_dark=INK_1,
    body_text_color_subdued=INK_3,
    body_text_color_subdued_dark=INK_3,
    background_fill_primary="#000000",
    background_fill_primary_dark="#000000",
    background_fill_secondary="#000000",
    background_fill_secondary_dark="#000000",
    block_background_fill="#000000",
    block_background_fill_dark="#000000",
    block_border_color=LINE,
    block_border_color_dark=LINE,
    block_border_width="1px",
    block_label_background_fill="#000000",
    block_label_background_fill_dark="#000000",
    block_label_text_color=INK_3,
    block_label_text_color_dark=INK_3,
    block_title_text_color=INK_2,
    block_title_text_color_dark=INK_2,
    border_color_primary=LINE,
    border_color_primary_dark=LINE,
    input_background_fill="#000000",
    input_background_fill_dark="#000000",
    input_border_color=LINE,
    input_border_color_dark=LINE,
    input_border_color_focus=GOLD,
    input_border_color_focus_dark=GOLD,
    button_primary_background_fill=INK_1,
    button_primary_background_fill_dark=INK_1,
    button_primary_background_fill_hover=INK_1,
    button_primary_background_fill_hover_dark=INK_1,
    button_primary_text_color="#000000",
    button_primary_text_color_dark="#000000",
    button_primary_border_color=INK_1,
    button_primary_border_color_dark=INK_1,
    button_secondary_background_fill="#000000",
    button_secondary_background_fill_dark="#000000",
    button_secondary_background_fill_hover="rgba(242,239,233,0.06)",
    button_secondary_background_fill_hover_dark="rgba(242,239,233,0.06)",
    button_secondary_text_color=INK_1,
    button_secondary_text_color_dark=INK_1,
    button_secondary_border_color=INK_4,
    button_secondary_border_color_dark=INK_4,
    checkbox_background_color_selected=GOLD,
    checkbox_background_color_selected_dark=GOLD,
    slider_color=GOLD,
    slider_color_dark=GOLD,
    color_accent=GOLD,
    color_accent_soft="rgba(201,168,106,0.16)",
    color_accent_soft_dark="rgba(201,168,106,0.16)",
    error_background_fill="#000000",
    error_background_fill_dark="#000000",
    error_border_color=LINE,
    error_border_color_dark=LINE,
    error_text_color=INK_1,
    error_text_color_dark=INK_1,
    error_icon_color=INK_2,
    error_icon_color_dark=INK_2,
)

CSS = """
.gradio-container { background: #000 !important; }
#page { max-width: 1120px; margin: 0 auto; padding: 4px 8px 40px; }
/* one left edge for header, controls, stage and examples: the HTML blocks carry no padding of their own */
#header-block, #examples-head, #header-block .html-container, #examples-head .html-container, #stage-overlay .html-container { padding: 0 !important; margin: 0 !important; }
#header-block .prose, #examples-head .prose { max-width: none; }

#nf2-header { padding: 32px 0 32px; border-bottom: 1px solid rgba(242,239,233,0.08); margin-bottom: 32px; }
#nf2-header h1 {
  margin: 0; font: 300 clamp(2.25rem, 1.75rem + 1.4vw, 3rem)/1.05 'Spectral', Georgia, serif;
  letter-spacing: -.012em; color: #F2EFE9;
}
#nf2-header .paper { margin: 14px 0 0; max-width: 46ch; font: italic 400 1.0625rem/1.5 'Spectral', Georgia, serif; color: #B3AFA7; text-wrap: balance; }
#nf2-header .meta { margin-top: 24px; display: flex; justify-content: space-between; align-items: center; flex-wrap: wrap; gap: 16px; }
#nf2-header .tags { font: 500 11px/1 'Work Sans', sans-serif; text-transform: uppercase; letter-spacing: .18em; color: #B3AFA7; }
#nf2-header .links { display: flex; gap: 8px; flex-wrap: wrap; }
#nf2-header .pill {
  height: 32px; padding: 0 14px; border-radius: 999px; border: 1px solid rgba(242,239,233,0.14);
  display: inline-flex; align-items: center;
  font: 500 12px/1 'Work Sans', sans-serif; letter-spacing: .02em; color: #DEDAD2; text-decoration: none;
  transition: border-color 160ms ease-out, color 160ms ease-out, background-color 160ms ease-out;
}
#nf2-header .pill:hover { border-color: #DEDAD2; color: #F2EFE9; background: rgba(242,239,233,0.06); }

.sec-head { display: flex; align-items: center; gap: 16px; margin: 48px 0 4px; }
.sec-head h2 { margin: 0; font: 500 11px/1.9 'Work Sans', sans-serif; text-transform: uppercase; letter-spacing: .18em; color: #B3AFA7 !important; }
.sec-head .rule { flex: 1; height: 1px; background: rgba(242,239,233,0.14); }
#examples { display: grid !important; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 0 24px; counter-reset: ex; padding: 0; }
#examples .example {
  counter-increment: ex; position: relative; display: block; width: 100%; min-width: 0; height: auto;
  padding: 34px 0 16px; text-align: left; background: transparent !important; box-shadow: none !important;
  border: 0 !important; border-top: 1px solid rgba(242,239,233,0.14) !important; border-radius: 0 !important;
  font: 300 1.25rem/1.25 'Spectral', Georgia, serif !important; color: #DEDAD2 !important;
  transition: color 160ms ease-out, border-color 160ms ease-out;
}
#examples .example::before {
  content: counter(ex, decimal-leading-zero); position: absolute; top: 14px; left: 0;
  font: 500 11px/1 'Work Sans', sans-serif; letter-spacing: .18em; color: #7A766F; transition: color 160ms ease-out;
}
@media (hover: hover) {
  #examples .example:hover { color: #F2EFE9 !important; border-top-color: #F2EFE9 !important; }
  #examples .example:hover::before { color: #C9A86A; }
}
#examples .example:focus-visible { color: #F2EFE9 !important; }
@media (max-width: 1023px) { #examples .example { font-size: 1.125rem !important; padding: 30px 0 14px; } }
@media (max-width: 767px) { #examples { grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 0 20px; } }

#controls { gap: 24px !important; }
#controls .block, #controls .form, #advanced .block, #advanced .form {
  border: 0 !important; background: transparent !important; box-shadow: none !important; padding-left: 0 !important; padding-right: 0 !important;
}
#controls .form { gap: 24px !important; }
#controls .block, #advanced .block { overflow-x: hidden !important; }
#controls span[data-testid="block-info"] {
  font: 500 11px/1.9 'Work Sans', sans-serif !important; text-transform: uppercase; letter-spacing: .18em; color: #DEDAD2 !important;
}
#prompt textarea {
  font: 300 1.25rem/1.45 'Spectral', Georgia, serif !important; color: #F2EFE9 !important;
  background: rgba(242,239,233,0.04) !important; border: 1px solid rgba(242,239,233,0.18) !important;
  border-radius: 14px !important; box-shadow: none !important; padding: 14px 16px !important;
  transition: border-color 160ms ease-out, background-color 160ms ease-out;
}
#prompt textarea:hover { border-color: rgba(242,239,233,0.32) !important; }
#prompt textarea:focus { border-color: #C9A86A !important; background: rgba(242,239,233,0.055) !important; }
#prompt textarea::placeholder { color: #B3AFA7; font-style: italic; }
/* long prompts: grow to ~9 lines, then scroll behind a hairline scrollbar */
#prompt textarea { max-height: 277px !important; overflow-y: auto !important; scrollbar-width: thin; scrollbar-color: rgba(242,239,233,0.18) transparent; }
#prompt textarea::-webkit-scrollbar { width: 6px; }
#prompt textarea::-webkit-scrollbar-thumb { background: rgba(242,239,233,0.18); border-radius: 999px; }
#prompt textarea::-webkit-scrollbar-track { background: transparent; }
@media (max-width: 767px) { #prompt textarea { max-height: 190px !important; } }
/* command: one lighter line under Generate, a pill with Apply inside at its right; Enter applies too */
#command-wrap { gap: 0 !important; }
#command span[data-testid="block-info"] { color: #B3AFA7 !important; }
#command .input-container {
  align-items: center; height: 40px; padding: 0 5px 0 16px; border: 1px solid rgba(242,239,233,0.14); border-radius: 999px;
  transition: border-color 160ms ease-out;
}
#command .input-container:hover { border-color: rgba(242,239,233,0.32); }
#command .input-container:has(input:focus) { border-color: #C9A86A; }
#command input {
  font: 300 1.0625rem/1.4 'Spectral', Georgia, serif !important; color: #F2EFE9 !important; text-overflow: ellipsis;
  background: transparent !important; border: 0 !important; box-shadow: none !important; padding: 0 !important;
  transition: opacity 160ms ease-out;
}
#command input:focus-visible { outline: none !important; }
#command input::placeholder { color: #7A766F; font-style: italic; }
#command .submit-button {
  height: 28px; min-width: 0; margin-left: 8px; padding: 0 12px; border-radius: 999px; border: 1px solid rgba(242,239,233,0.18);
  background: transparent !important; color: #DEDAD2;
  font: 500 11px/1 'Work Sans', sans-serif; text-transform: uppercase; letter-spacing: .12em;
  transition: border-color 160ms ease-out, color 160ms ease-out, opacity 160ms ease-out, box-shadow 160ms ease-out;
}
#command .submit-button:hover { border-color: #DEDAD2; color: #F2EFE9; }
/* a typed command invites Apply: border, ink and a faint gold halo breathe (no size change); hover/focus takes over */
#command:has(input:not(:placeholder-shown)) .submit-button { animation: nf2-invite 2.4s ease-in-out infinite; }
#command:has(input:not(:placeholder-shown)) .submit-button:hover,
#command:has(input:not(:placeholder-shown)) .submit-button:focus-visible { animation: none; border-color: #DEDAD2; color: #F2EFE9; box-shadow: none; }
@keyframes nf2-invite {
  0%, 100% { border-color: rgba(242,239,233,0.18); color: #DEDAD2; box-shadow: 0 0 0 0 rgba(201,168,106,0); }
  50% { border-color: rgba(201,168,106,0.85); color: #F2EFE9; box-shadow: 0 0 10px 0 rgba(201,168,106,0.22); }
}
@media (prefers-reduced-motion: reduce) {
  #command:has(input:not(:placeholder-shown)) .submit-button { border-color: rgba(201,168,106,0.7); color: #F2EFE9; }
}
/* rewriting: Apply pulses and cannot fire again, prompt and command step back, Generate and the examples wait */
#page:has(#command [data-testid="status-tracker"]:not(.hide) .progress-text) #command .submit-button {
  opacity: .45; pointer-events: none; animation: nf2-pulse 1.8s ease-in-out infinite;
  border-color: rgba(242,239,233,0.18); color: #DEDAD2; box-shadow: none;
}
#page:has(#command [data-testid="status-tracker"]:not(.hide) .progress-text) #prompt textarea,
#page:has(#command [data-testid="status-tracker"]:not(.hide) .progress-text) #command input { opacity: .5; pointer-events: none; }
#page:has(#command [data-testid="status-tracker"]:not(.hide) .progress-text) #generate-btn { opacity: .4; pointer-events: none; }
#page:has(#command [data-testid="status-tracker"]:not(.hide) .progress-text) #examples .example { pointer-events: none; }
/* generating: Apply waits for the image */
#page:has(#result-image [data-testid="status-tracker"]:not(.hide) .progress-text) #command .submit-button {
  opacity: .4; pointer-events: none; animation: none; border-color: rgba(242,239,233,0.18); color: #DEDAD2; box-shadow: none;
}
.seg .wrap { gap: 8px !important; }
/* two rows of formats: standard on top, 2K beneath, columns aligned */
#resolution .wrap { display: grid !important; grid-template-columns: repeat(4, max-content); }
@media (max-width: 640px) { #resolution .wrap { grid-template-columns: repeat(2, max-content); } }
.seg label {
  height: 32px; padding: 0 14px !important; border: 1px solid rgba(242,239,233,0.14) !important; border-radius: 999px !important;
  background: transparent !important; box-shadow: none !important; cursor: pointer;
  font: 500 12px/1 'Work Sans', sans-serif !important; letter-spacing: .02em; color: #DEDAD2 !important;
  transition: border-color 160ms ease-out, color 160ms ease-out;
}
/* the radio stays in the tab order: visually hidden, not display:none */
.seg label input { position: absolute !important; opacity: 0 !important; width: 1px !important; height: 1px !important; margin: 0 !important; pointer-events: none; }
.seg label { position: relative; }
.seg label span { margin: 0 !important; color: inherit !important; }
.seg label:hover { border-color: #DEDAD2 !important; color: #F2EFE9 !important; }
.seg label.selected { border-color: #F2EFE9 !important; color: #F2EFE9 !important; }
.seg label:has(input:focus-visible) { outline: 2px solid #F2EFE9; outline-offset: 3px; }
/* keyboard focus: the project page's ring, everywhere except the prompt (gold border is its focus) */
.gradio-container :focus-visible { outline: 2px solid #F2EFE9 !important; outline-offset: 3px !important; }
#prompt textarea:focus-visible { outline: none !important; }
#generate-btn { height: 48px; border-radius: 999px !important; font: 500 14px/1 'Work Sans', sans-serif !important; letter-spacing: .06em; transition: opacity 160ms ease-out, filter 160ms ease-out; }
#generate-btn:hover { filter: brightness(1.06); }
/* generating: the button steps back and cannot be queued twice */
/* "running" = a live tracker that carries progress text; the error tracker carries .clear-status instead (:has() cannot nest) */
#page:has(#result-image [data-testid="status-tracker"]:not(.hide) .progress-text) #generate-btn { opacity: .4; pointer-events: none; }
#page:has(#result-image [data-testid="status-tracker"]:not(.hide) .progress-text) #examples .example { pointer-events: none; }
#controls #advanced { border-top: 1px solid rgba(242,239,233,0.14) !important; border-radius: 0 !important; padding-top: 10px !important; margin-top: 4px; }
#advanced .label-wrap { display: flex !important; align-items: center !important; justify-content: space-between; padding: 4px 2px 4px 0 !important; cursor: pointer; }
#advanced .label-wrap span { font: 500 11px/1.9 'Work Sans', sans-serif; text-transform: uppercase; letter-spacing: .18em; color: #DEDAD2; transition: color 160ms ease-out; }
#advanced .label-wrap .icon { font-size: 0 !important; transform: none !important; width: 28px; height: 28px; border: 1px solid rgba(242,239,233,0.18); border-radius: 999px; display: inline-flex; align-items: center; justify-content: center; transition: border-color 160ms ease-out; }
#advanced .label-wrap .icon::after { content: ""; width: 7px; height: 7px; border-right: 1.5px solid #DEDAD2; border-bottom: 1.5px solid #DEDAD2; transform: translateY(-2px) rotate(45deg); transition: transform 200ms ease-out, border-color 160ms ease-out; }
#advanced .label-wrap.open .icon::after { transform: translateY(2px) rotate(-135deg); }
#advanced .label-wrap:hover span { color: #F2EFE9; }
#advanced .label-wrap:hover .icon { border-color: #DEDAD2; }
#advanced .label-wrap:hover .icon::after { border-color: #F2EFE9; }
/* second tier: Advanced labels one ink step down from PROMPT / RESOLUTION / STEPS */
#advanced span[data-testid="block-info"] { color: #B3AFA7 !important; }
#advanced .tab-like-container { border: 0 !important; border-bottom: 1px solid rgba(242,239,233,0.14) !important; border-radius: 0 !important; background: transparent !important; width: 72px; }
#advanced input[type="number"] { background: transparent !important; color: #F2EFE9 !important; font: 400 12px/1 'Work Sans', sans-serif !important; text-align: right; box-shadow: none !important; }
#advanced .tab-like-container input[type="number"] { width: 100%; padding: 4px 0 !important; }
#advanced input[type="number"]:focus, #advanced input[type="number"]:focus-visible { outline: none !important; }
#advanced .tab-like-container:has(input:focus) { border-bottom-color: #C9A86A !important; }
#advanced input[type="number"]::-webkit-inner-spin-button, #advanced input[type="number"]::-webkit-outer-spin-button { -webkit-appearance: none; margin: 0; }
#advanced input[type="number"] { -moz-appearance: textfield; appearance: textfield; }
#advanced input[type="range"]:focus-visible { outline-offset: 6px !important; }
#seed { padding: 0 !important; }
#seed label { display: flex; flex-direction: column; gap: 6px; }
#seed input[type="number"] {
  width: 100%; border: 0 !important; border-bottom: 1px solid rgba(242,239,233,0.14) !important; border-radius: 0 !important;
  padding: 6px 0 !important; text-align: left; height: auto !important; transition: border-color 160ms ease-out, opacity 160ms ease-out;
}
#seed input[type="number"]:hover { border-bottom-color: rgba(242,239,233,0.32) !important; }
#seed input[type="number"]:focus { border-bottom-color: #C9A86A !important; outline: none !important; }
#seed input[type="number"]:disabled { opacity: .45; cursor: default; }
#seed input[type="number"]::-webkit-inner-spin-button, #seed input[type="number"]::-webkit-outer-spin-button { -webkit-appearance: none; margin: 0; }
#seed input[type="number"] { -moz-appearance: textfield; }
#seed-row .form { display: flex !important; flex-direction: row !important; align-items: flex-end !important; gap: 24px !important; }
#seed-row .form > .block { flex: 1 1 0 !important; min-width: 0 !important; }
#randomize { padding-bottom: 8px !important; }
#randomize label { padding: 0 !important; gap: 10px; }
#advanced input[type="range"]::-webkit-slider-runnable-track { height: 2px; background: linear-gradient(to right, #C9A86A var(--range_progress), rgba(242,239,233,0.14) var(--range_progress)); }
#advanced input[type="range"]::-webkit-slider-thumb { height: 12px; width: 12px; margin-top: -5px; background-color: #F2EFE9; box-shadow: none; }
#advanced input[type="range"]::-moz-range-track { height: 2px; background: rgba(242,239,233,0.14); }
#advanced input[type="range"]::-moz-range-progress { height: 2px; background: #C9A86A; }
#advanced input[type="range"]::-moz-range-thumb { height: 12px; width: 12px; border: 0; background: #F2EFE9; }
#advanced .min_value, #advanced .max_value { font: 400 11px 'Work Sans', sans-serif; color: #B3AFA7 !important; }
#advanced label:has(input[type="checkbox"]) span { font: 500 12px/1 'Work Sans', sans-serif; color: #DEDAD2; }
#advanced button.reset-button { display: none !important; }
/* empty frame is square, the default output shape, so a 1024² result does not resize it */
#result-image { border: 1px solid rgba(242,239,233,0.14) !important; border-radius: 0 !important; background: #000 !important; min-height: 0; aspect-ratio: 1 / 1; }
#result-image:has(img) { aspect-ratio: auto; }
#result-image:fullscreen { border: 0 !important; }
#stage { position: relative; }
/* scroll target 40px above the image: scroll-margin is ignored when the Space page scrolls the iframe */
#stage-anchor { position: absolute; top: -40px; left: 0; width: 1px; height: 1px; pointer-events: none; }
#stage-overlay { position: absolute !important; inset: 0; pointer-events: none; border: 0 !important; background: transparent !important; padding: 0 !important; margin: 0 !important; }
#stage:has(#result-image img) #stage-overlay { display: none !important; }
#stage:has(#result-image [data-testid="status-tracker"]:not(.hide) .progress-text) .stage-empty p { visibility: hidden; }
#result-image [data-testid="status-tracker"] {
  background: transparent !important; z-index: 2; display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 0;
}
#result-image [data-testid="status-tracker"] .eta-bar { display: none !important; }
#result-image:has(img) [data-testid="status-tracker"] { background: rgba(0,0,0,0.78) !important; }
#result-image [data-testid="status-tracker"] .wrap, #result-image [data-testid="status-tracker"] svg, #result-image [data-testid="status-tracker"] .loader { display: none !important; }
#result-image .progress-level { order: 1; display: flex; flex-direction: column; align-items: center; gap: 16px; width: 200px; }
#result-image .progress-level-inner { font: 300 2.75rem/1 'Spectral', Georgia, serif !important; color: #F2EFE9 !important; letter-spacing: -.01em; }
#result-image .progress-bar-wrap { width: 200px !important; height: 1px !important; background: rgba(242,239,233,0.14) !important; border: 0 !important; border-radius: 0 !important; }
#result-image .progress-bar { height: 1px !important; background: #C9A86A !important; border-radius: 0 !important; }
#result-image .progress-text {
  order: 2; position: static !important; margin-top: 16px; background: transparent !important;
  font: 500 11px/1.9 'Work Sans', sans-serif !important; text-transform: uppercase; letter-spacing: .18em; color: #DEDAD2 !important;
}
#result-image [data-testid="status-tracker"]:not(:has(.progress-level)) .progress-text { animation: nf2-pulse 1.8s ease-in-out infinite; }
@keyframes nf2-pulse { 0%, 100% { opacity: .45; } 50% { opacity: 1; } }
#result-image .empty .icon { display: none; }
/* result toolbar: quiet hairline buttons, no share */
#result-image .icon-button-wrapper { top: 12px !important; right: 12px !important; gap: 6px; padding: 0 !important; background: transparent !important; border: 0 !important; box-shadow: none !important; }
#result-image .icon-button-wrapper button[aria-label="Share"], #result-image .icon-button-wrapper button[aria-label="Clear"] { display: none !important; }
#result-image .icon-button-wrapper button {
  width: 32px; height: 32px; padding: 7px !important; border-radius: 999px !important; border: 1px solid rgba(242,239,233,0.18) !important;
  background: rgba(0,0,0,0.72) !important; color: #DEDAD2 !important; box-shadow: none !important; transition: border-color 160ms ease-out, color 160ms ease-out;
}
#result-image .icon-button-wrapper button:hover { border-color: #DEDAD2 !important; color: #F2EFE9 !important; }
#result-image .icon-button-wrapper > *::after, #result-image .icon-button-wrapper > *::before { display: none !important; }
/* error state: toast and in-frame status share the tokens, no red */
.toast-wrap .toast-body { border-radius: 0 !important; box-shadow: none !important; }
.toast-wrap .toast-title { font: 500 11px/1.9 'Work Sans', sans-serif !important; text-transform: uppercase; letter-spacing: .18em; color: #DEDAD2 !important; }
.toast-wrap .toast-message-text, .toast-wrap .toast-text { font: 400 14px/1.5 'Work Sans', sans-serif !important; color: #F2EFE9 !important; }
.toast-wrap .toast-icon, .toast-wrap .chevron { color: #B3AFA7 !important; }
.toast-wrap .toast-close { color: #B3AFA7 !important; }
.toast-wrap .timer { background: rgba(242,239,233,0.32) !important; }
/* the error pill lives in the empty frame only; over a result the toast carries the message and the image stays visible */
#result-image [data-testid="status-tracker"] .clear-status { display: none !important; }
#result-image [data-testid="status-tracker"]:has(.clear-status) { background: transparent !important; pointer-events: none; }
/* seed is an output too: its own tracker never shows */
#controls [data-testid="status-tracker"] { display: none !important; }
#result-image [data-testid="status-tracker"] .error { display: none !important; }
#result-image [data-testid="status-tracker"] .error, #result-image [data-testid="status-tracker"] span.error {
  font: 500 11px/1.9 'Work Sans', sans-serif !important; text-transform: uppercase; letter-spacing: .18em; color: #DEDAD2 !important;
  border: 1px solid rgba(242,239,233,0.18) !important; border-radius: 999px !important; background: #000 !important; padding: 4px 14px !important; box-shadow: none !important;
}
.stage-empty { position: absolute; inset: 0; display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 10px; text-align: center; padding: 24px; }
.stage-empty .corner { position: absolute; width: 22px; height: 22px; border-color: rgba(242,239,233,0.32); border-style: solid; border-width: 0; }
.stage-empty .tl { top: 18px; left: 18px; border-top-width: 1px; border-left-width: 1px; }
.stage-empty .tr { top: 18px; right: 18px; border-top-width: 1px; border-right-width: 1px; }
.stage-empty .bl { bottom: 18px; left: 18px; border-bottom-width: 1px; border-left-width: 1px; }
.stage-empty .br { bottom: 18px; right: 18px; border-bottom-width: 1px; border-right-width: 1px; }
.stage-empty .stage-title { margin: 0; font: 500 11px/1.9 'Work Sans', sans-serif; text-transform: uppercase; letter-spacing: .18em; color: #DEDAD2 !important; }
.stage-empty .stage-hint { margin: 0; font: italic 400 1.0625rem/1.5 'Spectral', Georgia, serif; color: #B3AFA7 !important; }
footer { display: none !important; }
/* Gradio scopes this CSS under .gradio-container .contain itself; a .gradio-container prefix here would never match */
@media (prefers-reduced-motion: reduce) {
  *, *::before, *::after { transition: none !important; animation: none !important; }
}
.gradio-container .prose a { color: #C9A86A; }
"""

with gr.Blocks(title="NitroFusion2") as demo:
    with gr.Column(elem_id="page"):
        gr.HTML(
            elem_id="header-block",
            value="""
            <header id="nf2-header">
              <h1>NitroFusion2</h1>
              <p class="paper">Bringing Adversarial Loss to the Text Stream of Multimodal Diffusion Transformers</p>
              <div class="meta">
                <span class="tags">One- and two-step generation</span>
                <div class="links">
                  <a class="pill" href="https://arxiv.org/abs/2611.05648" target="_blank" rel="noopener">arXiv</a>
                  <a class="pill" href="https://chendaryen.github.io/NitroFusion2.github.io/" target="_blank" rel="noopener">Project page</a>
                  <a class="pill" href="https://github.com/ChenDarYen/NitroFusion2" target="_blank" rel="noopener">GitHub</a>
                  <a class="pill" href="https://huggingface.co/ChenDY/NitroFusion2" target="_blank" rel="noopener">Model</a>
                </div>
              </div>
            </header>
            """
        )

        with gr.Row(equal_height=False):
            with gr.Column(scale=5, elem_id="controls"):
                prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe your image.", autofocus=True, elem_id="prompt")
                run = gr.Button("Generate", variant="primary", elem_id="generate-btn")
                # refine after a result: the command rewrites the prompt, then generates
                with gr.Column(elem_id="command-wrap"):
                    command = gr.Textbox(
                        label="Command", lines=1, max_lines=1, placeholder='e.g. "make it a film still" or "cat → zebra"',
                        submit_btn="Apply", elem_id="command",
                    )
                resolution = gr.Radio(list(RESOLUTIONS.keys()), value="Square · 1:1", label="Resolution", elem_id="resolution", elem_classes="seg")
                steps = gr.Radio([1, 2, 3, 4], value=2, label="Steps", elem_classes="seg")

                with gr.Accordion("Advanced", open=False, elem_id="advanced"):
                    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(elem_id="seed-row"):
                        seed = gr.Number(value=2026, minimum=0, maximum=MAX_SEED, precision=0, label="Seed", interactive=False, elem_id="seed")
                        randomize = gr.Checkbox(value=True, label="Randomize seed", elem_id="randomize")

            with gr.Column(scale=6, elem_id="stage"):
                output = gr.Image(label="Result", show_label=False, format="png", elem_id="result-image")
                gr.HTML(STAGE_HTML, elem_id="stage-overlay")

        gr.HTML('<div class="sec-head"><h2>Examples</h2><span class="rule"></span></div>', elem_id="examples-head")
        with gr.Row(elem_id="examples"):
            example_btns = [gr.Button(label, elem_classes="example") for label in EXAMPLE_LABELS]

    # autofocus is for desktops; on a phone it would raise the keyboard on load
    demo.load(None, None, None, js="() => { if (window.innerWidth < 768) setTimeout(() => { const t = document.querySelector('#prompt textarea'); t && t.blur(); }, 400); }")
    resolution.change(on_resolution_change, [resolution, prompt, steps, seed], [width, height, output])
    randomize.change(lambda r: gr.update(interactive=not r), randomize, seed)

    inputs = [prompt, steps, width, height, seed, randomize]
    # on a phone the stage is below the controls: bring it into view when a run starts
    TO_STAGE = "const s = document.getElementById('stage'); if (s) { let a = document.getElementById('stage-anchor'); if (!a) { a = document.createElement('div'); a.id = 'stage-anchor'; s.prepend(a); } a.scrollIntoView({ behavior: 'smooth', block: 'start' }); }"
    SCROLL_JS = "(...args) => { if (window.innerWidth < 768 && args[0] && args[0].trim()) { setTimeout(() => { " + TO_STAGE + " }, 250); } return args; }"
    run.click(generate, inputs, [output, seed], js=SCROLL_JS)
    # an example's image is ready at once: bring the stage to the top of the view
    STAGE_JS = "() => { " + TO_STAGE + " }"
    for i, btn in enumerate(example_btns):
        btn.click(lambda i=i: show_example(i), None, [prompt, resolution, steps, width, height, seed, output], js=STAGE_JS)
    prompt.submit(generate, inputs, [output, seed], js=SCROLL_JS)

    # command: the rewrite replaces the prompt, then the image is regenerated
    command.submit(apply_command, [prompt, command], [prompt, command]).success(generate, inputs, [output, seed], js=SCROLL_JS)

# client-side rendering: server-side rendering races the API-info fetch on Spaces and flashes an error page
demo.launch(theme=theme, css=CSS, head=HEAD, ssr_mode=False)