app.py

import os
import json
import math
import time
import random
import numpy as np
from typing import Iterable

import torch
from PIL import Image
import gradio as gr
import spaces
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from huggingface_hub import HfFileSystem, ModelCard
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# ── THEME ────────────────────────────────────────────────────────────────────

THEME_PRESETS = {
    "orange-red": {"primary": "#c2410c", "secondary": "#ea580c"},
    "violet":     {"primary": "#7c3aed", "secondary": "#8b5cf6"},
    "ocean":      {"primary": "#0369a1", "secondary": "#0ea5e9"},
    "emerald":    {"primary": "#226105", "secondary": "#73ff00"},
    "rose":       {"primary": "#be185d", "secondary": "#ec4899"},
    "amber":      {"primary": "#b45309", "secondary": "#f59e0b"},
    "cyan":       {"primary": "#0e7490", "secondary": "#06b6d4"},
}

def make_color(name: str, hex_val: str) -> colors.Color:
    """Generate a Gradio Color object from a single hex value (simplified ramp)."""
    return colors.Color(
        name=name,
        c50=hex_val + "15",
        c100=hex_val + "25",
        c200=hex_val + "40",
        c300=hex_val + "60",
        c400=hex_val + "80",
        c500=hex_val,
        c600=hex_val,
        c700=hex_val,
        c800=hex_val,
        c900=hex_val,
        c950=hex_val,
    )

class DRexTheme(Soft):
    def __init__(
        self,
        primary_hue=colors.gray,
        accent_hex: str = "#c2410c",
        *,
        neutral_hue=colors.slate,
        text_size=sizes.text_md,
        font: Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Space Grotesk"), "Arial", "sans-serif"
        ),
        font_mono: Iterable[fonts.Font | str] = (
            fonts.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        a = accent_hex
        super().set(
            # backgrounds
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="#0f0d0c",
            body_background_fill_dark="#0b0a09",
            # buttons
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill=a,
            button_primary_background_fill_hover=a,
            button_primary_background_fill_dark=a,
            button_primary_background_fill_hover_dark=a,
            button_secondary_background_fill="*primary_200",
            button_secondary_background_fill_hover="*primary_300",
            # slider
            slider_color=a,
            slider_color_dark=a,
            # block
            block_title_text_weight="600",
            block_border_width="1px",
            block_shadow="none",
            block_label_background_fill="*primary_100",
            # input
            input_background_fill="#1a1816",
            input_background_fill_dark="#1a1816",
            input_border_color="#272422",
            input_border_color_dark="#272422",
        )

drex_theme = DRexTheme(accent_hex="#c2410c")

# ── GPU / DEVICE ─────────────────────────────────────────────────────────────

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16

print(f"[D-REX] Using device: {device}")

# ── LORAS ────────────────────────────────────────────────────────────────────

loras = [
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Studio-Realism/resolve/main/images/2.png",
        "title": "Studio Realism",
        "repo": "prithivMLmods/Qwen-Image-Studio-Realism",
        "weights": "qwen-studio-realism.safetensors",
        "trigger_word": "Studio Realism",
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Sketch-Smudge/resolve/main/images/1.png",
        "title": "Sketch Smudge",
        "repo": "prithivMLmods/Qwen-Image-Sketch-Smudge",
        "weights": "qwen-sketch-smudge.safetensors",
        "trigger_word": "Sketch Smudge",
    },
    {
        "image": "https://huggingface.co/Shakker-Labs/AWPortrait-QW/resolve/main/images/08fdaf6b644b61136340d5c908ca37993e47f34cdbe2e8e8251c4c72.jpg",
        "title": "AWPortrait QW",
        "repo": "Shakker-Labs/AWPortrait-QW",
        "weights": "AWPortrait-QW_1.0.safetensors",
        "trigger_word": "Portrait",
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Anime-LoRA/resolve/main/images/1.png",
        "title": "Qwen Anime",
        "repo": "prithivMLmods/Qwen-Image-Anime-LoRA",
        "weights": "qwen-anime.safetensors",
        "trigger_word": "Qwen Anime",
    },
    {
        "image": "https://huggingface.co/flymy-ai/qwen-image-realism-lora/resolve/main/assets/flymy_realism.png",
        "title": "Image Realism",
        "repo": "flymy-ai/qwen-image-realism-lora",
        "weights": "flymy_realism.safetensors",
        "trigger_word": "Super Realism Portrait",
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Fragmented-Portraiture/resolve/main/images/3.png",
        "title": "Fragmented Portraiture",
        "repo": "prithivMLmods/Qwen-Image-Fragmented-Portraiture",
        "weights": "qwen-fragmented-portraiture.safetensors",
        "trigger_word": "Fragmented Portraiture",
    },
    {
        "image": "https://huggingface.co/prithivMLmods/Qwen-Image-Synthetic-Face/resolve/main/images/2.png",
        "title": "Synthetic Face",
        "repo": "prithivMLmods/Qwen-Image-Synthetic-Face",
        "weights": "qwen-synthetic-face.safetensors",
        "trigger_word": "Synthetic Face",
    },
    {
        "image": "https://huggingface.co/itspoidaman/qwenglitch/resolve/main/images/GyZTwJIbkAAhS4h.jpeg",
        "title": "Qwen Glitch",
        "repo": "itspoidaman/qwenglitch",
        "weights": "qwenglitch1.safetensors",
        "trigger_word": "qwenglitch",
    },
    {
        "image": "https://huggingface.co/alfredplpl/qwen-image-modern-anime-lora/resolve/main/sample1.jpg",
        "title": "Modern Anime",
        "repo": "alfredplpl/qwen-image-modern-anime-lora",
        "weights": "lora.safetensors",
        "trigger_word": "Japanese modern anime style",
    },
]

