app.py

import hashlib
import inspect
import json
import os
from pathlib import Path
from typing import Any

import gradio as gr
import spaces
import torch
from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    DiffusionPipeline,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    HeunDiscreteScheduler,
    KDPM2AncestralDiscreteScheduler,
    KDPM2DiscreteScheduler,
    LMSDiscreteScheduler,
    UniPCMultistepScheduler,
)
from PIL import Image, ImageColor, ImageDraw, ImageFont

CATALOG_PATH = Path("loras.json")
COVER_CACHE_DIR = Path("images/auto-covers")
COVER_CACHE_DIR.mkdir(parents=True, exist_ok=True)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
HF_TOKEN = os.environ.get("HF_TOKEN")
DEFAULT_NEGATIVE = (
    "low quality, bad anatomy, bad hands, text, watermark, blurry, jpeg artifacts"
)

FAMILY_LABELS = {
    "sdxl": "SDXL",
    "sd15": "SD 1.5",
    "flux": "FLUX",
    "qwen-image": "Qwen-Image",
    "z-image": "Z-Image",
    "other": "Other",
    "all": "All",
}

FAMILY_BASE_MODELS = {
    "sdxl": "stabilityai/stable-diffusion-xl-base-1.0",
    "sd15": "runwayml/stable-diffusion-v1-5",
    "flux": "black-forest-labs/FLUX.2-klein-9B",
    "qwen-image": "Qwen/Qwen-Image-2512",
    "z-image": "Tongyi-MAI/Z-Image-Turbo",
}

FAMILY_DEFAULTS = {
    "sdxl": {"steps": 30, "cfg": 7.0, "width": 1024, "height": 1024},
    "sd15": {"steps": 28, "cfg": 7.5, "width": 768, "height": 768},
    "flux": {"steps": 30, "cfg": 3.5, "width": 1024, "height": 1024},
    "qwen-image": {"steps": 28, "cfg": 4.0, "width": 1024, "height": 1024},
    "z-image": {"steps": 28, "cfg": 4.0, "width": 1024, "height": 1024},
    "other": {"steps": 30, "cfg": 6.0, "width": 1024, "height": 1024},
}

SCHEDULER_CHOICES = [
    "Auto",
    "DPM++ 2M",
    "DPM++ 2M Karras",
    "DPM++ 2M SDE",
    "DPM++ 2M SDE Karras",
    "DPM++ SDE",
    "DPM++ SDE Karras",
    "DPM2",
    "DPM2 Karras",
    "DPM2 a",
    "DPM2 a Karras",
    "Euler",
    "Euler a",
    "Heun",
    "LMS",
    "LMS Karras",
    "DEIS",
    "UniPC",
]

