import glob
import os
import random
import re
import subprocess
import sys
import spaces
import torch
import gradio as gr
from huggingface_hub import HfApi, ModelCard, hf_hub_download, login
from safetensors.torch import load_file

TOKEN = os.environ.get("HF_TOKEN")
if TOKEN:
    login(token=TOKEN)
api = HfApi(token=TOKEN)

from diffusers import Krea2Pipeline

DTYPE = torch.bfloat16
BASE_MODEL_REPO = "krea/Krea-2-Turbo"
# Local bucket path, or an HF model safetensors reference:
#   /data/foo.safetensors
#   krea/Krea-2-Turbo:turbo.safetensors
#   krea/Krea-2-Turbo   (auto-picks the root .safetensors weight file)
BASE_CHECKPOINT = os.environ.get(
    "BASE_CHECKPOINT", "/data/krea2TurboNSFWAIO_v10.safetensors"
)
MAX_SEED = 2**31 - 1
_REPO_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$")

_ATTN_MAP = {"wq": "to_q", "wk": "to_k", "wv": "to_v", "wo": "to_out.0", "gate": "to_gate"}
_FF_MAP = {"gate": "ff.gate", "up": "ff.up", "down": "ff.down"}

# AoTI: the Krea2TransformerBlock is served from a prebuilt .pt2 (kernels only, no
# weights) with a single rank-64 LoRA hotswap slot. Any adapter with rank <=
# TARGET_RANK targeting the same layers swaps into it with no recompilation
# (diffusers #9453). Artifact is public; the Space runs eager if it's missing.
TARGET_RANK = 64
BLOCK_NAME = "Krea2TransformerBlock"
ARTIFACT_REPO = "multimodalart/krea2-aoti-kernels"
ARTIFACT_FILE = f"Krea2TransformerBlock-lora-r{TARGET_RANK}/package.pt2"

# Built-in LoRAs to feature in the gallery, as (repo, display title). Each
# adapter's weight file and trigger word are resolved from its repo at startup,
# so adding a LoRA is just one line here.
LORA_REPOS = [
    ("krea/Krea-2-LoRA-retroanime", "Retro Anime"),
    ("krea/Krea-2-LoRA-rainywindow", "Rainy Window"),
    ("krea/Krea-2-LoRA-vintagetarot", "Vintage Tarot"),
    ("krea/Krea-2-LoRA-sunsetblur", "Sunset Blur"),
    ("krea/Krea-2-LoRA-dotmatrix", "Dot Matrix"),
    ("krea/Krea-2-LoRA-neondrip", "Neon Drip"),
    ("krea/Krea-2-LoRA-darkbrush", "Dark Brush"),
    ("krea/Krea-2-LoRA-kidsdrawing", "Kids Drawing"),
    ("krea/Krea-2-LoRA-softwatercolor", "Soft Watercolor Art Deco"),

]
DEFAULT_SCALE = 1.0


def _read_trigger(repo):
    try:
        text = ModelCard.load(repo, token=TOKEN).text
        m = re.search(r"[Tt]rigger word[:\*\s]*`([^`]+)`", text)
        if m:
            return m.group(1).strip()
    except Exception:
        pass
    return ""


def _strip_community_prefix(key):
    for prefix in ("base_model.model.", "model.diffusion_model.", "diffusion_model.", "model."):
        if key.startswith(prefix):
            key = key[len(prefix) :]
    return key


def _community_attn_suffix(block_prefix, kind, sub, tail):
    if kind == "attn":
        if sub == "qknorm":
            return None
        if sub in _ATTN_MAP:
            return f"{block_prefix}.attn.{_ATTN_MAP[sub]}.{tail}"
    if kind == "mlp" and sub in _FF_MAP:
        return f"{block_prefix}.{_FF_MAP[sub]}.{tail}"
    return None