# ── MODEL INIT ────────────────────────────────────────────────────────────────

scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}

scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

# ZeroGPU: load on CPU — moved to cuda inside @spaces.GPU function only
pipe = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", scheduler=scheduler, torch_dtype=dtype
)

LIGHTNING_REPO   = "lightx2v/Qwen-Image-Lightning"
LIGHTNING_WEIGHT = "Qwen-Image-Lightning-8steps-V1.0.safetensors"
MAX_SEED         = np.iinfo(np.int32).max

ENHANCE_SUFFIXES = [
    ", ultra detailed, cinematic lighting, 8k resolution, professional photography",
    ", masterpiece, intricate details, dramatic composition, volumetric light",
    ", photorealistic, sharp focus, studio lighting, high contrast, award winning",
    ", concept art, highly detailed, trending on artstation, vivid colors, epic",
    ", hyperrealistic, golden hour lighting, bokeh, atmospheric depth, stunning",
]

# ── HELPERS ───────────────────────────────────────────────────────────────────

def aspect_to_wh(aspect: str):
    mapping = {
        "1:1": (1024, 1024), "16:9": (1152, 640), "9:16": (640, 1152),
        "4:3": (1024, 768),  "3:4": (768, 1024),  "3:2": (1024, 688),
        "2:3": (688, 1024),
    }
    return mapping.get(aspect, (1024, 1024))

def build_metadata_html(model_title: str, seed: int, steps: int, cfg: float, aspect: str) -> str:
    w, h = aspect_to_wh(aspect)
    return f"""
    <div style="
        display:grid;grid-template-columns:repeat(5,1fr);gap:8px;
        background:#131110;border:1px solid #272422;border-radius:8px;
        padding:10px 14px;font-family:'JetBrains Mono',monospace;margin-top:6px">
      <div style="text-align:center">
        <div style="font-size:13px;font-weight:600;color:var(--drex-acc,#c2410c)">{model_title.split()[0]}</div>
        <div style="font-size:9px;color:#524e4a;letter-spacing:1px;margin-top:2px">MODEL</div>
      </div>
      <div style="text-align:center">
        <div style="font-size:13px;font-weight:600;color:var(--drex-acc,#c2410c)">{seed}</div>
        <div style="font-size:9px;color:#524e4a;letter-spacing:1px;margin-top:2px">SEED</div>
      </div>
      <div style="text-align:center">
        <div style="font-size:13px;font-weight:600;color:var(--drex-acc,#c2410c)">{steps}</div>
        <div style="font-size:9px;color:#524e4a;letter-spacing:1px;margin-top:2px">STEPS</div>
      </div>
      <div style="text-align:center">
        <div style="font-size:13px;font-weight:600;color:var(--drex-acc,#c2410c)">{cfg:.1f}</div>
        <div style="font-size:9px;color:#524e4a;letter-spacing:1px;margin-top:2px">CFG</div>
      </div>
      <div style="text-align:center">
        <div style="font-size:13px;font-weight:600;color:var(--drex-acc,#c2410c)">{w}×{h}</div>
        <div style="font-size:9px;color:#524e4a;letter-spacing:1px;margin-top:2px">SIZE</div>
      </div>
    </div>"""