SCHEDULER_MAP = {
    "DPM++ 2M": lambda cfg: DPMSolverMultistepScheduler.from_config(cfg),
    "DPM++ 2M Karras": lambda cfg: DPMSolverMultistepScheduler.from_config(cfg, use_karras_sigmas=True),
    "DPM++ 2M SDE": lambda cfg: DPMSolverMultistepScheduler.from_config(cfg, algorithm_type="sde-dpmsolver++"),
    "DPM++ 2M SDE Karras": lambda cfg: DPMSolverMultistepScheduler.from_config(cfg, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"),
    "DPM++ SDE": lambda cfg: DPMSolverSinglestepScheduler.from_config(cfg),
    "DPM++ SDE Karras": lambda cfg: DPMSolverSinglestepScheduler.from_config(cfg, use_karras_sigmas=True),
    "DPM2": lambda cfg: KDPM2DiscreteScheduler.from_config(cfg),
    "DPM2 Karras": lambda cfg: KDPM2DiscreteScheduler.from_config(cfg, use_karras_sigmas=True),
    "DPM2 a": lambda cfg: KDPM2AncestralDiscreteScheduler.from_config(cfg),
    "DPM2 a Karras": lambda cfg: KDPM2AncestralDiscreteScheduler.from_config(cfg, use_karras_sigmas=True),
    "Euler": lambda cfg: EulerDiscreteScheduler.from_config(cfg),
    "Euler a": lambda cfg: EulerAncestralDiscreteScheduler.from_config(cfg),
    "Heun": lambda cfg: HeunDiscreteScheduler.from_config(cfg),
    "LMS": lambda cfg: LMSDiscreteScheduler.from_config(cfg),
    "LMS Karras": lambda cfg: LMSDiscreteScheduler.from_config(cfg, use_karras_sigmas=True),
    "DEIS": lambda cfg: DEISMultistepScheduler.from_config(cfg),
    "UniPC": lambda cfg: UniPCMultistepScheduler.from_config(cfg),
}

def load_lora_catalog() -> list[dict[str, Any]]:
    if not CATALOG_PATH.exists():
        raise RuntimeError("loras.json not found. Run scripts/update_loras_catalog.py first.")
    with CATALOG_PATH.open("r", encoding="utf-8") as file:
        catalog = json.load(file)
    return catalog

LORAS = load_lora_catalog()
LORA_BY_REPO = {entry["repo"]: entry for entry in LORAS}

def family_to_label(family: str) -> str:
    return FAMILY_LABELS.get(family, family.upper())

def fallback_cover_path(entry: dict[str, Any]) -> str:
    key = f"{entry['repo']}::{entry.get('family', 'other')}"
    digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:16]
    output_path = COVER_CACHE_DIR / f"{digest}.png"
    if output_path.exists():
        return str(output_path)

    family = entry.get("family", "other")
    title = entry.get("title", "LoRA")

    base_colors = {
        "sdxl": ("#1d4ed8", "#38bdf8"),
        "sd15": ("#0f766e", "#2dd4bf"),
        "flux": ("#a21caf", "#f97316"),
        "qwen-image": ("#3f3f46", "#60a5fa"),
        "z-image": ("#065f46", "#22c55e"),
        "other": ("#1f2937", "#9ca3af"),
    }
    color_a, color_b = base_colors.get(family, base_colors["other"])

    width, height = 1024, 1024
    image = Image.new("RGB", (width, height), color_a)
    draw = ImageDraw.Draw(image)

    for y in range(height):
        alpha = y / max(1, height - 1)
        r1, g1, b1 = ImageColor.getrgb(color_a)
        r2, g2, b2 = ImageColor.getrgb(color_b)
        color = (
            int(r1 * (1 - alpha) + r2 * alpha),
            int(g1 * (1 - alpha) + g2 * alpha),
            int(b1 * (1 - alpha) + b2 * alpha),
        )
        draw.line([(0, y), (width, y)], fill=color)

    font_title = ImageFont.load_default()
    font_sub = ImageFont.load_default()
    draw.text((60, 120), title[:60], fill="white", font=font_title)
    draw.text((60, 180), family_to_label(family), fill="white", font=font_sub)
    draw.text((60, 860), "artificialguybr", fill="white", font=font_sub)

    image.save(output_path, format="PNG")
    return str(output_path)

def cover_for_entry(entry: dict[str, Any]) -> str:
    image_url = (entry.get("image") or "").strip()
    return image_url if image_url else fallback_cover_path(entry)

def filter_loras(family: str, search: str) -> list[dict[str, Any]]:
    term = (search or "").strip().lower()
    filtered: list[dict[str, Any]] = []
    for row in LORAS:
        if family != "all" and row.get("family") != family:
            continue
        haystack = f"{row.get('title', '')} {row.get('repo', '')} {row.get('family', '')}".lower()
        if term and term not in haystack:
            continue
        filtered.append(row)
    filtered.sort(key=lambda item: item.get("title", "").lower())
    return filtered

def family_defaults(family: str) -> dict[str, Any]:
    return FAMILY_DEFAULTS.get(family, FAMILY_DEFAULTS["other"])