def _convert_community_key(key, tensor):
    key = _strip_community_prefix(key)
    if key.startswith(
        ("img_in.", "transformer_blocks.", "time_embed.", "time_mod_proj.",
         "text_fusion.", "txt_in.", "final_layer.")
    ):
        return key, tensor

    parts = key.split(".")
    tail = parts[-1]

    if key == "first.weight" or key == "first.bias":
        return f"img_in.{tail}", tensor

    m = re.fullmatch(r"blocks\.(\d+)\.mod\.lin", key)
    if m:
        hidden = tensor.numel() // 6
        return f"transformer_blocks.{m.group(1)}.scale_shift_table", tensor.reshape(6, hidden)

    m = re.fullmatch(r"blocks\.(\d+)\.(prenorm|postnorm)\.scale", key)
    if m:
        norm = "norm1" if m.group(2) == "prenorm" else "norm2"
        return f"transformer_blocks.{m.group(1)}.{norm}.weight", tensor

    m = re.fullmatch(r"blocks\.(\d+)\.attn\.qknorm\.(qnorm|knorm)\.scale", key)
    if m:
        norm = "norm_q" if m.group(2) == "qnorm" else "norm_k"
        return f"transformer_blocks.{m.group(1)}.attn.{norm}.weight", tensor

    m = re.fullmatch(r"blocks\.(\d+)\.(attn|mlp)\.(\w+)\.(.+)", key)
    if m:
        converted = _community_attn_suffix(f"transformer_blocks.{m.group(1)}", m.group(2), m.group(3), m.group(4))
        if converted:
            return converted, tensor

    if key.startswith("tmlp.0."):
        return "time_embed.linear_1." + key.split(".", 2)[2], tensor
    if key.startswith("tmlp.2."):
        return "time_embed.linear_2." + key.split(".", 2)[2], tensor
    if key.startswith("tproj.1."):
        return "time_mod_proj." + key.split(".", 2)[2], tensor

    m = re.fullmatch(r"txtfusion\.(layerwise_blocks|refiner_blocks)\.(\d+)\.(prenorm|postnorm)\.scale", key)
    if m:
        norm = "norm1" if m.group(3) == "prenorm" else "norm2"
        return f"text_fusion.{m.group(1)}.{m.group(2)}.{norm}.weight", tensor

    m = re.fullmatch(r"txtfusion\.(layerwise_blocks|refiner_blocks)\.(\d+)\.attn\.qknorm\.(qnorm|knorm)\.scale", key)
    if m:
        norm = "norm_q" if m.group(3) == "qnorm" else "norm_k"
        return f"text_fusion.{m.group(1)}.{m.group(2)}.attn.{norm}.weight", tensor

    m = re.fullmatch(r"txtfusion\.(layerwise_blocks|refiner_blocks)\.(\d+)\.(attn|mlp)\.(\w+)\.(.+)", key)
    if m:
        converted = _community_attn_suffix(
            f"text_fusion.{m.group(1)}.{m.group(2)}", m.group(3), m.group(4), m.group(5)
        )
        if converted:
            return converted, tensor

    if key.startswith("txtfusion.projector."):
        return "text_fusion.projector." + key.split(".", 2)[2], tensor

    if key == "txtmlp.0.scale":
        return "txt_in.norm.weight", tensor
    if key.startswith("txtmlp.1."):
        return "txt_in.linear_1." + key.split(".", 2)[2], tensor
    if key.startswith("txtmlp.3."):
        return "txt_in.linear_2." + key.split(".", 2)[2], tensor

    if key == "last.norm.scale":
        return "final_layer.norm.weight", tensor
    if key.startswith("last.linear."):
        return "final_layer.linear." + key.split(".", 2)[2], tensor
    if key == "last.modulation.lin":
        return "final_layer.scale_shift_table", tensor

    return None, None