def format_history_html(history: list) -> str:
    if not history:
        return "<div style='font-size:11px;font-family:monospace;color:#524e4a;padding:12px;text-align:center'>no history yet</div>"
    rows = ""
    for h in reversed(history[-10:]):
        rows += f"""
        <div style="padding:6px 8px;background:#1a1816;border:1px solid #272422;
             border-radius:5px;margin-bottom:5px;cursor:pointer"
             onclick="document.querySelector('textarea').value='{h['prompt'].replace("'","")}'">
          <div style="font-size:10px;font-family:monospace;color:#9a9088;
               overflow:hidden;text-overflow:ellipsis;white-space:nowrap">{h['prompt']}</div>
          <div style="font-size:9px;font-family:monospace;color:#524e4a;margin-top:2px">{h['model']} · {h['time']}</div>
        </div>"""
    return f"<div style='max-height:180px;overflow-y:auto'>{rows}</div>"

# ── CORE LOGIC ────────────────────────────────────────────────────────────────

@spaces.GPU(duration=200)
def generate(
    prompt, neg_prompt, cfg, steps, selected_index,
    randomize_seed, seed, aspect, lora_scale, speed_mode,
    history_state,
    oauth_token: gr.OAuthToken | None = None,
    progress=gr.Progress(track_tqdm=True),
):
    if oauth_token is None:
        raise gr.Error("Please sign in with your HuggingFace account to generate images.")
    if selected_index is None:
        raise gr.Error("Select a LoRA from the gallery first.")
    if not prompt.strip():
        raise gr.Error("Write a prompt before generating.")

    lora      = loras[selected_index]
    trigger   = lora["trigger_word"]
    prompt_in = f"{trigger} {prompt}" if trigger else prompt

    # ZeroGPU: LoRA loading must happen inside GPU scope
    pipe.to("cuda")
    try:
        pipe.unload_lora_weights()

        if speed_mode == "Fast · 8 steps":
            pipe.load_lora_weights(LIGHTNING_REPO, weight_name=LIGHTNING_WEIGHT, adapter_name="lightning")
            pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style")
            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
        else:
            pipe.load_lora_weights(lora["repo"], weight_name=lora.get("weights"), low_cpu_mem_usage=True, adapter_name="style")
            pipe.set_adapters(["style"], adapter_weights=[lora_scale])

        if randomize_seed:
            seed = random.randint(0, MAX_SEED)

        w, h = aspect_to_wh(aspect)

        generator = torch.Generator(device="cuda").manual_seed(seed)
        image = pipe(
            prompt=prompt_in,
            negative_prompt=neg_prompt,
            num_inference_steps=steps,
            true_cfg_scale=cfg,
            width=w,
            height=h,
            generator=generator,
        ).images[0]
    finally:
        pipe.to("cpu")
        torch.cuda.empty_cache()

    # update history
    history_state = history_state or []
    history_state.append({
        "prompt": prompt[:80],
        "model": lora["title"],
        "time": time.strftime("%H:%M"),
    })
    history_state = history_state[-20:]

    meta_html    = build_metadata_html(lora["title"], seed, steps, cfg, aspect)
    history_html = format_history_html(history_state)

    return image, seed, meta_html, history_html, history_state

