Update app.py

app.py

@@ -10,33 +10,33 @@ import concurrent.futures
 from pathlib import Path
 from typing import List, Optional
 
-import numpy as np
 import spaces
+import numpy as np
 import torch
 from PIL import Image
-from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 
 from gradio import Server
 from fastapi import Request, UploadFile, File, Form
 from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
+from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# -------------------------------------------------------------------------
-# App Setup & Directories
-# -------------------------------------------------------------------------
 app = Server()
 
 BASE_DIR = Path(__file__).resolve().parent
+STATIC_DIR = BASE_DIR / "static"
 OUTPUT_DIR = BASE_DIR / "outputs"
 EXAMPLES_DIR = BASE_DIR / "examples"
 
+STATIC_DIR.mkdir(exist_ok=True)
 OUTPUT_DIR.mkdir(exist_ok=True)
-EXAMPLES_DIR.mkdir(exist_ok=True)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-dtype  = torch.bfloat16
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16
 
 if torch.cuda.is_available():
     print("current device:", torch.cuda.current_device())
@@ -45,99 +45,96 @@ if torch.cuda.is_available():
 else:
     DEVICE_LABEL = str(DEVICE).lower()
 
-# -------------------------------------------------------------------------
-# Model Loading
-# -------------------------------------------------------------------------
 print("Loading 4B Distilled model (Standard VAE)...")
 pipe_standard = Flux2KleinPipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
     torch_dtype=dtype,
-).to(DEVICE)
-# pipe_standard.enable_model_cpu_offload() # Uncomment if memory is tight
+    token=HF_TOKEN,
+)
+pipe_standard.enable_model_cpu_offload()
 
 print("Loading Small Decoder VAE...")
 vae_small = AutoencoderKLFlux2.from_pretrained(
     "black-forest-labs/FLUX.2-small-decoder",
     torch_dtype=dtype,
-).to(DEVICE)
+    token=HF_TOKEN,
+)
 
 print("Loading 4B Distilled model (Small Decoder VAE)...")
 pipe_small_decoder = Flux2KleinPipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
     vae=vae_small,
     torch_dtype=dtype,
-).to(DEVICE)
-# pipe_small_decoder.enable_model_cpu_offload() # Uncomment if memory is tight
+    token=HF_TOKEN,
+)
+pipe_small_decoder.enable_model_cpu_offload()
 
 pipe_lock_standard = threading.Lock()
-pipe_lock_small    = threading.Lock()
+pipe_lock_small = threading.Lock()
 
 
-# -------------------------------------------------------------------------
-# Image Processing Utilities
-# -------------------------------------------------------------------------
-def save_image(img: Image.Image, prefix: str = "output") -> str:
-    filename = f"{prefix}_{uuid.uuid4().hex}.png"
-    path = OUTPUT_DIR / filename
-    img.save(path, format="PNG")
-    return filename
-
 def calc_dimensions(pil_img: Image.Image):
     iw, ih = pil_img.size
     aspect = iw / ih
 
-    if aspect >= 1:          # landscape / square
-        new_width  = 1024
+    if aspect >= 1:
+        new_width = 1024
         new_height = int(round(1024 / aspect))
-    else:                    # portrait
+    else:
         new_height = 1024
-        new_width  = int(round(1024 * aspect))
+        new_width = int(round(1024 * aspect))
 
-    # snap to 8-pixel grid with round(), clamp to [256, 1024]
-    new_width  = max(256, min(1024, round(new_width  / 8) * 8))
+    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_and_resize_images(input_paths: List[str], width: int, height: int):
-    if not input_paths:
+
+def parse_and_resize_images(input_images, width: int, height: int):
+    if not input_images:
         return None
 
     raw_list = []
-    for src in input_paths:
+    for item in input_images:
         try:
-            if os.path.exists(src):
-                raw_list.append(Image.open(src).convert("RGB"))
+            if isinstance(item, str):
+                raw_list.append(Image.open(item).convert("RGB"))
+            elif isinstance(item, Image.Image):
+                raw_list.append(item.convert("RGB"))
+            else:
+                raw_list.append(Image.open(item).convert("RGB"))
         except Exception as e:
             print(f"Skipping invalid image: {e}")
 
     if not raw_list:
         return None
 
-    resized = [
-        img.resize((width, height), Image.LANCZOS)
-        for img in raw_list
-    ]
-    return resized
+    return [img.resize((width, height), Image.LANCZOS) for img in raw_list]
+
 
 def run_pipeline(pipe, lock, kwargs, seed):
     with lock:
-        gen    = torch.Generator(device=DEVICE).manual_seed(seed)
+        gen = torch.Generator(device="cpu").manual_seed(seed)
         result = pipe(**kwargs, generator=gen).images[0]
     return result
 
-# -------------------------------------------------------------------------
-# Inference Function
-# -------------------------------------------------------------------------
+
+def save_image(img: Image.Image, prefix: str = "output") -> str:
+    filename = f"{prefix}_{uuid.uuid4().hex}.png"
+    path = OUTPUT_DIR / filename
+    img.save(path, format="PNG")
+    return filename
+
+
 @spaces.GPU(duration=120)
 def infer(
-    prompt: str,
-    input_paths: List[str],
-    seed: int,
-    randomize_seed: bool,
-    width: int,
-    height: int,
-    num_inference_steps: int,
-    guidance_scale: float,
+    prompt,
+    input_images=None,
+    seed=42,
+    randomize_seed=False,
+    width=1024,
+    height=1024,
+    num_inference_steps=4,
+    guidance_scale=1.0,
 ):
     gc.collect()
     if torch.cuda.is_available():