def _resolve_checkpoint_source(spec):
    """Return a local path to a .safetensors file (download from HF if needed)."""
    spec = (spec or "").strip()
    if not spec:
        return None
    if os.path.isfile(spec):
        return spec

    repo, _, filename = spec.partition(":")
    if _REPO_RE.match(repo):
        if not filename:
            files = api.list_repo_files(repo)
            root_safes = [f for f in files if f.endswith(".safetensors") and "/" not in f]
            if not root_safes:
                raise FileNotFoundError(f"No root .safetensors found in {repo}")
            filename = max(root_safes, key=lambda f: "turbo" in f.lower() or "raw" in f.lower())
        print(f"[base] fetching {repo}/{filename}")
        return hf_hub_download(repo, filename, token=TOKEN)

    raise FileNotFoundError(f"Checkpoint not found: {spec}")


def _load_transformer_checkpoint(checkpoint_path):
    """Load a Krea 2 Turbo .safetensors checkpoint into diffusers transformer keys."""
    raw = load_file(checkpoint_path, device="cpu")
    # Drop FP8 scale sidecars (ComfyUI) and keep only the actual weight tensors.
    raw = {k: v for k, v in raw.items() if not k.endswith(".weight_scale")}

    stripped_keys = [_strip_community_prefix(k) for k in raw]
    diffusers_native = sum(
        1 for k in stripped_keys
        if k.startswith(("img_in.", "transformer_blocks.", "time_embed.", "text_fusion.", "txt_in.", "final_layer."))
    )
    community_keys = sum(
        1 for k in stripped_keys
        if k.startswith(("blocks.", "first.", "txtfusion.", "tmlp.", "tproj.", "last.", "txtmlp."))
    )

    if diffusers_native > community_keys:
        print(f"[base] checkpoint keys look diffusers-native; loading {checkpoint_path} directly")
        return raw

    converted = {}
    skipped = []
    for key, tensor in raw.items():
        new_key, new_tensor = _convert_community_key(key, tensor)
        if new_key is None:
            skipped.append(key)
            continue
        converted[new_key] = new_tensor

    if skipped:
        print(f"[base] skipped {len(skipped)} non-transformer keys (e.g. {skipped[:3]})")
    print(f"[base] converted {len(converted)} community keys from {checkpoint_path}")
    return converted


def _load_pipeline():
    pipe = Krea2Pipeline.from_pretrained(BASE_MODEL_REPO, torch_dtype=DTYPE)
    try:
        checkpoint_path = _resolve_checkpoint_source(BASE_CHECKPOINT)
    except FileNotFoundError as e:
        print(f"[base] {e}; using hub transformer")
        return pipe
    if checkpoint_path:
        print(f"[base] swapping transformer weights from {checkpoint_path}")
        state_dict = _load_transformer_checkpoint(checkpoint_path)
        missing, unexpected = pipe.transformer.load_state_dict(state_dict, strict=False)
        if missing:
            print(f"[base] missing {len(missing)} keys (first: {missing[:5]})")
        if unexpected:
            print(f"[base] unexpected {len(unexpected)} keys (first: {unexpected[:5]})")
    return pipe


def resolve_custom_lora(repo):
    """Find the LoRA weight file and trigger word for an arbitrary Krea-2 LoRA repo."""
    files = api.list_repo_files(repo)
    safes = [f for f in files if f.endswith(".safetensors")]
    if not safes:
        raise gr.Error(f"No .safetensors weights found in {repo}.")
    # Prefer the final weights at the repo root over training checkpoints in
    # subfolders (e.g. checkpoint-1000/...), then prefer a lora-named file.
    weight_name = max(safes, key=lambda f: ("/" not in f, "lora" in f.lower()))
    return weight_name, _read_trigger(repo)


pipe = _load_pipeline()

# Resolve built-in LoRA metadata (weight file + trigger) without loading each as
# its own adapter: AoTI hotswap uses a single padded "style" slot, so styles are
# swapped into that one slot on demand. A missing/renamed repo just drops its
# tile instead of crashing the Space.
LORAS = []
for repo, title in LORA_REPOS:
    key = repo.split("LoRA-")[-1]
    try:
        weight_name, trigger = resolve_custom_lora(repo)
    except Exception as e:
        print(f"[lora] resolve failed for {repo}: {e}")
        continue
    LORAS.append(
        {"key": key, "title": title, "repo": repo,
         "weight_name": weight_name, "trigger": trigger, "scale": DEFAULT_SCALE}
    )