def enhance_prompt(prompt: str) -> str:
    if not prompt.strip():
        return prompt
    suffix = random.choice(ENHANCE_SUFFIXES)
    base   = prompt.rstrip(".").rstrip(",").strip()
    return base + suffix

def on_lora_select(evt: gr.SelectData, aspect):
    lora        = loras[evt.index]
    placeholder = f"Describe your image for {lora['title']}..."
    info_md     = f"### [{lora['repo']}](https://huggingface.co/{lora['repo']}) ✅"
    new_aspect  = aspect
    if "aspect" in lora:
        new_aspect = {"portrait": "9:16", "landscape": "16:9"}.get(lora["aspect"], aspect)
    return gr.update(placeholder=placeholder), info_md, evt.index, new_aspect

def on_speed_change(speed):
    if speed == "Fast · 8 steps":
        return gr.update(value="Fast · 8 steps with Lightning LoRA"), 8, 1.0
    return gr.update(value="Base · 50 steps — best quality"), 50, 4.0

def fetch_hf_lora(link: str):
    parts = link.strip("/").split("/")
    if len(parts) < 2:
        raise ValueError("Invalid repo path.")
    repo = "/".join(parts[-2:])
    card = ModelCard.load(repo)
    base = card.data.get("base_model", "")
    bases = base if isinstance(base, list) else [base]
    if not any("Qwen/Qwen-Image" in b for b in bases):
        raise ValueError("Not a Qwen-Image LoRA.")
    trigger   = card.data.get("instance_prompt", "")
    img_path  = card.data.get("widget", [{}])[0].get("output", {}).get("url")
    image_url = f"https://huggingface.co/{repo}/resolve/main/{img_path}" if img_path else None
    fs        = HfFileSystem()
    files     = fs.ls(repo, detail=False)
    weight    = next((f.split("/")[-1] for f in files if f.endswith(".safetensors")), None)
    if not weight:
        raise ValueError("No .safetensors found.")
    return parts[-1], repo, weight, trigger, image_url

def add_custom_lora(custom_text: str):
    global loras
    if not custom_text.strip():
        return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None
    try:
        link = custom_text.strip()
        if "huggingface.co/" in link:
            link = link.split("huggingface.co/")[-1]
        title, repo, weight, trigger, image = fetch_hf_lora(link)
        existing = next((i for i, l in enumerate(loras) if l["repo"] == repo), None)
        if existing is None:
            loras.append({"image": image, "title": title, "repo": repo, "weights": weight, "trigger_word": trigger})
            existing = len(loras) - 1
        card_html = f"""
        <div style="background:#1a1816;border:1px solid #272422;border-radius:8px;padding:10px;
                    display:flex;align-items:center;gap:10px;margin-top:6px">
          {'<img src="'+image+'" style="width:48px;height:48px;object-fit:cover;border-radius:5px">' if image else ''}
          <div>
            <div style="font-size:13px;font-weight:600;color:#f0ebe5">{title}</div>
            <div style="font-size:10px;font-family:monospace;color:#9a9088;margin-top:2px">
              {('trigger: <b>'+trigger+'</b>') if trigger else 'no trigger word'}
            </div>
          </div>
        </div>"""
        new_gallery = [(l["image"], l["title"]) for l in loras]
        return (gr.update(visible=True, value=card_html), gr.update(visible=True), 
                gr.update(value=new_gallery, selected_index=None), f"Custom: {weight}", existing)
    except Exception as e:
        gr.Warning(str(e))
        return gr.update(visible=True, value=f"<span style='color:#e24b4a'>{e}</span>"), gr.update(visible=True), gr.update(), "", None

def remove_custom_lora():
    gallery_reset = [(l["image"], l["title"]) for l in loras]
    return gr.update(visible=False), gr.update(visible=False), gr.update(value=gallery_reset), "", None

generate.zerogpu = True

# ── CSS ───────────────────────────────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

