Update app.py

app.py

@@ -5,91 +5,24 @@ import numpy as np
 import random
 import spaces
 import torch
+import time
+import base64
+import json
+import html as html_lib
+from io import BytesIO
 from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 from PIL import Image
 from pathlib import Path
-import threading
 from typing import Iterable
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
-
-# ── Theme ─────────────────────────────────────────────────────────────────────
-
-colors.orange_red = colors.Color(
-    name="orange_red",
-    c50="#FFF0E5",
-    c100="#FFE0CC",
-    c200="#FFC299",
-    c300="#FFA366",
-    c400="#FF8533",
-    c500="#FF4500",
-    c600="#E63E00",
-    c700="#CC3700",
-    c800="#B33000",
-    c900="#992900",
-    c950="#802200",
-)
-
-class OrangeRedTheme(Soft):
-    def __init__(
-        self,
-        *,
-        primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red,
-        neutral_hue: colors.Color | str = colors.slate,
-        text_size: sizes.Size | str = sizes.text_lg,
-        font: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
-        ),
-        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
-        ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue,
-            secondary_hue=secondary_hue,
-            neutral_hue=neutral_hue,
-            text_size=text_size,
-            font=font,
-            font_mono=font_mono,
-        )
-        super().set(
-            background_fill_primary="*primary_50",
-            background_fill_primary_dark="*primary_900",
-            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
-            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
-            button_primary_text_color="white",
-            button_primary_text_color_hover="white",
-            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_secondary_text_color="black",
-            button_secondary_text_color_hover="white",
-            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
-            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
-            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
-            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
-            slider_color="*secondary_500",
-            slider_color_dark="*secondary_600",
-            block_title_text_weight="600",
-            block_border_width="3px",
-            block_shadow="*shadow_drop_lg",
-            button_primary_shadow="*shadow_drop_lg",
-            button_large_padding="11px",
-            color_accent_soft="*primary_100",
-            block_label_background_fill="*primary_200",
-        )
-
-orange_red_theme = OrangeRedTheme()
 
 # ── Config ────────────────────────────────────────────────────────────────────
 
-dtype  = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype   = torch.bfloat16
+device  = "cuda" if torch.cuda.is_available() else "cpu"
 
 MAX_SEED       = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
+LANCZOS        = getattr(Image, "Resampling", Image).LANCZOS
 
 # ── Models ────────────────────────────────────────────────────────────────────
 
@@ -114,68 +47,166 @@ pipe_small_decoder = Flux2KleinPipeline.from_pretrained(
 )
 pipe_small_decoder.enable_model_cpu_offload()
 
+# ── Examples ──────────────────────────────────────────────────────────────────
+
+EXAMPLES_CONFIG = [
+    {
+        "images": ["examples/1.jpg"],
+        "prompt": "Change the weather to stormy.",
+    },
+    {
+        "images": ["examples/2.jpg"],
+        "prompt": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.",
+    },
+    {
+        "images": ["examples/3.jpg"],
+        "prompt": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.",
+    },
+    {
+        "images": ["examples/4.jpg"],
+        "prompt": "Make the texture high-resolution.",
+    },
+]
+
 # ── Helpers ───────────────────────────────────────────────────────────────────
 
-def update_dimensions_from_image(image_list):
-    if image_list is None or len(image_list) == 0:
-        return 1024, 1024
+def make_thumb_b64(path, max_dim=220):
+    if not os.path.exists(path):
+        return ""
+    try:
+        img = Image.open(path).convert("RGB")
+        img.thumbnail((max_dim, max_dim), LANCZOS)
+        buf = BytesIO()
+        img.save(buf, format="JPEG", quality=65)
+        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
+    except Exception as e:
+        print(f"Thumbnail error for {path}: {e}")
+        return ""
+
+
+def encode_full_image(path):
+    if not os.path.exists(path):
+        return ""
+    try:
+        with open(path, "rb") as f:
+            data = f.read()
+        ext = path.rsplit(".", 1)[-1].lower()
+        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg",
+                "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
+        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
+    except Exception as e:
+        print(f"Encode error for {path}: {e}")
+        return ""
+
+
+def build_example_cards_html():
+    cards = ""
+    for i, ex in enumerate(EXAMPLES_CONFIG):
+        thumbs_html = ""
+        for path in ex["images"]:
+            thumb = make_thumb_b64(path)
+            if thumb:
+                thumbs_html += f'<img src="{thumb}" alt="">'
+            else:
+                thumbs_html += '<div class="example-thumb-placeholder">Preview</div>'
+        n = len(ex["images"])
+        badge = f'{n} image{"s" if n > 1 else ""}'
+        prompt_short = html_lib.escape(ex["prompt"][:90])
+        if len(ex["prompt"]) > 90:
+            prompt_short += "..."
+        cards += f'''<div class="example-card" data-idx="{i}">
+            <div class="example-thumbs">{thumbs_html}</div>
+            <div class="example-meta"><span class="example-badge">{badge}</span></div>
+            <div class="example-prompt-text">{prompt_short}</div>
+        </div>'''
+    return cards
+
+
+def load_example_data(idx_str):
+    try:
+        idx = int(float(idx_str)) if idx_str and idx_str.strip() else -1
+    except (ValueError, TypeError):
+        idx = -1
+    if idx < 0 or idx >= len(EXAMPLES_CONFIG):
+        return json.dumps({"images": [], "prompt": "", "names": [], "status": "error"})
+    ex = EXAMPLES_CONFIG[idx]
+    b64_list, names = [], []
+    for path in ex["images"]:
+        b64 = encode_full_image(path)
+        if b64:
+            b64_list.append(b64)
+            names.append(os.path.basename(path))
+    return json.dumps({"images": b64_list, "prompt": ex["prompt"],
+                        "names": names, "status": "ok"})
 
-    item = image_list[0]
-    img  = item[0] if isinstance(item, tuple) else item
 
-    if isinstance(img, str):
-        img = Image.open(img).convert("RGB")
+def b64_to_pil_list(b64_json_str):
+    if not b64_json_str or b64_json_str.strip() in ("", "[]"):
+        return []
+    try:
+        b64_list = json.loads(b64_json_str)
+    except Exception:
+        return []
+    pil_images = []
+    for b64_str in b64_list:
+        if not b64_str or not isinstance(b64_str, str):
+            continue
+        try:
+            if b64_str.startswith("data:image"):
+                _, data = b64_str.split(",", 1)
+            else:
+                data = b64_str
+            image_data = base64.b64decode(data)
+            pil_images.append(Image.open(BytesIO(image_data)).convert("RGB"))
+        except Exception as e:
+            print(f"Error decoding image: {e}")
+    return pil_images
 
-    iw, ih       = img.size
-    aspect_ratio = iw / ih
 
-    if aspect_ratio >= 1:
-        new_width  = 1024
-        new_height = int(1024 / aspect_ratio)
+def update_dimensions_on_upload(pil_images):
+    if not pil_images:
+        return 1024, 1024
+    img = pil_images[0]
+    iw, ih = img.size
+    if iw >= ih:
+        nw = 1024
+        nh = int(1024 * ih / iw)
     else:
-        new_height = 1024
-        new_width  = int(1024 * aspect_ratio)
-
-    new_width  = max(256, min(1024, round(new_width  / 8) * 8))
-    new_height = max(256, min(1024, round(new_height / 8) * 8))
-    return new_width, new_height
-
-
-def parse_input_images(input_images):
-    """Safely parse gallery / filepath / PIL inputs → list[PIL.Image] or None."""
-    if input_images is None:
-        return None
-    if isinstance(input_images, str):
-        return [Image.open(input_images).convert("RGB")] if os.path.exists(input_images) else None
-    if isinstance(input_images, list) and len(input_images) > 0:
-        parsed = []
-        for item in input_images:
-            try:
-                src = item[0] if isinstance(item, tuple) else item
-                if isinstance(src, str):
-                    parsed.append(Image.open(src).convert("RGB"))
-                elif isinstance(src, Image.Image):
-                    parsed.append(src.convert("RGB"))
-                elif hasattr(src, "name"):
-                    parsed.append(Image.open(src.name).convert("RGB"))
-            except Exception as e:
-                print(f"Skipping invalid image: {e}")
-        return parsed or None
-    return None
+        nh = 1024
+        nw = int(1024 * iw / ih)
+    return max(256, (nw // 8) * 8), max(256, (nh // 8) * 8)
+
+
+def format_time(seconds: float) -> str:
+    if seconds < 60:
+        return f"{seconds:.2f}s"
+    minutes = int(seconds // 60)
+    secs    = seconds % 60
+    return f"{minutes}m {secs:.2f}s"
+
+
+def pil_to_b64(img: Image.Image) -> str:
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
 
+print("Building example thumbnails...")
+EXAMPLE_CARDS_HTML = build_example_cards_html()
+print(f"Built {len(EXAMPLES_CONFIG)} example cards.")
 
 # ── Inference ─────────────────────────────────────────────────────────────────
 
 @spaces.GPU(duration=240)
 def infer(
+    images_b64_json,
     prompt,
-    input_images=None,
-    seed=42,
-    randomize_seed=False,
-    width=1024,
-    height=1024,
-    num_inference_steps=4,
-    guidance_scale=1.0,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    num_inference_steps,
+    guidance_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
     gc.collect()
@@ -187,7 +218,11 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    image_list = parse_input_images(input_images)
+    pil_images = b64_to_pil_list(images_b64_json)
+    image_list = pil_images if pil_images else None
+
+    if image_list:
+        width, height = update_dimensions_on_upload(image_list)
 
     shared_kwargs = dict(
         prompt=prompt,
@@ -199,249 +234,1007 @@ def infer(
     if image_list is not None:
         shared_kwargs["image"] = image_list
 
-    # ── Step 1: Standard VAE ──────────────────────────────────────────────────
-    progress(0.05, desc="🟦 Running Standard VAE generation...")
-    print("Starting Standard VAE generation...")
-
-    gen_std      = torch.Generator(device="cpu").manual_seed(seed)
+    # Standard VAE
+    progress(0.05, desc="Running Standard VAE generation...")
+    t0 = time.perf_counter()
+    gen_std = torch.Generator(device="cpu").manual_seed(seed)
     out_standard = pipe_standard(**shared_kwargs, generator=gen_std).images[0]
-
-    print("Standard VAE generation complete.")
-    progress(0.55, desc="🟦 Standard VAE done — now running 🟩 Small Decoder VAE...")
+    time_std = format_time(time.perf_counter() - t0)
+    progress(0.55, desc=f"Standard VAE done in {time_std} — now running Small Decoder VAE...")
 
     gc.collect()
     torch.cuda.empty_cache()
 
-    # ── Step 2: Small Decoder VAE ─────────────────────────────────────────────
-    print("Starting Small Decoder VAE generation...")
-
+    # Small Decoder VAE
+    t0 = time.perf_counter()
     gen_small = torch.Generator(device="cpu").manual_seed(seed)
     out_small = pipe_small_decoder(**shared_kwargs, generator=gen_small).images[0]
-
-    print("Small Decoder VAE generation complete.")
-    progress(1.0, desc="✅ Both generations complete!")
+    time_small = format_time(time.perf_counter() - t0)
+    progress(1.0, desc=f"Both done! Standard: {time_std} | Small Decoder: {time_small}")
 
     gc.collect()
     torch.cuda.empty_cache()
 
-    return out_standard, out_small, seed
+    std_b64   = pil_to_b64(out_standard)
+    small_b64 = pil_to_b64(out_small)
 
+    result_json = json.dumps({
+        "standard_b64":   std_b64,
+        "small_b64":      small_b64,
+        "seed":           seed,
+        "time_std":       time_std,
+        "time_small":     time_small,
+        "status":         "ok",
+    })
 
-@spaces.GPU(duration=240)
-def infer_example(images, prompt):
-    if not images:
-        images_list = None
-    elif isinstance(images, str):
-        images_list = [images]
-    else:
-        images_list = images
-
-    out_std, out_small, seed_used = infer(
-        prompt=prompt,
-        input_images=images_list,
-        seed=0,
-        randomize_seed=True,
-        width=1024,
-        height=1024,
-        num_inference_steps=4,
-        guidance_scale=1.0,
-    )
-    return out_std, out_small, seed_used
+    return out_standard, out_small, seed, result_json
 
 
 # ── CSS ────────────────────��──────────────────────────────────────────────────
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1000px;
-}
-#main-title h1 {
-    font-size: 2.4em !important;
-}
-.vae-badge {
-    font-weight: 700;
-    font-size: 0.95em;
-    text-align: center;
-    padding: 4px 16px;
-    border-radius: 20px;
-    display: block;
-    margin-bottom: 6px;
-}
-.output-order-note {
-    text-align: center;
-    font-size: 0.85em;
-    color: #777;
-    font-style: italic;
-    margin-top: 2px;
-    margin-bottom: 4px;
+css = r"""
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=JetBrains+Mono:wght@400;500;600&display=swap');
+*{box-sizing:border-box;margin:0;padding:0}
+body,.gradio-container{
+    background:#0f0f13!important;font-family:'Inter',system-ui,-apple-system,sans-serif!important;
+    font-size:14px!important;color:#e4e4e7!important;min-height:100vh;
+}
+footer{display:none!important}
+.hidden-input{display:none!important;height:0!important;overflow:hidden!important;margin:0!important;padding:0!important}
+
+#example-load-btn{
+    position:absolute!important;left:-9999px!important;top:-9999px!important;
+    width:1px!important;height:1px!important;opacity:0.01!important;pointer-events:none!important;
+}
+#gradio-run-btn{
+    position:absolute;left:-9999px;top:-9999px;width:1px;height:1px;
+    opacity:0.01;pointer-events:none;overflow:hidden;
+}
+
+/* ── Shell ── */
+.app-shell{
+    background:#18181b;border:1px solid #27272a;border-radius:16px;
+    margin:12px auto;max-width:1400px;overflow:hidden;
+    box-shadow:0 25px 50px -12px rgba(0,0,0,.6),0 0 0 1px rgba(255,255,255,.03);
+}
+.app-header{
+    background:linear-gradient(135deg,#18181b,#1e1e24);border-bottom:1px solid #27272a;
+    padding:14px 24px;display:flex;align-items:center;justify-content:space-between;flex-wrap:wrap;gap:12px;
+}
+.app-header-left{display:flex;align-items:center;gap:12px}
+.app-logo{
+    width:36px;height:36px;background:linear-gradient(135deg,#FF6B00,#FF4500,#FF8C00);
+    border-radius:10px;display:flex;align-items:center;justify-content:center;
+    box-shadow:0 4px 12px rgba(255,69,0,.35);
+}
+.app-logo svg{width:20px;height:20px;fill:#fff;flex-shrink:0}
+.app-title{
+    font-size:18px;font-weight:700;background:linear-gradient(135deg,#e4e4e7,#a1a1aa);
+    -webkit-background-clip:text;-webkit-text-fill-color:transparent;letter-spacing:-.3px;
+}
+.app-badge{
+    font-size:11px;font-weight:600;padding:3px 10px;border-radius:20px;
+    background:rgba(255,69,0,.15);color:#FF6B35;border:1px solid rgba(255,69,0,.25);letter-spacing:.3px;
+}
+.app-badge.fast{background:rgba(34,197,94,.12);color:#4ade80;border:1px solid rgba(34,197,94,.25)}
+
+/* ── Toolbar ── */
+.app-toolbar{
+    background:#18181b;border-bottom:1px solid #27272a;padding:8px 16px;
+    display:flex;gap:4px;align-items:center;flex-wrap:wrap;
+}
+.tb-sep{width:1px;height:28px;background:#27272a;margin:0 8px}
+.modern-tb-btn{
+    display:inline-flex;align-items:center;justify-content:center;gap:6px;
+    min-width:32px;height:34px;background:transparent;border:1px solid transparent;
+    border-radius:8px;cursor:pointer;font-size:13px;font-weight:600;padding:0 12px;
+    font-family:'Inter',sans-serif;color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;
+    transition:all .15s ease;
+}
+.modern-tb-btn:hover{background:rgba(255,69,0,.15);border-color:rgba(255,69,0,.3)}
+.modern-tb-btn:active{background:rgba(255,69,0,.25);border-color:rgba(255,69,0,.45)}
+.modern-tb-btn .tb-label{font-size:13px;color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;font-weight:600}
+.modern-tb-btn .tb-svg{width:15px;height:15px;flex-shrink:0}
+.modern-tb-btn .tb-svg,.modern-tb-btn .tb-svg *{stroke:#ffffff!important;fill:none!important}
+.tb-info{font-family:'JetBrains Mono',monospace;font-size:12px;color:#71717a;padding:0 8px;display:flex;align-items:center}
+
+/* ── Layout ── */
+.app-main-row{display:flex;gap:0;flex:1;overflow:hidden}
+.app-main-left{flex:1;display:flex;flex-direction:column;min-width:0;border-right:1px solid #27272a}
+.app-main-right{width:460px;display:flex;flex-direction:column;flex-shrink:0;background:#18181b}
+
+/* ── Drop Zone ── */
+#gallery-drop-zone{position:relative;background:#09090b;min-height:300px;overflow:auto}
+#gallery-drop-zone.drag-over{outline:2px solid #FF4500;outline-offset:-2px;background:rgba(255,69,0,.04)}
+.upload-prompt-modern{position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);z-index:20}
+.upload-click-area{
+    display:flex;flex-direction:column;align-items:center;justify-content:center;
+    cursor:pointer;padding:36px 52px;border:2px dashed #3f3f46;border-radius:16px;
+    background:rgba(255,69,0,.03);transition:all .2s ease;gap:8px;
+}
+.upload-click-area:hover{background:rgba(255,69,0,.08);border-color:#FF4500;transform:scale(1.03)}
+.upload-main-text{color:#71717a;font-size:14px;font-weight:500;margin-top:4px}
+.upload-sub-text{color:#52525b;font-size:12px;text-align:center}
+
+/* ── Gallery ── */
+.image-gallery-grid{
+    display:grid;grid-template-columns:repeat(auto-fill,minmax(140px,1fr));
+    gap:12px;padding:16px;align-content:start;
+}
+.gallery-thumb{
+    position:relative;aspect-ratio:1;border-radius:10px;overflow:hidden;
+    cursor:pointer;border:2px solid #27272a;transition:all .2s ease;background:#18181b;
+}
+.gallery-thumb:hover{border-color:#3f3f46;transform:translateY(-2px);box-shadow:0 4px 12px rgba(0,0,0,.4)}
+.gallery-thumb.selected{border-color:#FF4500!important;box-shadow:0 0 0 3px rgba(255,69,0,.2)}
+.gallery-thumb img{width:100%;height:100%;object-fit:cover}
+.thumb-badge{
+    position:absolute;top:6px;left:6px;background:#FF4500;color:#fff;
+    padding:2px 8px;border-radius:4px;font-family:'JetBrains Mono',monospace;font-size:11px;font-weight:600;
+}
+.thumb-remove{
+    position:absolute;top:6px;right:6px;width:24px;height:24px;background:rgba(0,0,0,.75);
+    color:#fff;border:1px solid rgba(255,255,255,.15);border-radius:50%;cursor:pointer;
+    display:none;align-items:center;justify-content:center;font-size:12px;transition:all .15s;line-height:1;
+}
+.gallery-thumb:hover .thumb-remove{display:flex}
+.thumb-remove:hover{background:#FF4500;border-color:#FF4500}
+.gallery-add-card{
+    aspect-ratio:1;border-radius:10px;border:2px dashed #3f3f46;
+    display:flex;flex-direction:column;align-items:center;justify-content:center;
+    cursor:pointer;transition:all .2s ease;background:rgba(255,69,0,.03);gap:4px;
+}
+.gallery-add-card:hover{border-color:#FF4500;background:rgba(255,69,0,.08)}
+.gallery-add-card .add-icon{font-size:28px;color:#71717a;font-weight:300}
+.gallery-add-card .add-text{font-size:12px;color:#71717a;font-weight:500}
+
+/* ── Hint Bar ── */
+.hint-bar{
+    background:rgba(255,69,0,.06);border-top:1px solid #27272a;border-bottom:1px solid #27272a;
+    padding:10px 20px;font-size:13px;color:#a1a1aa;line-height:1.7;
+}
+.hint-bar b{color:#FF8C55;font-weight:600}
+.hint-bar kbd{
+    display:inline-block;padding:1px 6px;background:#27272a;border:1px solid #3f3f46;
+    border-radius:4px;font-family:'JetBrains Mono',monospace;font-size:11px;color:#a1a1aa;
+}
+
+/* ── Suggestions ── */
+.suggestions-section{border-top:1px solid #27272a;padding:12px 16px}
+.suggestions-title,.examples-title{
+    font-size:12px;font-weight:600;color:#71717a;text-transform:uppercase;
+    letter-spacing:.8px;margin-bottom:10px;
+}
+.suggestions-wrap{display:flex;flex-wrap:wrap;gap:6px}
+.suggestion-chip{
+    display:inline-flex;align-items:center;gap:4px;padding:5px 12px;
+    background:rgba(255,69,0,.08);border:1px solid rgba(255,69,0,.2);border-radius:20px;
+    color:#FF8C55;font-size:12px;font-weight:500;font-family:'Inter',sans-serif;
+    cursor:pointer;transition:all .15s;white-space:nowrap;
+}
+.suggestion-chip:hover{background:rgba(255,69,0,.15);border-color:rgba(255,69,0,.35);color:#FF6B35;transform:translateY(-1px)}
+
+/* ── Examples ── */
+.examples-section{border-top:1px solid #27272a;padding:12px 16px}
+.examples-scroll{display:flex;gap:10px;overflow-x:auto;padding-bottom:8px}
+.examples-scroll::-webkit-scrollbar{height:6px}
+.examples-scroll::-webkit-scrollbar-track{background:#09090b;border-radius:3px}
+.examples-scroll::-webkit-scrollbar-thumb{background:#27272a;border-radius:3px}
+.examples-scroll::-webkit-scrollbar-thumb:hover{background:#3f3f46}
+.example-card{
+    flex-shrink:0;width:210px;background:#09090b;border:1px solid #27272a;
+    border-radius:10px;overflow:hidden;cursor:pointer;transition:all .2s ease;
+}
+.example-card:hover{border-color:#FF4500;transform:translateY(-2px);box-shadow:0 4px 12px rgba(255,69,0,.15)}
+.example-card.loading{opacity:.5;pointer-events:none}
+.example-thumbs{display:flex;height:110px;overflow:hidden;background:#18181b}
+.example-thumbs img{flex:1;object-fit:cover;min-width:0;border-bottom:1px solid #27272a}
+.example-thumb-placeholder{
+    flex:1;display:flex;align-items:center;justify-content:center;
+    background:#18181b;color:#3f3f46;font-size:11px;min-width:0;
+}
+.example-meta{padding:6px 10px;display:flex;align-items:center;gap:6px}
+.example-badge{
+    display:inline-flex;padding:2px 7px;background:rgba(255,69,0,.1);border-radius:4px;
+    font-size:10px;font-weight:600;color:#FF6B35;font-family:'JetBrains Mono',monospace;white-space:nowrap;
+}
+.example-prompt-text{
+    padding:0 10px 8px;font-size:11px;color:#a1a1aa;line-height:1.4;
+    display:-webkit-box;-webkit-line-clamp:2;-webkit-box-orient:vertical;overflow:hidden;
+}
+
+/* ── Panel Cards ── */
+.panel-card{border-bottom:1px solid #27272a}
+.panel-card-title{
+    padding:12px 20px;font-size:12px;font-weight:600;color:#71717a;
+    text-transform:uppercase;letter-spacing:.8px;border-bottom:1px solid rgba(39,39,42,.6);
+}
+.panel-card-body{padding:16px 20px;display:flex;flex-direction:column;gap:8px}
+.modern-label{font-size:13px;font-weight:500;color:#a1a1aa;margin-bottom:4px;display:block}
+.modern-textarea{
+    width:100%;background:#09090b;border:1px solid #27272a;border-radius:8px;
+    padding:10px 14px;font-family:'Inter',sans-serif;font-size:14px;color:#e4e4e7;
+    resize:vertical;outline:none;min-height:42px;transition:border-color .2s;
+}
+.modern-textarea:focus{border-color:#FF4500;box-shadow:0 0 0 3px rgba(255,69,0,.15)}
+.modern-textarea::placeholder{color:#3f3f46}
+.modern-textarea.error-flash{
+    border-color:#ef4444!important;box-shadow:0 0 0 3px rgba(239,68,68,.2)!important;
+    animation:shake .4s ease;
+}
+@keyframes shake{0%,100%{transform:translateX(0)}20%,60%{transform:translateX(-4px)}40%,80%{transform:translateX(4px)}}
+
+/* ── Toast ── */
+.toast-notification{
+    position:fixed;top:24px;left:50%;transform:translateX(-50%) translateY(-120%);
+    z-index:9999;padding:10px 24px;border-radius:10px;font-family:'Inter',sans-serif;
+    font-size:14px;font-weight:600;display:flex;align-items:center;gap:8px;
+    box-shadow:0 8px 24px rgba(0,0,0,.5);
+    transition:transform .35s cubic-bezier(.34,1.56,.64,1),opacity .35s ease;opacity:0;pointer-events:none;
+}
+.toast-notification.visible{transform:translateX(-50%) translateY(0);opacity:1}
+.toast-notification.error{background:linear-gradient(135deg,#dc2626,#b91c1c);color:#fff;border:1px solid rgba(255,255,255,.15)}
+.toast-notification.warning{background:linear-gradient(135deg,#d97706,#b45309);color:#fff;border:1px solid rgba(255,255,255,.15)}
+.toast-notification.info{background:linear-gradient(135deg,#FF4500,#CC3700);color:#fff;border:1px solid rgba(255,255,255,.15)}
+
+/* ── Run Button ── */
+.btn-run{
+    display:flex;align-items:center;justify-content:center;gap:8px;width:100%;
+    background:linear-gradient(135deg,#FF4500,#CC3700);border:none;border-radius:10px;
+    padding:12px 24px;cursor:pointer;font-size:15px;font-weight:600;font-family:'Inter',sans-serif;
+    color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;transition:all .2s ease;
+    box-shadow:0 4px 16px rgba(255,69,0,.3),inset 0 1px 0 rgba(255,255,255,.1);
+}
+.btn-run:hover{
+    background:linear-gradient(135deg,#FF6B35,#FF4500);transform:translateY(-1px);
+    box-shadow:0 6px 24px rgba(255,69,0,.45),inset 0 1px 0 rgba(255,255,255,.15);
+}
+.btn-run:active{transform:translateY(0)}
+.btn-run svg{width:18px;height:18px;fill:#ffffff!important}
+
+/* ── Output Frame ── */
+.output-section{padding:16px 20px;display:flex;flex-direction:column;gap:12px;flex:1}
+.output-section-title{
+    font-size:12px;font-weight:600;color:#71717a;text-transform:uppercase;letter-spacing:.8px;
+}
+.output-compare-row{display:flex;gap:12px;flex:1}
+.output-col{flex:1;display:flex;flex-direction:column;gap:8px;min-width:0}
+.output-col-label{
+    display:flex;align-items:center;justify-content:space-between;
+    padding:6px 10px;border-radius:8px;font-size:12px;font-weight:600;
+}
+.output-col-label.std{background:rgba(30,100,255,.1);color:#6BA3FF;border:1px solid rgba(30,100,255,.2)}
+.output-col-label.small{background:rgba(34,197,94,.1);color:#4ade80;border:1px solid rgba(34,197,94,.2)}
+.output-img-wrap{
+    flex:1;background:#09090b;border:1px solid #27272a;border-radius:10px;
+    overflow:hidden;min-height:220px;display:flex;align-items:center;justify-content:center;
+    position:relative;
+}
+.output-img-wrap img{max-width:100%;max-height:340px;object-fit:contain;display:block}
+.output-placeholder{color:#3f3f46;font-size:12px;text-align:center;padding:20px}
+.timing-chip{
+    display:inline-flex;align-items:center;gap:4px;padding:3px 10px;border-radius:20px;
+    font-family:'JetBrains Mono',monospace;font-size:11px;font-weight:600;
+    background:#18181b;border:1px solid #27272a;color:#71717a;
+}
+.timing-chip.done{color:#4ade80;border-color:rgba(34,197,94,.3);background:rgba(34,197,94,.08)}
+.out-download-btn{
+    display:none;align-items:center;justify-content:center;background:rgba(255,69,0,.1);
+    border:1px solid rgba(255,69,0,.2);border-radius:6px;cursor:pointer;padding:3px 10px;
+    font-size:11px;font-weight:500;color:#FF8C55!important;gap:4px;height:24px;transition:all .15s;
+}
+.out-download-btn:hover{background:rgba(255,69,0,.2);border-color:rgba(255,69,0,.35);color:#fff!important}
+.out-download-btn.visible{display:inline-flex}
+.out-download-btn svg{width:12px;height:12px;fill:#FF8C55}
+
+/* ── Loader ── */
+.modern-loader{
+    position:absolute;top:0;left:0;right:0;bottom:0;background:rgba(9,9,11,.92);
+    z-index:15;flex-direction:column;align-items:center;justify-content:center;gap:12px;
+    backdrop-filter:blur(4px);display:none;
+}
+.modern-loader.active{display:flex}
+.loader-spinner{
+    width:32px;height:32px;border:3px solid #27272a;border-top-color:#FF4500;
+    border-radius:50%;animation:spin .8s linear infinite;
+}
+@keyframes spin{to{transform:rotate(360deg)}}
+.loader-text{font-size:12px;color:#a1a1aa;font-weight:500}
+.loader-bar-track{width:180px;height:4px;background:#27272a;border-radius:2px;overflow:hidden}
+.loader-bar-fill{
+    height:100%;background:linear-gradient(90deg,#FF4500,#FF8C00,#FF4500);
+    background-size:200% 100%;animation:shimmer 1.5s ease-in-out infinite;border-radius:2px;
+}
+@keyframes shimmer{0%{background-position:200% 0}100%{background-position:-200% 0}}
+
+/* ── Settings ── */
+.settings-group{border:1px solid #27272a;border-radius:10px;margin:12px 16px;overflow:hidden}
+.settings-group-title{
+    font-size:12px;font-weight:600;color:#71717a;text-transform:uppercase;letter-spacing:.8px;
+    padding:10px 16px;border-bottom:1px solid #27272a;background:rgba(24,24,27,.5);
+}
+.settings-group-body{padding:14px 16px;display:flex;flex-direction:column;gap:12px}
+.slider-row{display:flex;align-items:center;gap:10px;min-height:28px}
+.slider-row label{font-size:13px;font-weight:500;color:#a1a1aa;min-width:80px;flex-shrink:0}
+.slider-row input[type="range"]{
+    flex:1;-webkit-appearance:none;height:6px;background:#27272a;border-radius:3px;outline:none;
+}
+.slider-row input[type="range"]::-webkit-slider-thumb{
+    -webkit-appearance:none;width:16px;height:16px;background:linear-gradient(135deg,#FF4500,#CC3700);
+    border-radius:50%;cursor:pointer;box-shadow:0 2px 6px rgba(255,69,0,.4);transition:transform .15s;
+}
+.slider-row input[type="range"]::-webkit-slider-thumb:hover{transform:scale(1.2)}
+.slider-row .slider-val{
+    min-width:52px;text-align:right;font-family:'JetBrains Mono',monospace;font-size:12px;
+    padding:3px 8px;background:#09090b;border:1px solid #27272a;border-radius:6px;color:#a1a1aa;flex-shrink:0;
+}
+.checkbox-row{display:flex;align-items:center;gap:8px;font-size:13px;color:#a1a1aa}
+.checkbox-row input[type="checkbox"]{accent-color:#FF4500;width:16px;height:16px;cursor:pointer}
+
+/* ── Statusbar ── */
+.app-statusbar{
+    background:#18181b;border-top:1px solid #27272a;padding:6px 20px;
+    display:flex;gap:12px;height:34px;align-items:center;font-size:12px;
+}
+.sb-section{
+    padding:0 12px;flex:1;display:flex;align-items:center;font-family:'JetBrains Mono',monospace;
+    font-size:12px;color:#52525b;overflow:hidden;white-space:nowrap;
+}
+.sb-section.sb-fixed{
+    flex:0 0 auto;min-width:90px;text-align:center;justify-content:center;
+    padding:3px 12px;background:rgba(255,69,0,.08);border-radius:6px;color:#FF6B35;font-weight:500;
+}
+.exp-note{padding:10px 20px;font-size:12px;color:#52525b;border-top:1px solid #27272a;text-align:center}
+.exp-note a{color:#FF8C55;text-decoration:none}
+.exp-note a:hover{text-decoration:underline}
+
+::-webkit-scrollbar{width:8px;height:8px}
+::-webkit-scrollbar-track{background:#09090b}
+::-webkit-scrollbar-thumb{background:#27272a;border-radius:4px}
+::-webkit-scrollbar-thumb:hover{background:#3f3f46}
+
+@media(max-width:900px){
+    .app-main-row{flex-direction:column}
+    .app-main-right{width:100%}
+    .app-main-left{border-right:none;border-bottom:1px solid #27272a}
+    .output-compare-row{flex-direction:column}
 }
 """
 
-# ── UI ────────────────────────────────────────────────────────────────────────
-
-with gr.Blocks() as demo:
-
-    with gr.Column(elem_id="col-container"):
-
-        gr.Markdown(
-            "# ⚡ **Flux.2-4B-Encoder-Comparator**",
-            elem_id="main-title",
-        )
-        gr.Markdown(
-            "Compare **FLUX.2-klein-4B** side-by-side with two VAE decoders — generated **sequentially** from the **same seed**.\n\n"
-            "🟦 **Standard VAE** is generated **first**, then 🟩 **Small Decoder VAE** ([FLUX.2-small-decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)) · "
-            "[[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)]"
-        )
-
-        # ── Main two-column row ───────────────────────────────────────────────
-        with gr.Row(equal_height=True):
-
-            # ── Left: inputs ─────────────────────────────────────────────────
-            with gr.Column():
-                input_images = gr.Gallery(
-                    label="Input Image(s) for Editing (optional)",
-                    type="filepath",
-                    columns=2,
-                    rows=1,
-                    height=300,
-                    allow_preview=True,
-                )
-
-                prompt = gr.Text(
-                    label="Edit Prompt",
-                    show_label=True,
-                    placeholder="e.g., Transform the scene into a snowy winter day...",
-                )
-
-                run_button = gr.Button("⚡ Run Comparison", variant="primary")
-
-            # ── Right: outputs ────────────────────────────────────────────────
-            with gr.Column():
-                with gr.Row():
-
-                    # Standard VAE — Generated First
-                    with gr.Column():
-                        gr.HTML(
-                            '<span class="vae-badge" '
-                            'style="background:#FFE0CC;color:#CC3700;">'
-                            '🟦 Standard VAE</span>'
-                            '<p class="output-order-note">① Generated First</p>'
-                        )
-                        result_standard = gr.Image(
-                            label="Standard VAE",
-                            show_label=False,
-                            interactive=False,
-                            format="png",
-                            height=280,
-                        )
-
-                    # Small Decoder VAE — Generated Second
-                    with gr.Column():
-                        gr.HTML(
-                            '<span class="vae-badge" '
-                            'style="background:#FFF0E5;color:#E63E00;">'
-                            '🟩 Small Decoder VAE</span>'
-                            '<p class="output-order-note">② Generated Second</p>'
-                        )
-                        result_small = gr.Image(
-                            label="Small Decoder VAE",
-                            show_label=False,
-                            interactive=False,
-                            format="png",
-                            height=280,
-                        )
-
-                with gr.Accordion("Advanced Settings", open=False):
-                    seed_output = gr.Number(label="Seed Used", precision=0)
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=0,
-                    )
-                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-
-                    with gr.Row():
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-                        height_slider = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-
-                    with gr.Row():
-                        num_inference_steps = gr.Slider(
-                            label="Inference Steps",
-                            minimum=1,
-                            maximum=20,
-                            step=1,
-                            value=4,
-                        )
-                        guidance_scale = gr.Slider(
-                            label="Guidance Scale",
-                            minimum=0.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=1.0,
-                        )
-
-        # ── Examples ──────────────────────────────────────────────────────────
-        gr.Examples(
-            examples=[
-                [["examples/1.jpg"], "Change the weather to stormy."],
-                [["examples/2.jpg"], "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition."],
-                [["examples/3.jpg"], "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent."],
-                [["examples/4.jpg"], "Make the texture high-resolution."],
-            ],
-            inputs=[input_images, prompt],
-            outputs=[result_standard, result_small, seed_output],
-            fn=infer_example,
-            cache_examples=False,
-            label="Examples",
-        )
-
-        gr.Markdown(
-            "[*](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) "
-            "Experimental Space — FLUX.2 [klein] 4B VAE Decoder Comparison."
-        )
-
-    # ── Events ────────────────────────────────────────────────────────────────
-
-    input_images.upload(
-        fn=update_dimensions_from_image,
-        inputs=[input_images],
-        outputs=[width, height_slider],
-    )
+# ── SVGs ──────────────────────────────────────────────────────────────────────
+
+FLUX_LOGO_SVG = '''<svg viewBox="0 0 24 24" fill="white" xmlns="http://www.w3.org/2000/svg">
+<path d="M13 2L4.09 12.11a2 2 0 0 0-.09 2.16L8 21h8l4-6.73a2 2 0 0 0-.09-2.16L13 2z"/>
+</svg>'''
+
+UPLOAD_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4"/><polyline points="17 8 12 3 7 8"/><line x1="12" y1="3" x2="12" y2="15"/></svg>'
+REMOVE_SVG  = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="10"/><line x1="15" y1="9" x2="9" y2="15"/><line x1="9" y1="9" x2="15" y2="15"/></svg>'
+CLEAR_SVG   = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="3 6 5 6 21 6"/><path d="M19 6v14a2 2 0 01-2 2H7a2 2 0 01-2-2V6m3 0V4a2 2 0 012-2h4a2 2 0 012 2v2"/><line x1="10" y1="11" x2="10" y2="17"/><line x1="14" y1="11" x2="14" y2="17"/></svg>'
+DOWNLOAD_SVG = '<svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M12 16l-5-5h3V4h4v7h3l-5 5z"/><path d="M20 18H4v2h16v-2z"/></svg>'
+RUN_SVG = '<svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg" fill="#ffffff"><path d="M13 2L4.09 12.11a2 2 0 0 0-.09 2.16L8 21h8l4-6.73a2 2 0 0 0-.09-2.16L13 2z"/></svg>'
+
+# ── Gallery JS ────────────────────────────────────────────────────────────────
+
+gallery_js = r"""
+() => {
+function init() {
+    if (window.__fluxInitDone) return;
+
+    const galleryGrid  = document.getElementById('image-gallery-grid');
+    const dropZone     = document.getElementById('gallery-drop-zone');
+    const uploadPrompt = document.getElementById('upload-prompt');
+    const uploadClick  = document.getElementById('upload-click-area');
+    const fileInput    = document.getElementById('custom-file-input');
+    const btnUpload    = document.getElementById('tb-upload');
+    const btnRemove    = document.getElementById('tb-remove');
+    const btnClear     = document.getElementById('tb-clear');
+    const promptInput  = document.getElementById('custom-prompt-input');
+    const runBtnEl     = document.getElementById('custom-run-btn');
+    const imgCountTb   = document.getElementById('tb-image-count');
+    const imgCountSb   = document.getElementById('sb-image-count');
+
+    if (!galleryGrid || !fileInput || !dropZone) { setTimeout(init, 250); return; }
+    window.__fluxInitDone = true;
+
+    let images = [];
+    window.__uploadedImages = images;
+    let selectedIdx = -1;
+    let toastTimer  = null;
+
+    function showToast(message, type) {
+        let toast = document.getElementById('app-toast');
+        if (!toast) {
+            toast = document.createElement('div');
+            toast.id = 'app-toast';
+            toast.className = 'toast-notification';
+            toast.innerHTML = '<span class="toast-icon"></span><span class="toast-text"></span>';
+            document.body.appendChild(toast);
+        }
+        const icon = toast.querySelector('.toast-icon');
+        const text = toast.querySelector('.toast-text');
+        toast.className = 'toast-notification ' + (type || 'error');
+        icon.textContent = type === 'warning' ? '\u26A0' : type === 'info' ? '\u2139' : '\u2717';
+        text.textContent = message;
+        if (toastTimer) clearTimeout(toastTimer);
+        void toast.offsetWidth;
+        toast.classList.add('visible');
+        toastTimer = setTimeout(() => toast.classList.remove('visible'), 3500);
+    }
+    window.__showToast = showToast;
+
+    function flashPromptError() {
+        if (!promptInput) return;
+        promptInput.classList.add('error-flash');
+        promptInput.focus();
+        setTimeout(() => promptInput.classList.remove('error-flash'), 800);
+    }
+
+    function setGradioValue(containerId, value) {
+        const container = document.getElementById(containerId);
+        if (!container) return;
+        container.querySelectorAll('input, textarea').forEach(el => {
+            if (el.type === 'file' || el.type === 'range' || el.type === 'checkbox') return;
+            const proto = el.tagName === 'TEXTAREA' ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
+            const ns = Object.getOwnPropertyDescriptor(proto, 'value');
+            if (ns && ns.set) {
+                ns.set.call(el, value);
+                el.dispatchEvent(new Event('input',  {bubbles:true, composed:true}));
+                el.dispatchEvent(new Event('change', {bubbles:true, composed:true}));
+            }
+        });
+    }
+    window.__setGradioValue = setGradioValue;
+
+    function syncImagesToGradio() {
+        window.__uploadedImages = images;
+        const b64Array = images.map(img => img.b64);
+        setGradioValue('hidden-images-b64', JSON.stringify(b64Array));
+        updateCounts();
+    }
+
+    function syncPromptToGradio() {
+        if (promptInput) setGradioValue('prompt-gradio-input', promptInput.value);
+    }
+
+    function updateCounts() {
+        const n   = images.length;
+        const txt = n > 0 ? n + ' image' + (n > 1 ? 's' : '') : 'No images';
+        if (imgCountTb) imgCountTb.textContent = txt;
+        if (imgCountSb) imgCountSb.textContent = n > 0 ? txt + ' uploaded' : 'No images uploaded';
+    }
+
+    function addImage(b64, name) {
+        images.push({id: Date.now() + Math.random(), b64, name});
+        renderGallery();
+        syncImagesToGradio();
+    }
+    window.__addImage = addImage;
+
+    function removeImage(idx) {
+        images.splice(idx, 1);
+        if (selectedIdx === idx) selectedIdx = -1;
+        else if (selectedIdx > idx) selectedIdx--;
+        renderGallery();
+        syncImagesToGradio();
+    }
+
+    function clearAll() {
+        images = [];
+        window.__uploadedImages = images;
+        selectedIdx = -1;
+        renderGallery();
+        syncImagesToGradio();
+    }
+    window.__clearAll = clearAll;
+
+    function renderGallery() {
+        if (images.length === 0) {
+            galleryGrid.innerHTML = '';
+            galleryGrid.style.display = 'none';
+            if (uploadPrompt) uploadPrompt.style.display = '';
+            return;
+        }
+        if (uploadPrompt) uploadPrompt.style.display = 'none';
+        galleryGrid.style.display = 'grid';
+        let html = '';
+        images.forEach((img, i) => {
+            const sel = i === selectedIdx ? ' selected' : '';
+            html += '<div class="gallery-thumb' + sel + '" data-idx="' + i + '">'
+                  + '<img src="' + img.b64 + '" alt="">'
+                  + '<span class="thumb-badge">#' + (i+1) + '</span>'
+                  + '<button class="thumb-remove" data-remove="' + i + '">\u2715</button>'
+                  + '</div>';
+        });
+        html += '<div class="gallery-add-card" id="gallery-add-card"><span class="add-icon">+</span><span class="add-text">Add</span></div>';
+        galleryGrid.innerHTML = html;
+
+        galleryGrid.querySelectorAll('.gallery-thumb').forEach(thumb => {
+            thumb.addEventListener('click', (e) => {
+                if (e.target.closest('.thumb-remove')) return;
+                selectedIdx = selectedIdx === parseInt(thumb.dataset.idx) ? -1 : parseInt(thumb.dataset.idx);
+                renderGallery();
+            });
+        });
+        galleryGrid.querySelectorAll('.thumb-remove').forEach(btn => {
+            btn.addEventListener('click', (e) => { e.stopPropagation(); removeImage(parseInt(btn.dataset.remove)); });
+        });
+        const addCard = document.getElementById('gallery-add-card');
+        if (addCard) addCard.addEventListener('click', () => fileInput.click());
+    }
+
+    function processFiles(files) {
+        Array.from(files).forEach(file => {
+            if (!file.type.startsWith('image/')) return;
+            const reader = new FileReader();
+            reader.onload = (e) => addImage(e.target.result, file.name);
+            reader.readAsDataURL(file);
+        });
+    }
+
+    fileInput.addEventListener('change', (e) => { processFiles(e.target.files); e.target.value = ''; });
+    if (uploadClick) uploadClick.addEventListener('click', () => fileInput.click());
+    if (btnUpload)   btnUpload.addEventListener('click', () => fileInput.click());
+    if (btnRemove)   btnRemove.addEventListener('click', () => { if (selectedIdx >= 0) removeImage(selectedIdx); });
+    if (btnClear)    btnClear.addEventListener('click', clearAll);
+
+    dropZone.addEventListener('dragover',  (e) => { e.preventDefault(); dropZone.classList.add('drag-over'); });
+    dropZone.addEventListener('dragleave', (e) => { e.preventDefault(); dropZone.classList.remove('drag-over'); });
+    dropZone.addEventListener('drop',      (e) => { e.preventDefault(); dropZone.classList.remove('drag-over'); if (e.dataTransfer.files.length) processFiles(e.dataTransfer.files); });
+
+    if (promptInput) promptInput.addEventListener('input', syncPromptToGradio);
+    window.__setPrompt = function(text) { if (promptInput) { promptInput.value = text; syncPromptToGradio(); } };
+
+    // Example cards
+    document.querySelectorAll('.example-card[data-idx]').forEach(card => {
+        card.addEventListener('click', () => {
+            const idx = card.getAttribute('data-idx');
+            document.querySelectorAll('.example-card.loading').forEach(c => c.classList.remove('loading'));
+            card.classList.add('loading');
+            showToast('Loading example...', 'info');
+            setGradioValue('example-result-data', '');
+            setGradioValue('example-idx-input', idx);
+            setTimeout(() => {
+                const btn = document.getElementById('example-load-btn');
+                if (btn) { const b = btn.querySelector('button'); if (b) b.click(); else btn.click(); }
+            }, 150);
+            setTimeout(() => card.classList.remove('loading'), 12000);
+        });
+    });
+
+    // Sliders
+    function syncSlider(customId, gradioId) {
+        const slider = document.getElementById(customId);
+        const valSpan = document.getElementById(customId + '-val');
+        if (!slider) return;
+        slider.addEventListener('input', () => {
+            if (valSpan) valSpan.textContent = slider.value;
+            const container = document.getElementById(gradioId);
+            if (!container) return;
+            container.querySelectorAll('input[type="range"],input[type="number"]').forEach(el => {
+                const ns = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, 'value');
+                if (ns && ns.set) {
+                    ns.set.call(el, slider.value);
+                    el.dispatchEvent(new Event('input',  {bubbles:true, composed:true}));
+                    el.dispatchEvent(new Event('change', {bubbles:true, composed:true}));
+                }
+            });
+        });
+    }
+    syncSlider('custom-seed',     'gradio-seed');
+    syncSlider('custom-guidance', 'gradio-guidance');
+    syncSlider('custom-steps',    'gradio-steps');
+    syncSlider('custom-width',    'gradio-width');
+    syncSlider('custom-height',   'gradio-height');
+
+    const randCheck = document.getElementById('custom-randomize');
+    if (randCheck) {
+        randCheck.addEventListener('change', () => {
+            const container = document.getElementById('gradio-randomize');
+            if (!container) return;
+            const cb = container.querySelector('input[type="checkbox"]');
+            if (cb && cb.checked !== randCheck.checked) cb.click();
+        });
+    }
+
+    function showLoader(id) { const l = document.getElementById(id); if (l) l.classList.add('active'); }
+    function hideLoader(id) { const l = document.getElementById(id); if (l) l.classList.remove('active'); }
+    window.__showLoaders = () => { showLoader('loader-std'); showLoader('loader-small'); const sb = document.querySelector('.sb-fixed'); if (sb) sb.textContent = 'Processing...'; };
+    window.__hideLoaders = () => { hideLoader('loader-std'); hideLoader('loader-small'); const sb = document.querySelector('.sb-fixed'); if (sb) sb.textContent = 'Done'; };
+
+    function validateBeforeRun() {
+        const promptVal = promptInput ? promptInput.value.trim() : '';
+        if (!promptVal) { showToast('Please enter a prompt', 'warning'); flashPromptError(); return false; }
+        return true;
+    }
+
+    window.__clickGradioRunBtn = function() {
+        if (!validateBeforeRun()) return;
+        syncPromptToGradio(); syncImagesToGradio();
+        window.__showLoaders();
+        setTimeout(() => {
+            const gradioBtn = document.getElementById('gradio-run-btn');
+            if (!gradioBtn) return;
+            const btn = gradioBtn.querySelector('button');
+            if (btn) btn.click(); else gradioBtn.click();
+        }, 200);
+    };
+
+    if (runBtnEl) runBtnEl.addEventListener('click', () => window.__clickGradioRunBtn());
+
+    renderGallery();
+    updateCounts();
+}
+init();
+}
+"""
+
+wire_outputs_js = r"""
+() => {
+function watchOutputs() {
+    const resultContainer = document.getElementById('gradio-result-json');
+    if (!resultContainer) { setTimeout(watchOutputs, 500); return; }
+
+    let lastVal = '';
+
+    function checkResult() {
+        const el = resultContainer.querySelector('textarea') || resultContainer.querySelector('input');
+        if (!el || !el.value || el.value === lastVal || el.value.length < 20) return;
+        lastVal = el.value;
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
+        try {
+            const data = JSON.parse(el.value);
+            if (data.status !== 'ok') return;
+
+            // Standard VAE output
+            const stdWrap = document.getElementById('output-img-std');
+            const stdPh   = document.getElementById('output-ph-std');
+            const dlStd   = document.getElementById('dl-btn-std');
+            if (stdWrap && data.standard_b64) {
+                if (stdPh) stdPh.style.display = 'none';
+                let img = stdWrap.querySelector('img.out-img');
+                if (!img) { img = document.createElement('img'); img.className = 'out-img'; stdWrap.appendChild(img); }
+                img.src = data.standard_b64;
+                if (dlStd) { dlStd.classList.add('visible'); dlStd.onclick = () => { const a = document.createElement('a'); a.href = data.standard_b64; a.download = 'flux_standard.png'; a.click(); }; }
+            }
+
+            // Small Decoder output
+            const smWrap = document.getElementById('output-img-small');
+            const smPh   = document.getElementById('output-ph-small');
+            const dlSm   = document.getElementById('dl-btn-small');
+            if (smWrap && data.small_b64) {
+                if (smPh) smPh.style.display = 'none';
+                let img = smWrap.querySelector('img.out-img');
+                if (!img) { img = document.createElement('img'); img.className = 'out-img'; smWrap.appendChild(img); }
+                img.src = data.small_b64;
+                if (dlSm) { dlSm.classList.add('visible'); dlSm.onclick = () => { const a = document.createElement('a'); a.href = data.small_b64; a.download = 'flux_small_decoder.png'; a.click(); }; }
+            }
+
+            // Timing chips
+            const chipStd = document.getElementById('timing-chip-std');
+            const chipSm  = document.getElementById('timing-chip-small');
+            if (chipStd && data.time_std)   { chipStd.textContent  = '⏱ ' + data.time_std;   chipStd.classList.add('done'); }
+            if (chipSm  && data.time_small) { chipSm.textContent   = '⏱ ' + data.time_small; chipSm.classList.add('done'); }
+
+            // Seed
+            const seedSb = document.getElementById('sb-seed');
+            if (seedSb && data.seed !== undefined) seedSb.textContent = 'seed: ' + data.seed;
+
+            if (window.__hideLoaders) window.__hideLoaders();
+
+        } catch(e) { console.error('Output parse error:', e); }
+    }
+
+    const obs = new MutationObserver(checkResult);
+    obs.observe(resultContainer, {childList:true, subtree:true, characterData:true, attributes:true});
+    setInterval(checkResult, 600);
+}
+watchOutputs();
+
+function watchExampleResults() {
+    const container = document.getElementById('example-result-data');
+    if (!container) { setTimeout(watchExampleResults, 500); return; }
+    let lastProcessed = '';
+    function checkResult() {
+        const el = container.querySelector('textarea') || container.querySelector('input');
+        if (!el || !el.value || el.value === lastProcessed || el.value.length < 20) return;
+        try {
+            const data = JSON.parse(el.value);
+            if (data.status === 'ok' && data.images && data.images.length > 0) {
+                lastProcessed = el.value;
+                if (window.__clearAll) window.__clearAll();
+                if (window.__setPrompt && data.prompt) window.__setPrompt(data.prompt);
+                data.images.forEach((b64, i) => {
+                    if (b64 && window.__addImage) {
+                        const name = (data.names && data.names[i]) ? data.names[i] : ('example_' + (i+1) + '.jpg');
+                        window.__addImage(b64, name);
+                    }
+                });
+                document.querySelectorAll('.example-card.loading').forEach(c => c.classList.remove('loading'));
+                if (window.__showToast) window.__showToast('Example loaded — ' + data.images.length + ' image(s)', 'info');
+            }
+        } catch(e) {}
+    }
+    const obs = new MutationObserver(checkResult);
+    obs.observe(container, {childList:true, subtree:true, characterData:true, attributes:true});
+    setInterval(checkResult, 500);
+}
+watchExampleResults();
+}
+"""
+
+# ── Gradio App ────────────────────────────────────────────────────────────────
+
+with gr.Blocks(css=css) as demo:
+
+    # Hidden Gradio state components
+    hidden_images_b64  = gr.Textbox(value="[]",  elem_id="hidden-images-b64",   elem_classes="hidden-input", container=False)
+    prompt             = gr.Textbox(value="",    elem_id="prompt-gradio-input",  elem_classes="hidden-input", container=False)
+    seed               = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0,
+                                   elem_id="gradio-seed",       elem_classes="hidden-input", container=False)
+    randomize_seed     = gr.Checkbox(value=True, elem_id="gradio-randomize",     elem_classes="hidden-input", container=False)
+    guidance_scale     = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=1.0,
+                                   elem_id="gradio-guidance",   elem_classes="hidden-input", container=False)
+    steps              = gr.Slider(minimum=1, maximum=20, step=1, value=4,
+                                   elem_id="gradio-steps",      elem_classes="hidden-input", container=False)
+    width_gr           = gr.Slider(minimum=256, maximum=1024, step=8, value=1024,
+                                   elem_id="gradio-width",      elem_classes="hidden-input", container=False)
+    height_gr          = gr.Slider(minimum=256, maximum=1024, step=8, value=1024,
+                                   elem_id="gradio-height",     elem_classes="hidden-input", container=False)
+
+    result_std_gr      = gr.Image(elem_id="gradio-result-std",   elem_classes="hidden-input", container=False, format="png")
+    result_small_gr    = gr.Image(elem_id="gradio-result-small", elem_classes="hidden-input", container=False, format="png")
+    seed_out_gr        = gr.Number(elem_id="gradio-seed-out",    elem_classes="hidden-input", container=False)
+    result_json_gr     = gr.Textbox(value="",  elem_id="gradio-result-json",    elem_classes="hidden-input", container=False)
+
+    example_idx        = gr.Textbox(value="", elem_id="example-idx-input",  elem_classes="hidden-input", container=False)
+    example_result     = gr.Textbox(value="", elem_id="example-result-data", elem_classes="hidden-input", container=False)
+    example_load_btn   = gr.Button("Load Example", elem_id="example-load-btn")
+
+    gr.HTML(f"""
+    <div class="app-shell">
+
+      <!-- Header -->
+      <div class="app-header">
+        <div class="app-header-left">
+          <div class="app-logo">{FLUX_LOGO_SVG}</div>
+          <span class="app-title">FLUX.2-klein-4B VAE Comparator</span>
+          <span class="app-badge">4B Distilled</span>
+          <span class="app-badge fast">4-Step Fast</span>
+        </div>
+      </div>
+
+      <!-- Toolbar -->
+      <div class="app-toolbar">
+        <button id="tb-upload" class="modern-tb-btn" title="Upload images">
+          {UPLOAD_SVG}<span class="tb-label">Upload</span>
+        </button>
+        <button id="tb-remove" class="modern-tb-btn" title="Remove selected">
+          {REMOVE_SVG}<span class="tb-label">Remove</span>
+        </button>
+        <button id="tb-clear" class="modern-tb-btn" title="Clear all">
+          {CLEAR_SVG}<span class="tb-label">Clear All</span>
+        </button>
+        <div class="tb-sep"></div>
+        <span id="tb-image-count" class="tb-info">No images</span>
+        <div class="tb-sep"></div>
+        <span class="tb-info" style="color:#71717a;">
+          Compare 🟦 Standard VAE vs 🟩 Small Decoder VAE — same seed, sequential generation
+        </span>
+      </div>
+
+      <!-- Main -->
+      <div class="app-main-row">
+
+        <!-- Left: Upload + Examples -->
+        <div class="app-main-left">
+          <div id="gallery-drop-zone">
+            <div id="upload-prompt" class="upload-prompt-modern">
+              <div id="upload-click-area" class="upload-click-area">
+                <svg viewBox="0 0 80 80" fill="none" xmlns="http://www.w3.org/2000/svg">
+                  <rect x="8" y="14" width="64" height="52" rx="6" fill="none" stroke="#FF4500" stroke-width="2" stroke-dasharray="4 3"/>
+                  <polygon points="12,62 30,40 42,50 54,34 68,62" fill="rgba(255,69,0,0.15)" stroke="#FF4500" stroke-width="1.5"/>
+                  <circle cx="28" cy="30" r="6" fill="rgba(255,69,0,0.2)" stroke="#FF4500" stroke-width="1.5"/>
+                </svg>
+                <span class="upload-main-text">Click or drag reference image(s) here</span>
+                <span class="upload-sub-text">Optional — leave empty for text-to-image generation</span>
+              </div>
+            </div>
+            <input id="custom-file-input" type="file" accept="image/*" multiple style="display:none;" />
+            <div id="image-gallery-grid" class="image-gallery-grid" style="display:none;"></div>
+          </div>
+
+          <div class="hint-bar">
+            <b>Optional:</b> Upload reference image(s) for image-guided generation &nbsp;&middot;&nbsp;
+            Leave empty for <b>text-to-image</b> mode &nbsp;&middot;&nbsp;
+            <kbd>Remove</kbd> deletes selected &nbsp;&middot;&nbsp; <kbd>Clear All</kbd> removes everything
+          </div>
+
+          <div class="suggestions-section">
+            <div class="suggestions-title">Quick Prompts</div>
+            <div class="suggestions-wrap">
+              <button class="suggestion-chip" onclick="window.__setPrompt('Change the weather to stormy.')">Stormy Weather</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Transform the scene into a snowy winter day.')">Snowy Winter</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Relight with soft golden sunset lighting.')">Golden Sunset</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Make the texture high-resolution and ultra detailed.')">HD Texture</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Convert to cinematic film look with grain and vignette.')">Cinematic Film</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Transform into anime illustration style.')">Anime Style</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Apply oil painting effect with visible brush strokes.')">Oil Painting</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Add dramatic neon glow with cyberpunk aesthetic.')">Neon Cyberpunk</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('A futuristic city at night with flying cars and neon lights.')">Futuristic City</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('A serene mountain landscape at dawn with mist.')">Mountain Dawn</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Portrait of a warrior in ancient armor, epic lighting.')">Epic Portrait</button>
+              <button class="suggestion-chip" onclick="window.__setPrompt('Abstract geometric art with vibrant colors and patterns.')">Abstract Art</button>
+            </div>
+          </div>
+
+          <div class="examples-section">
+            <div class="examples-title">Quick Examples</div>
+            <div class="examples-scroll">
+              {EXAMPLE_CARDS_HTML}
+            </div>
+          </div>
+        </div>
+
+        <!-- Right: Prompt + Outputs + Settings -->
+        <div class="app-main-right">
+
+          <div class="panel-card">
+            <div class="panel-card-title">Generation Prompt</div>
+            <div class="panel-card-body">
+              <label class="modern-label" for="custom-prompt-input">Prompt</label>
+              <textarea id="custom-prompt-input" class="modern-textarea" rows="3"
+                placeholder="e.g., Transform the scene into a snowy winter day..."></textarea>
+            </div>
+          </div>
+
+          <div style="padding:12px 20px;">
+            <button id="custom-run-btn" class="btn-run">
+              {RUN_SVG}
+              <span>⚡ Run Comparison</span>
+            </button>
+          </div>
+
+          <!-- Output Compare -->
+          <div class="output-section">
+            <div class="output-section-title">Output Comparison</div>
+            <div class="output-compare-row">
+
+              <!-- Standard VAE -->
+              <div class="output-col">
+                <div class="output-col-label std">
+                  <span>🟦 Standard VAE</span>
+                  <span id="dl-btn-std" class="out-download-btn" title="Download">{DOWNLOAD_SVG}</span>
+                </div>
+                <div class="output-img-wrap" id="output-img-std">
+                  <div class="modern-loader" id="loader-std">
+                    <div class="loader-spinner"></div>
+                    <div class="loader-text">Generating...</div>
+                    <div class="loader-bar-track"><div class="loader-bar-fill"></div></div>
+                  </div>
+                  <div class="output-placeholder" id="output-ph-std">Result will appear here</div>
+                </div>
+                <span class="timing-chip" id="timing-chip-std">⏱ Waiting...</span>
+              </div>
+
+              <!-- Small Decoder VAE -->
+              <div class="output-col">
+                <div class="output-col-label small">
+                  <span>🟩 Small Decoder</span>
+                  <span id="dl-btn-small" class="out-download-btn" title="Download">{DOWNLOAD_SVG}</span>
+                </div>
+                <div class="output-img-wrap" id="output-img-small">
+                  <div class="modern-loader" id="loader-small">
+                    <div class="loader-spinner"></div>
+                    <div class="loader-text">Generating...</div>
+                    <div class="loader-bar-track"><div class="loader-bar-fill"></div></div>
+                  </div>
+                  <div class="output-placeholder" id="output-ph-small">Result will appear here</div>
+                </div>
+                <span class="timing-chip" id="timing-chip-small">⏱ Waiting...</span>
+              </div>
+
+            </div>
+          </div>
+
+          <!-- Advanced Settings -->
+          <div class="settings-group">
+            <div class="settings-group-title">Advanced Settings</div>
+            <div class="settings-group-body">
+              <div class="slider-row">
+                <label>Seed</label>
+                <input type="range" id="custom-seed" min="0" max="2147483647" step="1" value="0">
+                <span class="slider-val" id="custom-seed-val">0</span>
+              </div>
+              <div class="checkbox-row">
+                <input type="checkbox" id="custom-randomize" checked>
+                <label for="custom-randomize">Randomize seed</label>
+              </div>
+              <div class="slider-row">
+                <label>Guidance</label>
+                <input type="range" id="custom-guidance" min="0" max="10" step="0.1" value="1.0">
+                <span class="slider-val" id="custom-guidance-val">1.0</span>
+              </div>
+              <div class="slider-row">
+                <label>Steps</label>
+                <input type="range" id="custom-steps" min="1" max="20" step="1" value="4">
+                <span class="slider-val" id="custom-steps-val">4</span>
+              </div>
+              <div class="slider-row">
+                <label>Width</label>
+                <input type="range" id="custom-width" min="256" max="1024" step="8" value="1024">
+                <span class="slider-val" id="custom-width-val">1024</span>
+              </div>
+              <div class="slider-row">
+                <label>Height</label>
+                <input type="range" id="custom-height" min="256" max="1024" step="8" value="1024">
+                <span class="slider-val" id="custom-height-val">1024</span>
+              </div>
+            </div>
+          </div>
+
+        </div>
+      </div>
+
+      <div class="exp-note">
+        Experimental Space — <a href="https://huggingface.co/black-forest-labs/FLUX.2-klein-4B" target="_blank">FLUX.2-klein-4B</a>
+        &nbsp;·&nbsp; <a href="https://huggingface.co/black-forest-labs/FLUX.2-small-decoder" target="_blank">FLUX.2-small-decoder</a>
+        &nbsp;·&nbsp; Comparing Standard VAE vs Small Decoder VAE from the same seed
+      </div>
+
+      <div class="app-statusbar">
+        <div class="sb-section" id="sb-image-count">No images uploaded</div>
+        <div class="sb-section" id="sb-seed" style="color:#52525b;font-family:'JetBrains Mono',monospace;font-size:12px;">seed: —</div>
+        <div class="sb-section sb-fixed">Ready</div>
+      </div>
+
+    </div>
+    """)
+
+    run_btn = gr.Button("Run", elem_id="gradio-run-btn")
+
+    demo.load(fn=None, js=gallery_js)
+    demo.load(fn=None, js=wire_outputs_js)
+
+    run_btn.click(
         fn=infer,
-        inputs=[
-            prompt,
-            input_images,
-            seed,
-            randomize_seed,
-            width,
-            height_slider,
-            num_inference_steps,
-            guidance_scale,
-        ],
-        outputs=[result_standard, result_small, seed_output],
+        inputs=[hidden_images_b64, prompt, seed, randomize_seed,
+                width_gr, height_gr, steps, guidance_scale],
+        outputs=[result_std_gr, result_small_gr, seed_out_gr, result_json_gr],
+        js=r"""(imgs, p, s, rs, w, h, st, gs) => {
+            const images   = window.__uploadedImages || [];
+            const b64Array = images.map(img => img.b64);
+            const imgsJson = JSON.stringify(b64Array);
+            const promptEl = document.getElementById('custom-prompt-input');
+            const promptVal = promptEl ? promptEl.value : p;
+            const seedEl   = document.getElementById('custom-seed');
+            const guidEl   = document.getElementById('custom-guidance');
+            const stepsEl  = document.getElementById('custom-steps');
+            const wEl      = document.getElementById('custom-width');
+            const hEl      = document.getElementById('custom-height');
+            const randEl   = document.getElementById('custom-randomize');
+            return [
+                imgsJson, promptVal,
+                seedEl  ? parseFloat(seedEl.value)  : s,
+                randEl  ? randEl.checked             : rs,
+                wEl     ? parseFloat(wEl.value)      : w,
+                hEl     ? parseFloat(hEl.value)      : h,
+                stepsEl ? parseFloat(stepsEl.value)  : st,
+                guidEl  ? parseFloat(guidEl.value)   : gs,
+            ];
+        }""",
+    )
+
+    example_load_btn.click(
+        fn=load_example_data,
+        inputs=[example_idx],
+        outputs=[example_result],
+        queue=False,
     )
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(
-        css=css, theme=orange_red_theme,
+        css=css,
         ssr_mode=False,
         show_error=True,
+        allowed_paths=["examples"],
     )