print(f"[lora] {len(LORAS)}/{len(LORA_REPOS)} styles ready: {[l['key'] for l in LORAS]}")

# enable_lora_hotswap MUST run before the first adapter loads — it arms PEFT
# prepare_model_for_compiled_hotswap (pads adapters to target_rank, makes the
# structure compile-safe). diffusers raises if called after the first load.
pipe.transformer.enable_lora_hotswap(target_rank=TARGET_RANK)
FIRST = LORAS[0]
pipe.transformer.load_lora_adapter(
    FIRST["repo"], weight_name=FIRST["weight_name"], adapter_name="style", token=TOKEN
)
# Bake scaling=1.0 into the exported graph; the user scale is folded into the
# re-supplied lora_B constants per request, so the slider works without recompile.
pipe.transformer.set_adapters("style", weights=1.0)
pipe.to("cuda")

# Which LoRA repo currently occupies the "style" slot (+ the scale last folded in),
# so we only hotswap / re-patch when the request actually changes them.
CURRENT = {"repo": FIRST["repo"], "weight_name": FIRST["weight_name"], "scale": None}

AOTI_MODEL = None  # set at startup if the artifact exists; else eager fallback


def _full_weights(block, scale):
    """Full constant set for a block, with the user LoRA scale folded into lora_B."""
    w = {}
    for n, p in block.named_parameters(remove_duplicate=False):
        if scale != 1.0 and ".lora_B." in n:
            w[n] = p * scale
        else:
            w[n] = p
    for n, b in block.named_buffers(remove_duplicate=False):
        w[n] = b
    return w


def _patch_all(aoti_model, scale):
    n = 0
    for block in pipe.transformer.modules():
        if block.__class__.__name__ == BLOCK_NAME:
            block.forward = aoti_model.with_weights(_full_weights(block, scale))
            n += 1
    return n


def _load_artifact():
    """Load the prebuilt LoRA-aware .pt2 (kernels only, no JIT compile at runtime)."""
    global AOTI_MODEL
    from spaces.zero.torch.aoti import LazyAOTIModel
    pt2 = hf_hub_download(
        repo_id=ARTIFACT_REPO, filename=ARTIFACT_FILE, repo_type="dataset", token=TOKEN,
    )
    AOTI_MODEL = LazyAOTIModel(pt2)
    print(f"AoTI artifact loaded: {ARTIFACT_FILE}")


try:
    _load_artifact()
except Exception as e:
    print(f"AoTI artifact unavailable ({e}); running eager.")


def _ensure_adapter(repo, weight_name):
    """Hotswap the 'style' slot to a new LoRA only when it differs from the current."""
    if CURRENT["repo"] == repo and CURRENT["weight_name"] == weight_name:
        return False
    pipe.transformer.load_lora_adapter(
        repo, weight_name=weight_name, adapter_name="style", hotswap=True, token=TOKEN,
    )
    pipe.transformer.set_adapters("style", weights=1.0)
    CURRENT["repo"] = repo
    CURRENT["weight_name"] = weight_name
    CURRENT["scale"] = None
    return True

gallery_items = []
for lora in LORAS:
    try:
        img = hf_hub_download(lora["repo"], "images/05_turbo.png")
    except Exception:
        img = None
    gallery_items.append((img, lora["title"]))


def _mash_prompt(prompt, trigger):
    prompt = (prompt or "").strip()
    if trigger and trigger.lower() not in prompt.lower():
        return f"{prompt}, {trigger}" if prompt else trigger
    return prompt


CHIP = (
    "background:#171717;border:1px solid #262626;border-radius:5px;padding:2px 8px;"
    "font-family:'JetBrains Mono',ui-monospace,monospace;font-size:12px;color:#d4d4d5;"
)