body, .gradio-container { background: #0b0a09 !important; font-family: 'Space Grotesk', sans-serif !important; }

/* ── HEADER ── */
#drex-header {
  text-align: center;
  padding: 20px 0 8px;
  border-bottom: 1px solid #272422;
  margin-bottom: 16px;
}
#drex-header .drex-logo {
  display: inline-flex; align-items: center; gap: 12px;
}
#drex-header .drex-mark {
  width: 38px; height: 38px; background: #c2410c; border-radius: 8px;
  display: flex; align-items: center; justify-content: center;
}
#drex-header .drex-mark svg { width: 20px; height: 20px; fill: white; }
#drex-header h1 {
  font-family: 'Space Grotesk', sans-serif !important;
  font-size: 28px !important; font-weight: 700 !important;
  color: #f0ebe5 !important; letter-spacing: -1px; margin: 0;
}
#drex-header .drex-sub {
  font-family: 'JetBrains Mono', monospace;
  font-size: 11px; color: #524e4a; margin-top: 3px;
}

/* ── THEME BAR ── */
#theme-bar { display: flex; align-items: center; justify-content: center; gap: 8px; padding: 10px 0 14px; }
.theme-dot {
  width: 20px; height: 20px; border-radius: 50%; cursor: pointer;
  border: 2px solid transparent; transition: transform .15s;
  display: inline-block;
}
.theme-dot:hover { transform: scale(1.2); }