@@ -150,13 +147,20 @@ def infer(
         seed = random.randint(0, MAX_SEED)
 
     image_list = None
-    if input_paths:
-        first_pil = Image.open(input_paths[0]).convert("RGB")
-        width, height = calc_dimensions(first_pil)
-        image_list = parse_and_resize_images(input_paths, width, height)
+    if input_images and len(input_images) > 0:
+        first_img = input_images[0]
+        if isinstance(first_img, str):
+            first_pil = Image.open(first_img).convert("RGB")
+        elif isinstance(first_img, Image.Image):
+            first_pil = first_img.convert("RGB")
+        else:
+            first_pil = Image.open(first_img).convert("RGB")
 
-    width  = max(256, min(MAX_IMAGE_SIZE, round(int(width)  / 8) * 8))
-    height = max(256, min(MAX_IMAGE_SIZE, round(int(height) / 8) * 8))
+        width, height = calc_dimensions(first_pil)
+        image_list = parse_and_resize_images(input_images, width, height)
+    else:
+        width = max(256, min(MAX_IMAGE_SIZE, round(int(width) / 8) * 8))
+        height = max(256, min(MAX_IMAGE_SIZE, round(int(height) / 8) * 8))
 
     shared_kwargs = dict(
         prompt=prompt,
@@ -168,21 +172,13 @@ def infer(
     if image_list is not None:
         shared_kwargs["image"] = image_list
 
-    # Run both models concurrently
     with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
-        future_std   = executor.submit(
-            run_pipeline, pipe_standard,     pipe_lock_standard, shared_kwargs, seed
-        )
-        future_small = executor.submit(
-            run_pipeline, pipe_small_decoder, pipe_lock_small,    shared_kwargs, seed
-        )
-        concurrent.futures.wait(
-            [future_std, future_small],
-            return_when=concurrent.futures.ALL_COMPLETED,
-        )
+        future_std = executor.submit(run_pipeline, pipe_standard, pipe_lock_standard, shared_kwargs, seed)
+        future_small = executor.submit(run_pipeline, pipe_small_decoder, pipe_lock_small, shared_kwargs, seed)
+        concurrent.futures.wait([future_std, future_small], return_when=concurrent.futures.ALL_COMPLETED)
 
     out_standard = future_std.result()
-    out_small    = future_small.result()
+    out_small = future_small.result()
 
     gc.collect()
     if torch.cuda.is_available():
@@ -191,9 +187,36 @@ def infer(
     return out_standard, out_small, seed
 
 
-# -------------------------------------------------------------------------
-# API Endpoints
-# -------------------------------------------------------------------------
+def get_example_items():
+    example_prompts = {
+        "1.jpg": "Change the weather to stormy.",
+        "2.jpg": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.",
+        "3.jpg": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.",
+        "4.jpg": "Make the texture high-resolution.",
+    }
+
+    items = []
+    if EXAMPLES_DIR.exists():
+        for name in sorted(os.listdir(EXAMPLES_DIR)):
+            if name.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
+                items.append(
+                    {
+                        "file": name,
+                        "url": f"/example-file/{name}",
+                        "prompt": example_prompts.get(name, "Edit this image while preserving composition."),
+                    }
+                )
+    return items
+
+
+@app.get("/example-file/{filename}")
+async def example_file(filename: str):
+    path = EXAMPLES_DIR / filename
+    if not path.exists():
+        return JSONResponse({"error": "Example not found"}, status_code=404)
+    return FileResponse(path)
+
+
 @app.get("/download/{filename}")
 async def download_file(filename: str):
     path = OUTPUT_DIR / filename
@@ -201,8 +224,9 @@ async def download_file(filename: str):
         return JSONResponse({"error": "File not found"}, status_code=404)
     return FileResponse(path, filename=filename, media_type="image/png")
 
+
 @app.post("/api/compare")
-async def compare_images(
+async def compare_api(
     prompt: str = Form(...),
     seed: str = Form("0"),
     randomize_seed: str = Form("true"),
@@ -217,8 +241,6 @@ async def compare_images(
         image_paths = []
         if images:
             for upload in images:
-                if not upload.filename:
-                    continue
                 suffix = Path(upload.filename).suffix or ".png"
                 temp_name = f"upload_{uuid.uuid4().hex}{suffix}"
                 temp_path = OUTPUT_DIR / temp_name
@@ -229,26 +251,28 @@ async def compare_images(
                 image_paths.append(str(temp_path))
 
         out_std, out_small, used_seed = infer(
+            images=image_paths if image_paths else None,
             prompt=prompt,
-            input_paths=image_paths,
             seed=int(seed),
             randomize_seed=randomize_seed.lower() == "true",
             width=int(width),
             height=int(height),
             num_inference_steps=int(steps),
-            guidance_scale=float(guidance)
+            guidance_scale=float(guidance),
         )
 
         std_filename = save_image(out_std, prefix="std")
-        sml_filename = save_image(out_small, prefix="sml")
-
-        return JSONResponse({
-            "success": True,
-            "seed": used_seed,
-            "std_url": f"/download/{std_filename}",
-            "sml_url": f"/download/{sml_filename}",
-            "device": DEVICE_LABEL,
-        })
+        small_filename = save_image(out_small, prefix="small")
+
+        return JSONResponse(
+            {
+                "success": True,
+                "seed": used_seed,
+                "std_url": f"/download/{std_filename}",
+                "small_url": f"/download/{small_filename}",
+                "device": DEVICE_LABEL,
+            }
+        )
 
     except Exception as e:
         return JSONResponse({"success": False, "error": str(e)}, status_code=500)
@@ -261,448 +285,499 @@ async def compare_images(
                 pass
 
 
-# -------------------------------------------------------------------------
-# Frontend Template
-# -------------------------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def homepage(request: Request):
+    examples = get_example_items()
+    examples_json = json.dumps(examples)
+
     return f"""
 <!DOCTYPE html>
 <html lang="en">
 <head>
   <meta charset="UTF-8" />
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Flux.2-4B-Decoder-Comparator</title>
+  <title>Flux.2-4B VAE Decoder Comparison</title>
   <style>
     @import url('https://fonts.googleapis.com/css2?family=Ubuntu:wght@300;400;500;700&display=swap');
 
     :root {{
-      /* Ubuntu Brand Colors */
-      --ubuntu-orange: #E95420;
-      --ubuntu-orange-hover: #C7481B;
-      --ubuntu-aubergine: #77216F;
-      --ubuntu-aubergine-dark: #2C001E;
-      --ubuntu-warm-grey: #AEA79F;
-      
       --bg: #111111;
-      --panel: #1e1e1e;
+      --panel: #1E1E1E;
       --panel-2: #252525;
+      --panel-3: #2d2d2d;
       --border: #333333;
-      --muted: #9CA3AF;
-      --text: #F5F7FB;
-      --text-dim: #D1D5DB;
-      --input-bg: #181818;
-      --same-height: 760px;
+      --text: #ffffff;
+      --muted: #aaaaaa;
+      --ubuntu-orange: #E95420;
+      --ubuntu-orange-hover: #c74316;
+      --ubuntu-aubergine: #2C001E;
+      --ubuntu-aubergine-light: #44002e;
+      --input-bg: #151515;
+      --same-height: 800px;
     }}
 
-    * {{ box-sizing: border-box; border-radius: 4px; }}
-
+    * {{ box-sizing: border-box; font-family: 'Ubuntu', sans-serif; }}
+    
     html, body {{
       margin: 0; padding: 0;
-      background: var(--bg);
-      color: var(--text);
-      font-family: 'Ubuntu', sans-serif;
-      min-height: 100%;
-      overflow-x: hidden;
+      background: var(--bg); color: var(--text);
+      min-height: 100%; overflow-x: hidden;
     }}
 
     .topbar {{
       height: 60px;
+      background: var(--ubuntu-aubergine);
       border-bottom: 2px solid var(--ubuntu-orange);
-      background: var(--ubuntu-aubergine-dark);
-      display: flex;
-      align-items: center;
-      justify-content: center;
-      padding: 0 24px;
-      color: white;
-      font-size: 16px;
-      font-weight: 500;
-      letter-spacing: 0.5px;
+      display: flex; align-items: center; justify-content: center;
+      color: #fff; font-size: 16px; font-weight: 500; letter-spacing: 0.02em;
     }}
 
-    .container {{
-      max-width: 1440px;
-      margin: 0 auto;
-      padding: 28px;
-    }}
+    .container {{ max-width: 1440px; margin: 0 auto; padding: 28px; }}
 
     .hero {{ margin-bottom: 24px; padding-bottom: 20px; border-bottom: 1px solid var(--border); }}
-    .title {{ font-size: 38px; line-height: 1.2; font-weight: 700; margin: 0 0 8px 0; color: white; }}
-    .subtitle {{ color: var(--ubuntu-warm-grey); font-size: 16px; margin:0; }}
-    .subtitle a {{ color: var(--ubuntu-orange); text-decoration: none; }}
+    .title {{ font-size: 38px; line-height: 1.2; font-weight: 700; margin: 0 0 10px 0; }}
+    .subtitle {{ color: var(--muted); font-size: 16px; font-weight: 300; margin: 0; }}
 
     .layout {{
-      display: grid;
-      grid-template-columns: 450px 1fr;
-      gap: 24px;
-      align-items: stretch;
+      display: grid; grid-template-columns: 1fr 1.2fr; gap: 24px; align-items: stretch;
     }}
 
     .panel {{
-      background: var(--panel);
-      border: 1px solid var(--border);
-      min-height: var(--same-height);
-      display: flex;
-      flex-direction: column;
+      background: var(--panel); border: 1px solid var(--border);
+      min-height: var(--same-height); height: var(--same-height);
+      display: flex; flex-direction: column; overflow: hidden;
+      border-radius: 6px;
     }}
 
     .panel-header {{
-      height: 56px;
-      border-bottom: 1px solid var(--border);
+      height: 62px; min-height: 62px; border-bottom: 1px solid var(--border);
       display: flex; align-items: center; justify-content: space-between;
-      padding: 0 18px; background: var(--panel-2);
+      padding: 0 20px; background: #1a1a1a;
     }}
-    .panel-title {{ font-size: 18px; font-weight: 700; margin: 0; color: white; }}
-    
-    .panel-body {{ padding: 18px; flex: 1; overflow: auto; display: flex; flex-direction: column; gap: 16px; }}
 
-    .label {{ font-size: 14px; font-weight: 500; color: var(--text-dim); margin-bottom: 6px; display: block; }}
+    .panel-title {{ font-size: 20px; font-weight: 500; margin: 0; color: var(--text); }}
     
+    .status-pill {{
+      padding: 6px 12px; background: var(--panel-3); border: 1px solid var(--border);
+      color: var(--muted); font-size: 12px; font-weight: 500; border-radius: 20px;
+      transition: all 0.2s ease;
+    }}
+    .status-pill.active {{
+      background: rgba(233, 84, 32, 0.15); border-color: var(--ubuntu-orange); color: var(--ubuntu-orange);
+    }}
+
+    .panel-body {{ flex: 1; padding: 20px; overflow: auto; display: flex; flex-direction: column; gap: 20px; }}
+
+    .label {{ font-size: 14px; font-weight: 500; color: #ddd; margin-bottom: 8px; display: block; }}
+    .hint {{ color: var(--muted); font-size: 12px; margin-top: 6px; }}
+
     .input, .textarea {{
-      width: 100%;
-      background: var(--input-bg);
-      border: 1px solid var(--border);
-      color: var(--text);
-      outline: none; padding: 12px; font-size: 15px; font-family: 'Ubuntu', sans-serif;
+      width: 100%; background: var(--input-bg); border: 1px solid var(--border);
+      color: var(--text); outline: none; padding: 12px; font-size: 15px; border-radius: 4px;
     }}
-    .input:focus, .textarea:focus {{ border-color: var(--ubuntu-orange); }}
+    .input:focus, .textarea:focus {{ border-color: var(--ubuntu-orange); background: #111; }}
     .textarea {{ min-height: 120px; resize: vertical; }}
 
     .upload-wrap {{
-      background: var(--input-bg); border: 1px dashed var(--border);
-      min-height: 160px; display: flex; flex-direction: column;
-      align-items: center; justify-content: center; padding: 14px; cursor: pointer; text-align: center;
+      background: var(--input-bg); border: 2px dashed var(--border); border-radius: 6px;
+      min-height: 200px; display: flex; flex-direction: column; gap: 14px; padding: 14px; cursor: pointer;
+      transition: border-color 0.2s ease;
     }}
-    .upload-wrap:hover {{ border-color: var(--ubuntu-orange); }}
+    .upload-wrap.dragover {{ border-color: var(--ubuntu-orange); background: rgba(233, 84, 32, 0.05); }}
     .upload-wrap input[type="file"] {{ display: none; }}
     
-    .preview-grid {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(80px, 1fr)); gap: 8px; width:100%; display:none; }}
-    .thumb {{ position: relative; aspect-ratio: 1; overflow: hidden; border: 1px solid var(--border); }}
-    .thumb img {{ width: 100%; height: 100%; object-fit: cover; }}
+    .upload-placeholder {{
+      flex: 1; display: flex; flex-direction: column; align-items: center; justify-content: center;
+      gap: 12px; color: var(--muted); background: transparent; border: none; cursor: pointer;
+    }}
+    
+    .preview-grid {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(100px, 1fr)); gap: 10px; }}
+    .thumb {{ position: relative; aspect-ratio: 1/1; overflow: hidden; border: 1px solid var(--border); border-radius: 4px; }}
+    .thumb img {{ width: 100%; height: 100%; object-fit: cover; display: block; }}
     .thumb-remove {{
-      position: absolute; top: 4px; right: 4px; width: 20px; height: 20px;
-      background: rgba(0,0,0,0.8); color: white; border: none; cursor: pointer;
-      display: flex; align-items: center; justify-content: center; font-size: 12px; border-radius: 50%;
+      position: absolute; top: 4px; right: 4px; width: 24px; height: 24px;
+      background: rgba(0,0,0,0.7); border: none; color: white; border-radius: 50%;
+      cursor: pointer; font-size: 14px; display: flex; align-items: center; justify-content: center;
     }}
 
+    .advanced {{ border: 1px solid var(--border); background: var(--input-bg); border-radius: 4px; }}
     .advanced-toggle {{
-      width: 100%; height: 40px; border: 1px solid var(--border); background: var(--panel-2);
-      color: var(--text); display: flex; align-items: center; justify-content: space-between;
-      padding: 0 14px; cursor: pointer; font-weight: 500;
+      width: 100%; height: 48px; border: none; border-bottom: 1px solid var(--border);
+      background: transparent; color: var(--text); display: flex; align-items: center; justify-content: space-between;
+      padding: 0 16px; cursor: pointer; font-size: 14px; font-weight: 500;
     }}
-    .advanced-body {{ display: none; padding: 14px; background: var(--input-bg); border: 1px solid var(--border); border-top: none; }}
+    .advanced-body {{ display: none; padding: 16px; }}
     .advanced-body.open {{ display: block; }}
-    
-    .grid-2 {{ display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-bottom: 12px; }}
+    .advanced-grid {{ display: grid; grid-template-columns: 1fr 1fr; gap: 16px; }}
 
-    .actions {{ display: grid; grid-template-columns: 1fr; gap: 12px; margin-top: auto; }}
+    .actions {{ display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin-top: auto; }}
     .btn {{
-      height: 48px; border: none; font-size: 16px; font-weight: 700; cursor: pointer;
-      font-family: 'Ubuntu', sans-serif; transition: background 0.2s; color: white;
+      height: 48px; border: none; background: var(--panel-2); color: var(--text);
+      cursor: pointer; font-size: 15px; font-weight: 500; border-radius: 4px; transition: background 0.2s;
     }}
-    .btn-primary {{ background: var(--ubuntu-orange); }}
+    .btn:hover {{ background: var(--panel-3); }}
+    .btn-primary {{ background: var(--ubuntu-orange); color: white; }}
     .btn-primary:hover {{ background: var(--ubuntu-orange-hover); }}
 
-    /* --- Slider CSS --- */
-    .result-shell {{ position: relative; flex: 1; min-height: 400px; border: 1px solid var(--border); background: #000; overflow: hidden; display: flex; align-items: center; justify-content: center; }}
-    
-    .comparison-container {{
-      position: relative;
-      width: 100%;
-      height: 100%;
-      display: none; /* Hidden until loaded */
+    /* Result Stage & Slider */
+    .result-stage {{
+      position: relative; flex: 1; border: 1px solid var(--border); background: #0a0a0a;
+      border-radius: 6px; overflow: hidden; display: flex; align-items: center; justify-content: center;
     }}
     
-    .comparison-container img {{
-      position: absolute;
-      top: 0; left: 0;
-      width: 100%; height: 100%;
-      object-fit: contain;
-      pointer-events: none;
+    .slider-container {{
+      position: relative; width: 100%; height: 100%; display: none;
     }}
-    
-    #imgStd {{ z-index: 1; }}
-    #imgSml {{ z-index: 2; clip-path: polygon(0 0, 50% 0, 50% 100%, 0 100%); }}
-
-    .slider-input {{
-      position: absolute;
-      top: 0; left: 0;
-      width: 100%; height: 100%;
-      outline: none;
-      background: transparent;
-      -webkit-appearance: none;
-      margin: 0;
-      z-index: 4;
-      cursor: ew-resize;
+    .slider-container img {{
+      position: absolute; top: 0; left: 0; width: 100%; height: 100%;
+      object-fit: contain; pointer-events: none;
     }}
     
-    .slider-input::-webkit-slider-thumb {{
-      -webkit-appearance: none;
-      width: 40px; height: 1000px;
-      background: transparent;
-      cursor: ew-resize;
-    }}
-
-    .slider-line {{
-      position: absolute;
-      top: 0; bottom: 0; left: 50%;
-      width: 4px;
-      background: var(--ubuntu-orange);
-      z-index: 3;
-      transform: translateX(-50%);
-      pointer-events: none;
+    .resize-layer {{
+      position: absolute; top: 0; left: 0; width: 50%; height: 100%; overflow: hidden;
     }}
     
-    .slider-pill {{
-      position: absolute;
-      top: 50%; left: 50%;
-      transform: translate(-50%, -50%);
-      width: 32px; height: 32px;
-      background: var(--ubuntu-orange);
-      border-radius: 50%;
-      display: flex; align-items: center; justify-content: center;
-      z-index: 3; color: white;
+    .slider-handle {{
+      position: absolute; top: 0; bottom: 0; left: 50%; width: 2px;
+      background: var(--ubuntu-orange); cursor: ew-resize; transform: translateX(-50%); z-index: 10;
     }}
-
-    .decoder-label {{
-      position: absolute; top: 16px; padding: 6px 12px; background: rgba(0,0,0,0.6);
-      color: white; font-size: 13px; font-weight: 700; border-radius: 4px; z-index: 3;
-      backdrop-filter: blur(4px);
+    .slider-handle::after {{
+      content: '⟷'; position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%);
+      background: var(--ubuntu-orange); color: white; border-radius: 50%; width: 36px; height: 36px;
+      display: flex; align-items: center; justify-content: center; font-size: 16px; box-shadow: 0 0 10px rgba(0,0,0,0.5);
     }}
-    .decoder-label.left {{ left: 16px; }}
-    .decoder-label.right {{ right: 16px; }}
 
-    .empty-state {{
-      position: absolute; inset:0; display: flex; flex-direction: column; align-items: center;
-      justify-content: center; color: var(--text-dim); text-align: center; padding: 20px; z-index: 0;
+    .badge {{
+      position: absolute; bottom: 16px; padding: 6px 12px; background: rgba(0,0,0,0.7);
+      color: white; font-size: 13px; font-weight: 500; border-radius: 20px; z-index: 5;
+      border: 1px solid rgba(255,255,255,0.2); backdrop-filter: blur(4px);
     }}
+    .badge-left {{ left: 16px; }}
+    .badge-right {{ right: 16px; }}
 
+    .result-empty {{
+      display: flex; flex-direction: column; align-items: center; justify-content: center;
+      gap: 12px; color: var(--muted); text-align: center;
+    }}
     .loader {{
       position: absolute; inset: 0; display: none; align-items: center; justify-content: center;
-      flex-direction: column; gap: 14px; background: rgba(0,0,0,0.7); backdrop-filter: blur(5px); z-index: 10;
+      flex-direction: column; gap: 16px; background: rgba(17,17,17,0.7); backdrop-filter: blur(5px); z-index: 20;
+    }}
+    .spinner {{
+      width: 50px; height: 50px; border: 4px solid rgba(255,255,255,0.1);
+      border-top-color: var(--ubuntu-orange); border-radius: 50%; animation: spin 1s linear infinite;
     }}
-    .spinner {{ width: 40px; height: 40px; border: 4px solid rgba(255,255,255,0.2); border-top-color: var(--ubuntu-orange); border-radius: 50%; animation: spin 1s linear infinite; }}
+
+    .meta-bar {{ display: flex; gap: 16px; }}
+    .meta-item {{ flex: 1; padding: 12px; background: var(--panel-2); border: 1px solid var(--border); border-radius: 4px; }}
+    .meta-label {{ font-size: 12px; color: var(--muted); margin-bottom: 4px; }}
+    .meta-value {{ font-size: 14px; font-weight: 500; color: var(--text); }}
+
+    .examples-panel {{ margin-top: 24px; background: var(--panel); border: 1px solid var(--border); border-radius: 6px; overflow: hidden; }}
+    .examples-header {{ padding: 16px 20px; font-size: 18px; font-weight: 500; background: #1a1a1a; border-bottom: 1px solid var(--border); }}
+    .examples-grid {{ display: grid; grid-template-columns: repeat(4, 1fr); gap: 16px; padding: 20px; }}
+    .example-card {{ background: var(--input-bg); border: 1px solid var(--border); border-radius: 6px; cursor: pointer; overflow: hidden; transition: border-color 0.2s; }}
+    .example-card:hover {{ border-color: var(--ubuntu-orange); }}
+    .example-card img {{ width: 100%; aspect-ratio: 1; object-fit: cover; border-bottom: 1px solid var(--border); }}
+    .example-body {{ padding: 12px; font-size: 13px; color: var(--muted); line-height: 1.4; }}
+
     @keyframes spin {{ to {{ transform: rotate(360deg); }} }}
 
-    @media (max-width: 980px) {{
+    @media (max-width: 1024px) {{
       .layout {{ grid-template-columns: 1fr; }}
-      .panel {{ min-height: auto; }}
-      .result-shell {{ min-height: 500px; }}
+      .panel {{ height: auto; min-height: 500px; }}
+      .examples-grid {{ grid-template-columns: repeat(2, 1fr); }}
     }}
   </style>
 </head>
 <body>
-  <div class="topbar">FLUX.2 Experimental Space</div>
+
+  <div class="topbar">Flux.2 Klein 4B Comparison Suite</div>
 
   <div class="container">
-    <div class="hero">
-      <h1 class="title">Flux.2-4B-Decoder-Comparator</h1>
-      <p class="subtitle">Compare <b>FLUX.2-klein-4B</b> side-by-side with <a href="#" target="_blank">small decoder</a>.</p>
-    </div>
+    <section class="hero">
+      <h1 class="title">VAE Decoder Comparison</h1>
+      <p class="subtitle">Compare <strong>FLUX.2-klein-4B Standard VAE</strong> vs <strong>Small Decoder VAE</strong> seamlessly.</p>
+    </section>
 
-    <div class="layout">
+    <section class="layout">
+      
       <div class="panel">
-        <div class="panel-header"><h2 class="panel-title">Inputs</h2></div>
+        <div class="panel-header">
+          <h2 class="panel-title">Inputs</h2>
+          <span class="status-pill">Ready</span>
+        </div>
         <div class="panel-body">
           
           <div>
             <span class="label">Input Images</span>
             <div class="upload-wrap" id="uploadZone">
               <input id="fileInput" type="file" accept="image/*" multiple />
-              <div id="uploadPlaceholder">
-                <svg viewBox="0 0 24 24" width="32" height="32" fill="none" stroke="currentColor" stroke-width="2" style="margin-bottom:8px; color:var(--ubuntu-warm-grey);">
-                  <path d="M12 4v10"></path><path d="M8.5 7.5 12 4l3.5 3.5"></path><path d="M4 16.5h16"></path><path d="M6 20h12"></path>
+              <div class="upload-placeholder" id="uploadPlaceholder">
+                <svg width="32" height="32" fill="none" stroke="currentColor" stroke-width="2" viewBox="0 0 24 24">
+                  <path d="M12 4v16m-8-8h16"></path>
                 </svg>
-                <div style="font-weight:500;">Click or Drag images here</div>
+                <span>Click or drag images to upload</span>
               </div>
-              <div class="preview-grid" id="previewGrid"></div>
+              <div class="preview-grid" id="previewGrid" style="display:none;"></div>
             </div>
+            <div class="hint">Upload one or more images. The first image determines standard dimensions.</div>
           </div>
 
           <div>
-            <label class="label" for="prompt">Prompt</label>
-            <textarea id="prompt" class="textarea" placeholder="e.g., A black cat holding a sign that says hello world..."></textarea>
+            <label class="label" for="prompt">Generation Prompt</label>
+            <textarea id="prompt" class="textarea" placeholder="A black cat holding a sign that says hello world..."></textarea>
           </div>
 
-          <div>
-            <button class="advanced-toggle" id="advancedToggle" type="button">
-              <span>Advanced Settings</span><span id="advIcon">+</span>
+          <div class="advanced">
+            <button class="advanced-toggle" id="advancedToggle">
+              Advanced Configuration <span>+</span>
             </button>
             <div class="advanced-body" id="advancedBody">
-              <div class="grid-2">
-                <div><label class="label">Seed</label><input class="input" type="number" id="seed" value="0"></div>
-                <div><label class="label">Steps</label><input class="input" type="number" id="steps" value="4" min="1" max="20"></div>
-                <div><label class="label">Width</label><input class="input" type="number" id="width" value="1024" step="8"></div>
-                <div><label class="label">Height</label><input class="input" type="number" id="height" value="1024" step="8"></div>
-                <div style="grid-column: span 2;"><label class="label">Guidance</label><input class="input" type="number" id="guidance" value="1.0" step="0.1"></div>
+              <div class="advanced-grid">
+                <div>
+                  <label class="label">Width</label>
+                  <input id="width" class="input" type="number" step="8" value="1024" />
+                </div>
+                <div>
+                  <label class="label">Height</label>
+                  <input id="height" class="input" type="number" step="8" value="1024" />
+                </div>
+                <div>
+                  <label class="label">Inference Steps</label>
+                  <input id="steps" class="input" type="number" value="4" />
+                </div>
+                <div>
+                  <label class="label">Guidance Scale</label>
+                  <input id="guidance" class="input" type="number" step="0.1" value="1.0" />
+                </div>
+                <div>
+                  <label class="label">Seed</label>
+                  <input id="seed" class="input" type="number" value="0" />
+                </div>
+                <div style="display: flex; align-items: center; gap: 8px; margin-top: 24px;">
+                  <input id="randomizeSeed" type="checkbox" checked />
+                  <label for="randomizeSeed" style="font-size: 14px; color: #ddd;">Randomize Seed</label>
+                </div>
               </div>
-              <label style="display:flex; align-items:center; gap:8px; font-size:14px; color:var(--text-dim); cursor:pointer;">
-                <input type="checkbox" id="randSeed" checked> Randomize Seed
-              </label>
             </div>
           </div>
 
           <div class="actions">
-            <button class="btn btn-primary" id="runBtn" type="button">Run Comparison</button>
+            <button class="btn" id="clearBtn">Clear</button>
+            <button class="btn btn-primary" id="runBtn">Generate Comparison</button>
           </div>
         </div>
       </div>
 
       <div class="panel">
         <div class="panel-header">
-          <h2 class="panel-title">Comparison Result</h2>
-          <span style="font-size:13px; color:var(--muted);" id="metaText">Idle</span>
+          <h2 class="panel-title">Output</h2>
+          <span class="status-pill" id="resultStatus">Idle</span>
         </div>
-        <div class="panel-body" style="padding:0;">
-          <div class="result-shell">
-            
-            <div class="empty-state" id="emptyState">
-              <svg viewBox="0 0 24 24" width="48" height="48" fill="none" stroke="currentColor" stroke-width="1.5" style="margin-bottom:12px;">
-                <rect x="3" y="3" width="18" height="18" rx="2"></rect><circle cx="8.5" cy="8.5" r="1.5"></circle><path d="M21 15l-5-5L5 21"></path>
+        <div class="panel-body" style="padding-bottom: 0;">
+          
+          <div class="result-stage" id="resultStage">
+            <div class="result-empty" id="outputEmpty">
+              <svg width="48" height="48" fill="none" stroke="#555" stroke-width="1.5" viewBox="0 0 24 24">
+                <rect x="3" y="3" width="18" height="18" rx="2"></rect>
+                <circle cx="8.5" cy="8.5" r="1.5"></circle>
+                <path d="M21 15l-5-5L5 21"></path>
               </svg>
-              <div>Your outputs will appear here</div>
+              <span>Images will appear here</span>
             </div>
 
-            <div class="comparison-container" id="compContainer">
-              <img id="imgStd" alt="Standard Decoder">
-              <img id="imgSml" alt="Small Decoder">
+            <div class="slider-container" id="sliderContainer">
+              <img id="imgSmall" alt="Small Decoder Output" />
+              <div class="badge badge-right">Small Decoder</div>
               
-              <div class="decoder-label left">Standard VAE</div>
-              <div class="decoder-label right">Small Decoder</div>
-              
-              <div class="slider-line" id="sliderLine">
-                <div class="slider-pill">
-                  <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.5"><path d="M15 18l-6-6 6-6"/></svg>
-                  <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.5"><path d="M9 18l6-6-6-6"/></svg>
-                </div>
+              <div class="resize-layer" id="resizeLayer">
+                <img id="imgStandard" alt="Standard VAE Output" />
+                <div class="badge badge-left">Standard VAE</div>
               </div>
-              <input type="range" min="0" max="100" value="50" class="slider-input" id="compareSlider">
+              
+              <div class="slider-handle" id="sliderHandle"></div>
             </div>
 
-            <div class="loader" id="loader">
+            <div class="loader" id="loaderOverlay">
               <div class="spinner"></div>
-              <div style="font-weight:500;">Running parallel pipelines...</div>
+              <span style="font-weight: 500;">Processing Pipelines Concurrently...</span>
             </div>
+          </div>
 
+          <div class="meta-bar">
+            <div class="meta-item">
+              <div class="meta-label">Seed Used</div>
+              <div class="meta-value" id="usedSeed">-</div>
+            </div>
+            <div class="meta-item">
+              <div class="meta-label">Hardware</div>
+              <div class="meta-value" id="deviceValue">{DEVICE_LABEL}</div>
+            </div>
           </div>
+          <br/>
         </div>
       </div>
-    </div>
+    </section>
+
+    <section class="examples-panel">
+      <div class="examples-header">Example Prompts & Reference Images</div>
+      <div class="examples-grid" id="examplesGrid"></div>
+    </section>
   </div>
 
   <script>
-    const state = {{ files: [], advanced: false }};
-
-    const uploadZone = document.getElementById('uploadZone');
-    const fileInput = document.getElementById('fileInput');
-    const previewGrid = document.getElementById('previewGrid');
-    const uploadPlaceholder = document.getElementById('uploadPlaceholder');
+    const examples = {examples_json};
+    const state = {{ files: [] }};
+
+    // UI Elements
+    const uploadZone = document.getElementById("uploadZone");
+    const fileInput = document.getElementById("fileInput");
+    const uploadPlaceholder = document.getElementById("uploadPlaceholder");
+    const previewGrid = document.getElementById("previewGrid");
     
-    const advToggle = document.getElementById('advancedToggle');
-    const advBody = document.getElementById('advancedBody');
-    const advIcon = document.getElementById('advIcon');
+    const promptEl = document.getElementById("prompt");
+    const advancedToggle = document.getElementById("advancedToggle");
+    const advancedBody = document.getElementById("advancedBody");
     
-    const runBtn = document.getElementById('runBtn');
-    const loader = document.getElementById('loader');
+    const runBtn = document.getElementById("runBtn");
+    const clearBtn = document.getElementById("clearBtn");
     
-    const compContainer = document.getElementById('compContainer');
-    const emptyState = document.getElementById('emptyState');
-    const imgStd = document.getElementById('imgStd');
-    const imgSml = document.getElementById('imgSml');
-    const compareSlider = document.getElementById('compareSlider');
-    const sliderLine = document.getElementById('sliderLine');
-    const metaText = document.getElementById('metaText');
-
-    // -- File Upload Logic --
-    function renderFiles() {{
-      previewGrid.innerHTML = '';
-      if(state.files.length === 0) {{
-        uploadPlaceholder.style.display = 'block';
-        previewGrid.style.display = 'none';
+    const resultStatus = document.getElementById("resultStatus");
+    const outputEmpty = document.getElementById("outputEmpty");
+    const sliderContainer = document.getElementById("sliderContainer");
+    const imgStandard = document.getElementById("imgStandard");
+    const imgSmall = document.getElementById("imgSmall");
+    const loaderOverlay = document.getElementById("loaderOverlay");
+
+    // Toggle Advanced
+    advancedToggle.addEventListener("click", () => {{
+      advancedBody.classList.toggle("open");
+      advancedToggle.innerHTML = `Advanced Configuration <span>${{advancedBody.classList.contains('open') ? '−' : '+'}}</span>`;
+    }});
+
+    // Upload Logic
+    function renderPreviews() {{
+      previewGrid.innerHTML = "";
+      if (!state.files.length) {{
+        uploadPlaceholder.style.display = "flex";
+        previewGrid.style.display = "none";
         return;
       }}
-      uploadPlaceholder.style.display = 'none';
-      previewGrid.style.display = 'grid';
-      
+      uploadPlaceholder.style.display = "none";
+      previewGrid.style.display = "grid";
       state.files.forEach((file, idx) => {{
-        const div = document.createElement('div');
-        div.className = 'thumb';
-        const img = document.createElement('img');
-        img.src = URL.createObjectURL(file);
-        const btn = document.createElement('button');
-        btn.className = 'thumb-remove';
-        btn.innerHTML = '&times;';
-        btn.onclick = (e) => {{ e.stopPropagation(); state.files.splice(idx,1); renderFiles(); }};
-        div.appendChild(img); div.appendChild(btn);
-        previewGrid.appendChild(div);
+        const div = document.createElement("div"); div.className = "thumb";
+        const img = document.createElement("img"); img.src = URL.createObjectURL(file);
+        const btn = document.createElement("button"); btn.className = "thumb-remove"; btn.innerHTML = "×";
+        btn.onclick = (e) => {{ e.stopPropagation(); state.files.splice(idx, 1); renderPreviews(); }};
+        div.append(img, btn); previewGrid.appendChild(div);
       }});
     }}
 
-    uploadZone.onclick = (e) => {{ if(e.target === uploadZone || e.target.closest('#uploadPlaceholder')) fileInput.click(); }};
-    fileInput.onchange = (e) => {{ state.files.push(...Array.from(e.target.files)); renderFiles(); fileInput.value = ''; }};
-    uploadZone.ondragover = (e) => e.preventDefault();
-    uploadZone.ondrop = (e) => {{ e.preventDefault(); state.files.push(...Array.from(e.dataTransfer.files)); renderFiles(); }};
+    uploadZone.addEventListener("click", (e) => {{ if(e.target === uploadZone || e.target.closest('.upload-placeholder')) fileInput.click(); }});
+    fileInput.addEventListener("change", (e) => {{
+      const valid = Array.from(e.target.files).filter(f => f.type.startsWith("image/"));
+      state.files.push(...valid); renderPreviews(); fileInput.value = "";
+    }});
+    uploadZone.addEventListener("dragover", (e) => {{ e.preventDefault(); uploadZone.classList.add("dragover"); }});
+    uploadZone.addEventListener("dragleave", () => uploadZone.classList.remove("dragover"));
+    uploadZone.addEventListener("drop", (e) => {{
+      e.preventDefault(); uploadZone.classList.remove("dragover");
+      if(e.dataTransfer.files.length) {{
+        const valid = Array.from(e.dataTransfer.files).filter(f => f.type.startsWith("image/"));
+        state.files.push(...valid); renderPreviews();
+      }}
+    }});
 
-    // -- Advanced Toggle --
-    advToggle.onclick = () => {{
-      state.advanced = !state.advanced;
-      advBody.classList.toggle('open', state.advanced);
-      advIcon.textContent = state.advanced ? '−' : '+';
+    // Examples
+    async function loadExample(url, prompt) {{
+      try {{
+        const res = await fetch(url);
+        const blob = await res.blob();
+        state.files = [new File([blob], "example.jpg", {{ type: blob.type }})];
+        renderPreviews();
+        promptEl.value = prompt;
+      }} catch(err) {{ alert("Failed to load example."); }}
+    }}
+
+    const examplesGrid = document.getElementById("examplesGrid");
+    examples.forEach(ex => {{
+      const card = document.createElement("div"); card.className = "example-card";
+      card.innerHTML = `<img src="${{ex.url}}" alt="Example" /><div class="example-body">${{ex.prompt}}</div>`;
+      card.onclick = () => loadExample(ex.url, ex.prompt);
+      examplesGrid.appendChild(card);
+    }});
+
+    // Clear
+    clearBtn.onclick = () => {{
+      state.files = []; renderPreviews();
+      promptEl.value = ""; sliderContainer.style.display = "none";
+      outputEmpty.style.display = "flex"; resultStatus.innerText = "Idle";
+      resultStatus.classList.remove("active");
     }};
 
-    // -- Slider Logic --
-    compareSlider.addEventListener('input', (e) => {{
-      const val = e.target.value;
-      imgSml.style.clipPath = `polygon(0 0, ${{val}}% 0, ${{val}}% 100%, 0 100%)`;
-      sliderLine.style.left = `${{val}}%`;
+    // Slider Logic
+    const resizeLayer = document.getElementById("resizeLayer");
+    const sliderHandle = document.getElementById("sliderHandle");
+    let isDragging = false;
+    
+    sliderHandle.addEventListener("mousedown", () => isDragging = true);
+    document.addEventListener("mouseup", () => isDragging = false);
+    document.addEventListener("mousemove", (e) => {{
+      if (!isDragging) return;
+      const rect = sliderContainer.getBoundingClientRect();
+      let x = e.clientX - rect.left;
+      let percent = Math.max(0, Math.min(100, (x / rect.width) * 100));
+      resizeLayer.style.width = percent + "%";
+      sliderHandle.style.left = percent + "%";
     }});
 
-    // -- API Submission --
+    // Form Submission
     runBtn.onclick = async () => {{
-      const prompt = document.getElementById('prompt').value.trim();
-      if(!prompt) return alert('Please enter a prompt');
-
-      const formData = new FormData();
-      formData.append('prompt', prompt);
-      formData.append('seed', document.getElementById('seed').value);
-      formData.append('randomize_seed', document.getElementById('randSeed').checked);
-      formData.append('width', document.getElementById('width').value);
-      formData.append('height', document.getElementById('height').value);
-      formData.append('steps', document.getElementById('steps').value);
-      formData.append('guidance', document.getElementById('guidance').value);
+      if (!promptEl.value.trim()) return alert("Please enter a prompt.");
       
-      state.files.forEach(f => formData.append('images', f));
-
-      loader.style.display = 'flex';
+      const formData = new FormData();
+      formData.append("prompt", promptEl.value);
+      formData.append("width", document.getElementById("width").value);
+      formData.append("height", document.getElementById("height").value);
+      formData.append("steps", document.getElementById("steps").value);
+      formData.append("guidance", document.getElementById("guidance").value);
+      formData.append("seed", document.getElementById("seed").value);
+      formData.append("randomize_seed", document.getElementById("randomizeSeed").checked);
+      state.files.forEach(f => formData.append("images", f));
+
+      loaderOverlay.style.display = "flex";
+      resultStatus.innerText = "Processing..."; resultStatus.classList.add("active");
       runBtn.disabled = true;
 
       try {{
-        const res = await fetch('/api/compare', {{ method: 'POST', body: formData }});
+        const res = await fetch("/api/compare", {{ method: "POST", body: formData }});
         const data = await res.json();
         
-        if(!data.success) throw new Error(data.error);
+        if (!data.success) throw new Error(data.error);
+
+        // Reset Slider Position
+        resizeLayer.style.width = "50%";
+        sliderHandle.style.left = "50%";
 
-        // Load images into slider
-        imgStd.src = data.std_url + '?t=' + Date.now();
-        imgSml.src = data.sml_url + '?t=' + Date.now();
+        imgStandard.src = data.std_url + "?t=" + Date.now();
+        imgSmall.src = data.small_url + "?t=" + Date.now();
         
-        emptyState.style.display = 'none';
-        compContainer.style.display = 'block';
+        outputEmpty.style.display = "none";
+        sliderContainer.style.display = "block";
+        document.getElementById("usedSeed").innerText = data.seed;
+        resultStatus.innerText = "Completed";
         
-        // Reset slider to center
-        compareSlider.value = 50;
-        compareSlider.dispatchEvent(new Event('input'));
-
-        metaText.textContent = `Seed: ${{data.seed}} | Device: ${{data.device}}`;
-
-      }} catch(err) {{
-        alert('Error: ' + err.message);
+      }} catch (err) {{
+        alert("Error: " + err.message);
+        resultStatus.innerText = "Error";
       }} finally {{
-        loader.style.display = 'none';
+        loaderOverlay.style.display = "none";
         runBtn.disabled = false;
+        setTimeout(() => resultStatus.classList.remove("active"), 2000);
       }}
     }};
-
   </script>
 </body>
 </html>