def scheduler_interactive_for_family(family: str) -> bool:
    return family in {"sdxl", "sd15"}

def lora_dropdown_label(entry: dict[str, Any]) -> str:
    short_repo = entry.get("repo", "").split("/", 1)[-1]
    trigger = (entry.get("trigger_word") or "").strip()
    if trigger:
        return f"{entry['title']} · {short_repo} · trigger: {trigger}"
    return f"{entry['title']} · {short_repo}"

def selected_payload(selected_repo: str):
    if not selected_repo:
        return (
            gr.update(value=None),
            gr.update(value=""),
            gr.update(placeholder="Select a LoRA first"),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(value="Auto", interactive=False),
            gr.update(value=1.0),
        )

    selected = LORA_BY_REPO[selected_repo]
    family = selected.get("family", "other")
    defaults = family_defaults(family)
    base_model = selected.get("base_model") or FAMILY_BASE_MODELS.get(family, "")
    trigger = (selected.get("trigger_word") or "").strip()

    placeholder = f"Describe your image for {selected['title']}"
    if trigger:
        placeholder += f"  —  trigger word: «{trigger}»"

    badge_family = family_to_label(family)
    trigger_line = f"🔑 Trigger: `{trigger}`" if trigger else "🔑 No trigger word"
    info = (
        f"#### {selected['title']}\n"
        f"**`{selected_repo}`** · {badge_family}\n\n"
        f"{trigger_line}  \n"
        f"Base: `{base_model or 'Unknown'}`  \n"
        f"Weight: `{selected.get('weight_name') or 'auto'}`"
    )

    scheduler_enabled = scheduler_interactive_for_family(family)
    scheduler_value = "DPM++ 2M SDE Karras" if scheduler_enabled else "Auto"

    return (
        gr.update(value=cover_for_entry(selected)),
        info,
        gr.update(placeholder=placeholder),
        selected_repo,
        gr.update(value=defaults["steps"]),
        gr.update(value=defaults["cfg"]),
        gr.update(value=defaults["width"]),
        gr.update(value=defaults["height"]),
        gr.update(value=scheduler_value, interactive=scheduler_enabled),
        gr.update(value=1.0),
    )


def build_gallery_items(filtered: list[dict[str, Any]]) -> list[tuple[str, str]]:
    return [(cover_for_entry(item), item["repo"]) for item in filtered]


def refresh_lora_selector(family: str, search: str):
    filtered = filter_loras(family, search)
    first_repo = filtered[0]["repo"] if filtered else None
    count_text = f"**{len(filtered)}** LoRAs · {family_to_label(family)}"
    gallery_items = build_gallery_items(filtered)

    (
        preview_update,
        info_update,
        prompt_update,
        selected_repo_value,
        steps_update,
        cfg_update,
        width_update,
        height_update,
        scheduler_update,
        lora_scale_update,
    ) = selected_payload(first_repo)

    return (
        gr.update(value=gallery_items),
        count_text,
        preview_update,
        info_update,
        prompt_update,
        selected_repo_value,
        steps_update,
        cfg_update,
        width_update,
        height_update,
        scheduler_update,
        lora_scale_update,
    )


def on_gallery_select(evt: gr.SelectData, family: str, search: str):
    filtered = filter_loras(family, search)
    if evt.index >= len(filtered):
        return (gr.update(),) * 10
    repo = filtered[evt.index]["repo"]
    return selected_payload(repo)

CURRENT_PIPE: DiffusionPipeline | None = None
CURRENT_BASE_MODEL = ""
CURRENT_LOADED_REPO = ""

def pick_dtype_for_family(family: str) -> torch.dtype:
    if DEVICE != "cuda":
        return torch.float32
    if family in {"flux", "qwen-image", "z-image"}:
        return torch.bfloat16
    return torch.float16