/* ── BLOCKS ── */
.gradio-container .block {
  background: #131110 !important;
  border: 1px solid #272422 !important;
  border-radius: 10px !important;
}
label, .label-wrap span { font-family: 'JetBrains Mono', monospace !important; font-size: 10px !important; color: #524e4a !important; letter-spacing: 1px; text-transform: uppercase; }

/* ── GALLERY ── */
#lora-gallery .grid-wrap { height: 200px !important; }
#lora-gallery .thumbnail-item { border-radius: 6px !important; border: 1px solid #272422 !important; overflow: hidden; }
#lora-gallery .thumbnail-item.selected { border-color: #c2410c !important; box-shadow: 0 0 0 1px #c2410c !important; }

/* ── INPUTS ── */
textarea, input[type=text], input[type=number] {
  background: #1a1816 !important; border: 1px solid #272422 !important;
  border-radius: 6px !important; color: #f0ebe5 !important;
  font-family: 'Space Grotesk', sans-serif !important;
}
textarea:focus, input:focus { border-color: #c2410c !important; }

/* ── BUTTONS ── */
button.primary { background: #c2410c !important; border: none !important; border-radius: 6px !important; font-family: 'Space Grotesk', sans-serif !important; font-weight: 600 !important; }
button.primary:hover { filter: brightness(1.1); }
button.secondary { background: #1a1816 !important; border: 1px solid #272422 !important; border-radius: 6px !important; color: #9a9088 !important; }
button.secondary:hover { border-color: #c2410c !important; color: #c2410c !important; }

/* ── SLIDERS ── */
input[type=range] { accent-color: #c2410c; }

/* ── TABS ── */
.tab-nav button { font-family: 'JetBrains Mono', monospace !important; font-size: 11px !important; color: #524e4a !important; background: transparent !important; border: none !important; border-bottom: 2px solid transparent !important; border-radius: 0 !important; }
.tab-nav button.selected { color: #c2410c !important; border-bottom-color: #c2410c !important; }

/* ── STATUS ── */
#status-bar { font-family: 'JetBrains Mono', monospace !important; font-size: 11px !important; background: #131110 !important; border: 1px solid #272422 !important; border-radius: 6px !important; padding: 7px 12px !important; color: #9a9088 !important; }

/* ── METADATA ── */
#meta-panel { font-family: 'JetBrains Mono', monospace !important; }

/* ── GEN BTN ── */
#gen-btn { height: 44px !important; font-size: 14px !important; letter-spacing: 0.3px; }

/* ── ACCORDION ── */
.accordion { background: #131110 !important; border: 1px solid #272422 !important; }
.accordion .label-wrap { color: #9a9088 !important; }
"""

# ── JAVASCRIPT (theme switcher + prompt enhancer hint) ────────────────────────

JS_INIT = """
function() {
    // expose accent-color CSS var for metadata panel
    document.documentElement.style.setProperty('--drex-acc', '#c2410c');
}
"""

# ── GRADIO UI ─────────────────────────────────────────────────────────────────

with gr.Blocks(title="D-REX Studio") as app:
    selected_index = gr.State(value=None)
    history_state  = gr.State(value=[])

    # ── HEADER ──
    gr.HTML("""
    <div id="drex-header">
      <div class="drex-logo">
        <div class="drex-mark">
          <svg viewBox="0 0 16 16">
            <path d="M8 0L1 4v8l7 4 7-4V4L8 0zm0 2.4L13 5.5 8 8.6 3 5.5 8 2.4z
                     M2.5 6.8l4.7 2.7v5.1l-4.7-2.7V6.8zm6.3 7.8V9.5l4.7-2.7v5.1l-4.7 2.7z"/>
          </svg>
        </div>
        <div>
          <h1>D-REX</h1>
          <div class="drex-sub">LoRA Studio · Qwen-Image</div>
        </div>
      </div>
    </div>
    """)

    # ── LOGIN ──
    with gr.Row():
        gr.LoginButton(scale=0)

    # ── THEME BAR ──
    with gr.Row():
        theme_selector = gr.HTML("""
        <div id="theme-bar">
          <span style="font-size:10px;font-family:monospace;color:#524e4a;margin-right:4px">THEME</span>
          <span class="theme-dot" style="background:#c2410c" title="Orange Red"></span>
          <span class="theme-dot" style="background:#7c3aed" title="Violet"></span>
          <span class="theme-dot" style="background:#0369a1" title="Ocean"></span>
          <span class="theme-dot" style="background:#047857" title="Emerald"></span>
          <span class="theme-dot" style="background:#be185d" title="Rose"></span>
          <span class="theme-dot" style="background:#b45309" title="Amber"></span>
          <span style="font-size:10px;font-family:monospace;color:#524e4a;margin-left:8px">CUSTOM</span>
        </div>
        <div style="display:flex;justify-content:center;gap:8px;margin-bottom:12px">
          <input id="custom-theme-hex" type="text" maxlength="7" placeholder="#hex color"
            style="width:100px;background:#1a1816;border:1px solid #272422;border-radius:5px;
                   padding:4px 8px;font-size:11px;font-family:monospace;color:#f0ebe5;outline:none"/>
          <button onclick="
            const v=document.getElementById('custom-theme-hex').value.trim();
            if(/^#[0-9a-fA-F]{6}$/.test(v)){
              document.querySelectorAll('button.primary').forEach(b=>b.style.background=v);
              document.querySelectorAll('input[type=range]').forEach(r=>r.style.accentColor=v);
              document.documentElement.style.setProperty('--drex-acc',v);
            }"
            style="background:#1a1816;border:1px solid #272422;border-radius:5px;
                   padding:4px 12px;font-size:11px;font-family:monospace;color:#9a9088;cursor:pointer">
            apply
          </button>
        </div>
        """)

    # ── MAIN ──
    with gr.Row():

        # ── LEFT PANEL ──
        with gr.Column(scale=4):

            with gr.Tabs():

                with gr.Tab("Gallery"):
                    gallery = gr.Gallery(
                        value=[(l["image"], l["title"]) for l in loras],
                        label=None,
                        allow_preview=False,
                        columns=3,
                        elem_id="lora-gallery",
                        show_label=False,
                    )
                    selected_info = gr.Markdown("", elem_id="lora-info")

                with gr.Tab("Favorites"):
                    gr.Markdown("*Star a LoRA in the gallery to save it here.*",
                                elem_id="fav-placeholder")
                    fav_html = gr.HTML(
                        "<div style='font-size:11px;font-family:monospace;color:#524e4a;padding:12px;text-align:center'>no favorites yet</div>"
                    )
                    gr.Markdown(
                        "> Tip: Right-click a LoRA image → **Add to favorites** coming in next update.",
                        elem_id="fav-tip"
                    )

                with gr.Tab("History"):
                    history_html = gr.HTML(
                        "<div style='font-size:11px;font-family:monospace;color:#524e4a;padding:12px;text-align:center'>no history yet</div>",
                        elem_id="history-display",
                    )
                    clear_history_btn = gr.Button("Clear history", size="sm", variant="secondary")

                with gr.Tab("Custom LoRA"):
                    custom_lora_input = gr.Textbox(
                        label="HuggingFace repo",
                        placeholder="username/lora-model-name",
                        show_label=True,
                    )
                    gr.Markdown("[Browse Qwen-Image LoRAs →](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)")
                    custom_lora_info   = gr.HTML(visible=False)
                    custom_lora_remove = gr.Button("Remove custom LoRA", visible=False, size="sm", variant="secondary")

        # ── RIGHT PANEL ──
        with gr.Column(scale=5):

            with gr.Row():
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe your image...",
                    lines=2,
                    scale=5,
                )
                with gr.Column(scale=1, min_width=90):
                    enhance_btn = gr.Button("✦ Enhance", variant="secondary", size="sm")
                    gen_btn     = gr.Button("Generate", variant="primary", elem_id="gen-btn")

            neg_prompt = gr.Textbox(
                label="Negative prompt",
                placeholder="blur, watermark, low quality...",
                lines=1,
            )

            result = gr.Image(label="Output", format="png", elem_id="output-image")

            meta_panel = gr.HTML(
                "<div></div>",
                elem_id="meta-panel",
                visible=True,
            )

            with gr.Row():
                aspect_ratio = gr.Dropdown(
                    label="Aspect ratio",
                    choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
                    value="3:2",
                )
                speed_mode = gr.Dropdown(
                    label="Mode",
                    choices=["Base · 50 steps", "Fast · 8 steps"],
                    value="Base · 50 steps",
                )

            status_bar = gr.Textbox(
                value="D-REX ready · base · 50 steps",
                label=None,
                interactive=False,
                show_label=False,
                elem_id="status-bar",
            )

            with gr.Accordion("Advanced settings", open=False):
                with gr.Row():
                    cfg_scale = gr.Slider(
                        label="CFG Scale", minimum=1.0, maximum=5.0, step=0.1, value=4.0
                    )
                    steps = gr.Slider(
                        label="Steps", minimum=4, maximum=50, step=1, value=50
                    )
                with gr.Row():
                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
                    seed           = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                    lora_scale     = gr.Slider(label="LoRA scale", minimum=0, maximum=2, step=0.01, value=1.0)

    # ── EVENTS ───────────────────────────────────────────────────────────────

    gallery.select(
        on_lora_select,
        inputs=[aspect_ratio],
        outputs=[prompt, selected_info, selected_index, aspect_ratio],
    )

    speed_mode.change(
        on_speed_change,
        inputs=[speed_mode],
        outputs=[status_bar, steps, cfg_scale],
    )

    enhance_btn.click(
        enhance_prompt,
        inputs=[prompt],
        outputs=[prompt],
    )

    gr.on(
        triggers=[gen_btn.click, prompt.submit],
        fn=generate,
        inputs=[
            prompt, neg_prompt, cfg_scale, steps,
            selected_index, randomize_seed, seed,
            aspect_ratio, lora_scale, speed_mode,
            history_state,
        ],
        outputs=[result, seed, meta_panel, history_html, history_state],
    )
    # Note: oauth_token injected automatically by Gradio from LoginButton

    custom_lora_input.input(
        add_custom_lora,
        inputs=[custom_lora_input],
        outputs=[custom_lora_info, custom_lora_remove, gallery, selected_info, selected_index],
    )

    custom_lora_remove.click(
        remove_custom_lora,
        outputs=[custom_lora_info, custom_lora_remove, gallery, selected_info, selected_index],
    )

    clear_history_btn.click(
        lambda: ([], "<div style='font-size:11px;font-family:monospace;color:#524e4a;padding:12px;text-align:center'>history cleared</div>"),
        outputs=[history_state, history_html],
    )

# ── LAUNCH ────────────────────────────────────────────────────────────────────

app.queue()
app.launch(
    theme=drex_theme, 
    css=CSS, 
    js=JS_INIT,
)