def _info_html(title, trigger, scale=None, thumb=None):
    parts = ['<div style="display:flex;flex-wrap:wrap;align-items:center;gap:8px;font-size:14px;color:#a3a3a3;">']
    if thumb:
        parts.append(
            f'<img src="{thumb}" style="width:40px;height:40px;border-radius:6px;'
            'object-fit:cover;border:1px solid #262626;" />'
        )
    parts.append(f'<span style="font-weight:600;color:#f5f5f5;">{title}</span>')
    if trigger:
        parts.append(f'<span>Trigger</span><span style="{CHIP}">{trigger}</span>')
    if scale is not None:
        parts.append(f'<span>weight</span><span style="{CHIP}">{scale}</span>')
    parts.append("</div>")
    return "".join(parts)


def _preview_path(repo):
    """Repo-relative path of the first sample image (for a browser thumbnail URL)."""
    try:
        files = api.list_repo_files(repo)
    except Exception:
        return None
    imgs = sorted(f for f in files if f.lower().endswith((".jpg", ".jpeg", ".png", ".webp")))
    return imgs[0] if imgs else None


_NO_LORA_HTML = '<div style="color:#737373;font-size:14px;">No LoRA loaded</div>'


def preview_custom_lora(repo):
    """Live feedback when a custom LoRA path is entered: resolve trigger + thumbnail."""
    repo = (repo or "").strip()
    if not repo:
        return None, _NO_LORA_HTML, gr.update(placeholder="Select a LoRA, then type a prompt")
    if not _REPO_RE.match(repo):
        # Mid-typing or an incomplete paste — stay quiet rather than flashing errors.
        return gr.update(), gr.update(), gr.update()
    try:
        _weight, trigger = resolve_custom_lora(repo)
    except Exception:
        return None, _NO_LORA_HTML, gr.update()
    img = _preview_path(repo)
    thumb = f"https://huggingface.co/{repo}/resolve/main/{img}" if img else None
    info = _info_html(repo, trigger, thumb=thumb)
    placeholder = (
        f'Type a prompt. "{trigger}" is added automatically.' if trigger else "Type a prompt"
    )
    # Clear the gallery selection so it's clear the custom LoRA is what will be used.
    return None, info, gr.update(placeholder=placeholder)


def update_selection(evt: gr.SelectData):
    lora = LORAS[evt.index]
    chip = (
        "background:#171717;border:1px solid #262626;border-radius:5px;padding:2px 8px;"
        "font-family:'JetBrains Mono',ui-monospace,monospace;font-size:12px;color:#d4d4d5;"
    )
    info = (
        '<div style="display:flex;flex-wrap:wrap;align-items:center;gap:8px;font-size:14px;color:#a3a3a3;">'
        f'<span style="font-weight:600;color:#f5f5f5;">{lora["title"]}</span>'
        f'<span>Trigger</span><span style="{chip}">{lora["trigger"]}</span>'
        f'<span>weight</span><span style="{chip}">{lora["scale"]}</span>'
        "</div>"
    )
    placeholder = f'Type a prompt. "{lora["trigger"]}" is added automatically.'
    return evt.index, info, gr.update(placeholder=placeholder), gr.update(value=lora["scale"])


@spaces.GPU(duration=75, size="large")
def _generate(repo, weight_name, full_prompt, scale, steps, guidance, width, height, seed):
    scale = float(scale)
    swapped = _ensure_adapter(repo, weight_name)
    if AOTI_MODEL is not None:
        # Re-supply constants when the adapter or scale changed; the kernel holds
        # no weights, so the fold takes effect on the next forward.
        if swapped or CURRENT["scale"] != scale:
            _patch_all(AOTI_MODEL, scale)
            CURRENT["scale"] = scale
    else:
        pipe.transformer.set_adapters("style", weights=scale)
    generator = torch.Generator("cuda").manual_seed(int(seed))
    return pipe(
        prompt=full_prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=int(width),
        height=int(height),
        generator=generator,
    ).images[0]