def load_pipeline(base_model: str, family: str) -> DiffusionPipeline:
    global CURRENT_PIPE, CURRENT_BASE_MODEL, CURRENT_LOADED_REPO
    if CURRENT_PIPE is not None and CURRENT_BASE_MODEL == base_model:
        return CURRENT_PIPE

    if CURRENT_PIPE is not None:
        del CURRENT_PIPE
        CURRENT_PIPE = None
        CURRENT_BASE_MODEL = ""
        CURRENT_LOADED_REPO = ""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    dtype = pick_dtype_for_family(family)
    pipe = DiffusionPipeline.from_pretrained(
        base_model,
        torch_dtype=dtype,
        token=HF_TOKEN,
    )

    if DEVICE == "cuda":
        pipe = pipe.to("cuda")
    if hasattr(pipe, "enable_attention_slicing"):
        pipe.enable_attention_slicing()
    if hasattr(pipe, "enable_vae_slicing"):
        pipe.enable_vae_slicing()

    CURRENT_PIPE = pipe
    CURRENT_BASE_MODEL = base_model
    return pipe

def apply_scheduler(pipe: DiffusionPipeline, scheduler_name: str, family: str) -> None:
    if scheduler_name == "Auto":
        return
    if family not in {"sdxl", "sd15"}:
        return
    if not hasattr(pipe, "scheduler"):
        return
    try:
        scheduler_builder = SCHEDULER_MAP.get(scheduler_name)
        if scheduler_builder:
            pipe.scheduler = scheduler_builder(pipe.scheduler.config)
    except Exception:
        pass

def build_prompt(prompt: str, trigger_word: str) -> str:
    prompt = (prompt or "").strip()
    trigger_word = (trigger_word or "").strip()
    if not prompt and not trigger_word:
        return ""
    if prompt and trigger_word:
        return f"{prompt}, {trigger_word}"
    return prompt or trigger_word