def run_lora(prompt, custom_lora, selected_index, lora_scale, steps, guidance, width, height, seed, randomize,
             progress=gr.Progress(track_tqdm=True)):
    if custom_lora and custom_lora.strip():
        repo = custom_lora.strip()
        weight_name, trigger = resolve_custom_lora(repo)
    elif selected_index is not None:
        lora = LORAS[selected_index]
        repo, weight_name, trigger = lora["repo"], lora["weight_name"], lora["trigger"]
    else:
        raise gr.Error("Select a LoRA from the gallery or enter a custom LoRA path.")

    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    full_prompt = _mash_prompt(prompt, trigger)
    try:
        image = _generate(repo, weight_name, full_prompt, lora_scale, steps, guidance, width, height, seed)
    except Exception as e:
        msg = str(e)
        if "hotswap" in msg.lower() or "rank" in msg.lower() or "target" in msg.lower():
            raise gr.Error(
                f"This LoRA isn't hotswap-compatible (needs rank <= {TARGET_RANK} and the same "
                f"target layers as the Krea-2 built-ins). Details: {msg}"
            )
        raise
    return image, seed


# 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 (LoRA trigger word and weight). */
.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 + """; }
"""

HEADER = """
<header id="krea-header">
  <div class="eyebrow">KREA 2 · LORA EXPLORER</div>
  <h1>LoRA the Explorer</h1>
  <p class="subtitle">Explore the built-in styles, or load any Krea 2 LoRA from Hugging Face by path, then generate on Krea 2 Turbo.</p>
  <div class="meta">
    <div class="badges">
      <span class="badge">Turbo · few-step</span>
      <span class="badge">Any Krea 2 LoRA</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.Blocks(title="Krea 2 LoRA Explorer") as demo:
    with gr.Column(elem_id="page"):
        gr.HTML(HEADER)
        selected_index = gr.State(None)
        with gr.Row():
            with gr.Column(scale=3):
                prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Select a LoRA, then type a prompt")
            with gr.Column(scale=1):
                generate_button = gr.Button("Generate", variant="primary", elem_id="generate-btn")
        with gr.Row(equal_height=False):
            with gr.Column():
                selected_info = gr.HTML("")
                gallery = gr.Gallery(
                    value=gallery_items,
                    label="Styles",
                    columns=3,
                    height="auto",
                    object_fit="cover",
                    allow_preview=False,
                    elem_id="lora-gallery",
                )
                with gr.Group():
                    custom_lora = gr.Textbox(
                        label="Custom LoRA",
                        info="Any Krea 2 LoRA Hugging Face path. Overrides the gallery selection.",
                        placeholder="username/my-krea2-lora",
                    )
                    custom_info = gr.HTML(_NO_LORA_HTML)
                    gr.Markdown(
                        "[Browse Krea 2 LoRAs](https://huggingface.co/models?other=base_model:adapter:krea/Krea-2-Turbo)"
                    )
            with gr.Column():
                result = gr.Image(label="Result", format="png", elem_id="result-image")
                with gr.Accordion("Advanced", open=False):
                    lora_scale = gr.Slider(0.0, 2.0, value=0.8, step=0.01, label="LoRA scale")
                    steps = gr.Slider(1, 30, 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, 1536, value=1024, step=16, label="Width")
                        height = gr.Slider(512, 1536, 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")

        gallery.select(update_selection, outputs=[selected_index, selected_info, prompt, lora_scale])
        custom_lora.change(preview_custom_lora, custom_lora, [selected_index, custom_info, prompt])
        custom_lora.submit(preview_custom_lora, custom_lora, [selected_index, custom_info, prompt])
        inputs = [prompt, custom_lora, selected_index, lora_scale, steps, guidance, width, height, seed, randomize]
        gr.on([generate_button.click, prompt.submit], run_lora, inputs, [result, seed])

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