def round_dim(value: float | int) -> int:
    dim = int(value)
    return max(256, (dim // 8) * 8)

@spaces.GPU
def run_lora(
    prompt: str,
    negative_prompt: str,
    cfg_scale: float,
    steps: int,
    selected_repo: str,
    scheduler_name: str,
    seed: int,
    width: int,
    height: int,
    lora_scale: float,
):
    global CURRENT_LOADED_REPO

    if not selected_repo:
        raise gr.Error("Select a LoRA from the list before generating.")

    selected = LORA_BY_REPO.get(selected_repo)
    if not selected:
        raise gr.Error("Selected LoRA is not in the loaded catalog.")

    family = selected.get("family", "other")
    base_model = selected.get("base_model") or FAMILY_BASE_MODELS.get(family)
    if not base_model:
        raise gr.Error(f"No base model configured for {selected_repo}.")

    full_prompt = build_prompt(prompt, selected.get("trigger_word", ""))
    if not full_prompt:
        raise gr.Error("Prompt cannot be empty.")

    pipe = load_pipeline(base_model, family)
    apply_scheduler(pipe, scheduler_name, family)

    if CURRENT_LOADED_REPO and CURRENT_LOADED_REPO != selected_repo:
        try:
            pipe.unload_lora_weights()
        except Exception:
            pass

    load_kwargs = {}
    weight_name = (selected.get("weight_name") or "").strip()
    if weight_name:
        load_kwargs["weight_name"] = weight_name
    if HF_TOKEN:
        load_kwargs["token"] = HF_TOKEN
    pipe.load_lora_weights(selected_repo, **load_kwargs)
    CURRENT_LOADED_REPO = selected_repo

    params = inspect.signature(pipe.__call__).parameters
    kwargs: dict[str, Any] = {"prompt": full_prompt}

    if "num_inference_steps" in params:
        kwargs["num_inference_steps"] = int(steps)
    if "guidance_scale" in params:
        kwargs["guidance_scale"] = float(cfg_scale)
    if "width" in params:
        kwargs["width"] = round_dim(width)
    if "height" in params:
        kwargs["height"] = round_dim(height)

    if "negative_prompt" in params and negative_prompt:
        kwargs["negative_prompt"] = negative_prompt
    elif "negative_prompt" in params and family in {"sdxl", "sd15"}:
        kwargs["negative_prompt"] = DEFAULT_NEGATIVE

    if "cross_attention_kwargs" in params:
        kwargs["cross_attention_kwargs"] = {"scale": float(lora_scale)}
    elif "joint_attention_kwargs" in params:
        kwargs["joint_attention_kwargs"] = {"scale": float(lora_scale)}
    elif "lora_scale" in params:
        kwargs["lora_scale"] = float(lora_scale)

    generator_device = "cuda" if DEVICE == "cuda" else "cpu"
    kwargs["generator"] = torch.Generator(device=generator_device).manual_seed(int(seed))

    result = pipe(**kwargs)
    return result.images[0]

CSS = """
/* ── Google Fonts ── */
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;600;700;800&family=DM+Sans:wght@300;400;500&display=swap');

/* ── Reset & root ── */
:root {
    --bg:        #0a0a0f;
    --surface:   #111118;
    --card:      #16161f;
    --border:    #2a2a3a;
    --accent:    #7c6fff;
    --accent2:   #ff6fd8;
    --accent3:   #6fffd4;
    --text:      #e8e8f0;
    --muted:     #6b6b80;
    --radius:    14px;
    --radius-sm: 8px;
    --shadow:    0 8px 32px rgba(0,0,0,0.5);
}

* { box-sizing: border-box; }

body, .gradio-container {
    background: var(--bg) !important;
    font-family: 'DM Sans', sans-serif !important;
    color: var(--text) !important;
    min-height: 100vh;
}

/* subtle grid background */
.gradio-container::before {
    content: '';
    position: fixed;
    inset: 0;
    background-image:
        linear-gradient(rgba(124,111,255,0.03) 1px, transparent 1px),
        linear-gradient(90deg, rgba(124,111,255,0.03) 1px, transparent 1px);
    background-size: 40px 40px;
    pointer-events: none;
    z-index: 0;
}

/* ── Header ── */
.app-header {
    padding: 40px 0 28px;
    text-align: center;
    position: relative;
}

.app-header h1 {
    font-family: 'Syne', sans-serif !important;
    font-size: 2.4rem !important;
    font-weight: 800 !important;
    background: linear-gradient(135deg, #fff 0%, var(--accent) 50%, var(--accent2) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0 0 8px !important;
    letter-spacing: -0.03em;
}

.app-header p {
    color: var(--muted) !important;
    font-size: 0.95rem !important;
    margin: 0 !important;
}

/* ── Steps pill bar ── */
.steps-bar {
    display: flex;
    align-items: center;
    justify-content: center;
    gap: 8px;
    margin-bottom: 32px;
    flex-wrap: wrap;
}

.step-pill {
    display: inline-flex;
    align-items: center;
    gap: 8px;
    background: var(--card);
    border: 1px solid var(--border);
    border-radius: 100px;
    padding: 7px 16px 7px 10px;
    font-size: 0.82rem;
    font-weight: 500;
    color: var(--muted);
}

.step-pill .num {
    width: 22px; height: 22px;
    border-radius: 50%;
    background: var(--accent);
    color: #fff;
    font-family: 'Syne', sans-serif;
    font-weight: 700;
    font-size: 0.75rem;
    display: grid;
    place-items: center;
}

.step-arrow { color: var(--border); font-size: 1rem; }

/* ── Panel cards ── */
.panel-card {
    background: var(--card);
    border: 1px solid var(--border);
    border-radius: var(--radius);
    padding: 24px;
}

/* ── Section labels ── */
.section-label {
    font-family: 'Syne', sans-serif !important;
    font-size: 0.7rem !important;
    font-weight: 700 !important;
    letter-spacing: 0.12em !important;
    text-transform: uppercase !important;
    color: var(--accent) !important;
    margin-bottom: 10px !important;
    display: flex;
    align-items: center;
    gap: 6px;
}

/* ── Inputs ── */
label > span {
    font-family: 'DM Sans', sans-serif !important;
    font-size: 0.78rem !important;
    font-weight: 500 !important;
    color: var(--muted) !important;
    text-transform: uppercase;
    letter-spacing: 0.06em;
    margin-bottom: 5px !important;
}

textarea, input[type="text"], .gr-input, select {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius-sm) !important;
    color: var(--text) !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 0.92rem !important;
    transition: border-color 0.2s, box-shadow 0.2s;
}

textarea:focus, input[type="text"]:focus, .gr-input:focus {
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 3px rgba(124,111,255,0.15) !important;
    outline: none !important;
}

/* ── Dropdown ── */
.gr-dropdown, [data-testid="dropdown"] {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius-sm) !important;
    color: var(--text) !important;
}

/* ── Family radio tabs ── */
.family-tabs .gr-radio-group {
    display: flex !important;
    flex-wrap: wrap !important;
    gap: 6px !important;
    background: transparent !important;
    border: none !important;
    padding: 0 !important;
}

.family-tabs .gr-radio-group label {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: 100px !important;
    padding: 6px 14px !important;
    cursor: pointer !important;
    font-size: 0.82rem !important;
    font-weight: 500 !important;
    color: var(--muted) !important;
    transition: all 0.18s !important;
    white-space: nowrap !important;
}

.family-tabs .gr-radio-group label:hover {
    border-color: var(--accent) !important;
    color: var(--text) !important;
}

.family-tabs .gr-radio-group label:has(input:checked) {
    background: var(--accent) !important;
    border-color: var(--accent) !important;
    color: #fff !important;
}

/* ── LoRA info card ── */
.lora-info-box {
    background: linear-gradient(135deg, rgba(124,111,255,0.08), rgba(255,111,216,0.05));
    border: 1px solid rgba(124,111,255,0.25);
    border-radius: var(--radius);
    padding: 16px;
    min-height: 80px;
}

.lora-info-box p, .lora-info-box h4 {
    color: var(--text) !important;
}

/* ── Preview image ── */
.lora-preview img {
    border-radius: var(--radius) !important;
    object-fit: cover !important;
    border: 1px solid var(--border) !important;
    aspect-ratio: 1 / 1;
}

/* ── Generate button ── */
.generate-btn {
    background: linear-gradient(135deg, var(--accent) 0%, var(--accent2) 100%) !important;
    border: none !important;
    border-radius: var(--radius) !important;
    color: #fff !important;
    font-family: 'Syne', sans-serif !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
    letter-spacing: 0.04em !important;
    padding: 14px 28px !important;
    cursor: pointer !important;
    width: 100% !important;
    height: 54px !important;
    transition: opacity 0.2s, transform 0.15s, box-shadow 0.2s !important;
    box-shadow: 0 4px 24px rgba(124,111,255,0.35) !important;
}

.generate-btn:hover {
    opacity: 0.9 !important;
    transform: translateY(-1px) !important;
    box-shadow: 0 8px 32px rgba(124,111,255,0.5) !important;
}

.generate-btn:active {
    transform: translateY(0) !important;
}

/* ── Result image ── */
.result-image img {
    border-radius: var(--radius) !important;
    border: 1px solid var(--border) !important;
    box-shadow: var(--shadow) !important;
    width: 100% !important;
}

/* ── Sliders ── */
.gr-slider input[type=range] {
    accent-color: var(--accent) !important;
}

/* ── Accordion ── */
.gr-accordion {
    background: var(--card) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius) !important;
    margin-top: 16px !important;
}

.gr-accordion-header {
    font-family: 'Syne', sans-serif !important;
    font-size: 0.85rem !important;
    font-weight: 700 !important;
    letter-spacing: 0.06em !important;
    color: var(--muted) !important;
    text-transform: uppercase;
    padding: 14px 20px !important;
}

.gr-accordion-header:hover {
    color: var(--text) !important;
}

/* ── Count badge ── */
.count-badge {
    display: inline-block;
    background: rgba(124,111,255,0.15);
    color: var(--accent);
    border-radius: 100px;
    padding: 3px 10px;
    font-size: 0.78rem;
    font-weight: 600;
    font-family: 'Syne', sans-serif;
}


.lora-gallery {
    border-radius: var(--radius) !important;
    overflow: hidden !important;
}

.lora-gallery .grid-wrap {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius) !important;
    padding: 8px !important;
    gap: 8px !important;
}

.lora-gallery .thumbnail-item {
    border-radius: 10px !important;
    overflow: hidden !important;
    border: 2px solid transparent !important;
    cursor: pointer !important;
    transition: border-color 0.18s, transform 0.18s, box-shadow 0.18s !important;
    position: relative;
}

.lora-gallery .thumbnail-item:hover {
    border-color: var(--accent) !important;
    transform: scale(1.03) !important;
    box-shadow: 0 4px 20px rgba(124,111,255,0.4) !important;
    z-index: 2;
}

.lora-gallery .thumbnail-item.selected {
    border-color: var(--accent2) !important;
    box-shadow: 0 0 0 3px rgba(255,111,216,0.3) !important;
}

.lora-gallery .thumbnail-item img {
    object-fit: cover !important;
    width: 100% !important;
    height: 100% !important;
    display: block !important;
}

.lora-gallery .caption-label {
    position: absolute !important;
    bottom: 0 !important;
    left: 0 !important;
    right: 0 !important;
    background: linear-gradient(transparent, rgba(0,0,0,0.85)) !important;
    color: #fff !important;
    font-size: 0.72rem !important;
    font-family: 'DM Sans', sans-serif !important;
    padding: 18px 8px 6px !important;
    text-align: center !important;
    white-space: nowrap !important;
    overflow: hidden !important;
    text-overflow: ellipsis !important;
}

/* ── Scrollbar ── */
::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: var(--surface); }
::-webkit-scrollbar-thumb { background: var(--border); border-radius: 3px; }
::-webkit-scrollbar-thumb:hover { background: var(--accent); }

/* ── Responsive ── */
@media (max-width: 768px) {
    .app-header h1 { font-size: 1.7rem !important; }
    .steps-bar { gap: 4px; }
    .step-pill { font-size: 0.75rem; padding: 6px 12px 6px 8px; }
}
"""

INITIAL_FAMILY = "sdxl"
INITIAL_FILTERED = filter_loras(INITIAL_FAMILY, "")
INITIAL_SELECTED = INITIAL_FILTERED[0]["repo"] if INITIAL_FILTERED else None

with gr.Blocks() as app:

    gr.HTML("""
    <div class="app-header">
        <h1>✦ LoRA Playground</h1>
        <p>by artificialguybr · Generate images with custom LoRA adapters</p>
    </div>
    <div class="steps-bar">
        <span class="step-pill"><span class="num">1</span>Pick model family</span>
        <span class="step-arrow">›</span>
        <span class="step-pill"><span class="num">2</span>Choose a LoRA</span>
        <span class="step-arrow">›</span>
        <span class="step-pill"><span class="num">3</span>Write your prompt</span>
        <span class="step-arrow">›</span>
        <span class="step-pill"><span class="num">4</span>Generate</span>
    </div>
    """)

    selected_repo = gr.State(INITIAL_SELECTED)

    with gr.Row(equal_height=False):

        with gr.Column(scale=1, min_width=320):

            gr.HTML('<p class="section-label">⬡ Step 1 · Model Family</p>')
            family_filter = gr.Dropdown(
                label="",
                choices=[(label, key) for key, label in FAMILY_LABELS.items() if key != "all"],
                value=INITIAL_FAMILY,
                container=False,
            )

            gr.HTML('<div style="height:16px"></div>')

            gr.HTML('<p class="section-label">◈ Step 2 · Choose LoRA</p>')
            search_box = gr.Textbox(
                label="",
                placeholder="🔍  Search by name, style, keyword…",
                container=False,
            )
            gr.HTML('<div style="height:8px"></div>')
            catalog_stats = gr.Markdown(
                f"**{len(INITIAL_FILTERED)}** LoRAs · SDXL",
                elem_classes=["count-badge"],
            )
            gr.HTML('<div style="height:6px"></div>')
            lora_gallery = gr.Gallery(
                label="",
                value=build_gallery_items(INITIAL_FILTERED),
                columns=2,
                rows=4,
                height=420,
                object_fit="cover",
                allow_preview=False,
                container=False,
                elem_classes=["lora-gallery"],
            )

            gr.HTML('<div style="height:20px"></div>')

            gr.HTML('<p class="section-label">◎ LoRA Details</p>')
            lora_preview = gr.Image(
                label="",
                height=240,
                container=False,
                elem_classes=["lora-preview"],
            )
            gr.HTML('<div style="height:10px"></div>')
            selected_info = gr.Markdown("", elem_classes=["lora-info-box"])

        with gr.Column(scale=2):

            gr.HTML('<p class="section-label">✏ Step 3 · Your Prompt</p>')
            prompt = gr.Textbox(
                label="",
                lines=3,
                placeholder="Select a LoRA to begin…",
                container=False,
            )
            gr.HTML('<div style="height:8px"></div>')
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                lines=2,
                value=DEFAULT_NEGATIVE,
            )

            gr.HTML('<div style="height:14px"></div>')

            generate_button = gr.Button(
                "✦  Generate Image",
                variant="primary",
                elem_classes=["generate-btn"],
            )

            gr.HTML('<div style="height:16px"></div>')

            result = gr.Image(
                label="",
                height=640,
                container=False,
                elem_classes=["result-image"],
            )

    with gr.Accordion("⚙  Advanced Settings", open=False):
        with gr.Row():
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=7.0)
            steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=30)
        with gr.Row():
            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=8, value=1024)
            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=8, value=1024)
        with gr.Row():
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=2**32 - 1,
                step=1,
                value=0,
                randomize=True,
            )
            lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1.5, step=0.01, value=1.0)
        scheduler = gr.Dropdown(
            label="Scheduler (SD families only)",
            choices=SCHEDULER_CHOICES,
            value="Auto",
            interactive=False,
        )

    _selector_outputs = [
        lora_gallery,
        catalog_stats,
        lora_preview,
        selected_info,
        prompt,
        selected_repo,
        steps,
        cfg_scale,
        width,
        height,
        scheduler,
        lora_scale,
    ]

    _payload_outputs = [
        lora_preview,
        selected_info,
        prompt,
        selected_repo,
        steps,
        cfg_scale,
        width,
        height,
        scheduler,
        lora_scale,
    ]

    _initial_payload = selected_payload(INITIAL_SELECTED)

    family_filter.change(fn=refresh_lora_selector, inputs=[family_filter, search_box], outputs=_selector_outputs)
    search_box.change(fn=refresh_lora_selector, inputs=[family_filter, search_box], outputs=_selector_outputs)
    lora_gallery.select(fn=on_gallery_select, inputs=[family_filter, search_box], outputs=_payload_outputs)
    app.load(fn=lambda: _initial_payload, outputs=_payload_outputs)

    generate_button.click(
        fn=run_lora,
        inputs=[prompt, negative_prompt, cfg_scale, steps, selected_repo, scheduler, seed, width, height, lora_scale],
        outputs=[result],
    )


app.queue(max_size=20)
app.launch(theme=gr.themes.Base(), css=CSS)