Update app.py

app.py

@@ -1,248 +1,148 @@
 import os
+import io
 import gc
-import gradio as gr
-import numpy as np
+import uuid
+import json
+import base64
 import random
+import zipfile
+import threading
+import concurrent.futures
+from pathlib import Path
+from typing import List, Optional
+
 import spaces
+import numpy as np
 import torch
-from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 from PIL import Image
-from pathlib import Path
-import concurrent.futures
-import threading
-from typing import Iterable
-
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
-
-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()
+from gradio import Server
+from fastapi import Request, UploadFile, File, Form
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse, StreamingResponse
+from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 
-dtype  = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# --- App Configuration & Directories ---
+app = Server()
 
-MAX_SEED       = np.iinfo(np.int32).max
+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)
+
+MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
-EXAMPLES_DIR   = Path("examples")
 
+dtype  = torch.bfloat16
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if torch.cuda.is_available():
+    print("current device:", torch.cuda.current_device())
+    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
+    DEVICE_LABEL = torch.cuda.get_device_name(torch.cuda.current_device()).lower()
+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()
 
 print("Loading Small Decoder VAE...")
 vae_small = AutoencoderKLFlux2.from_pretrained(
     "black-forest-labs/FLUX.2-small-decoder",
     torch_dtype=dtype,
-)
+).to(device)
 
 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()
 
 pipe_lock_standard = threading.Lock()
 pipe_lock_small    = threading.Lock()
 
+# --- Utility Functions ---
 def calc_dimensions(pil_img: Image.Image):
-    """
-    Given a PIL image return (width, height) snapped to multiples of 8,
-    fitting within 1024 px on the long side, min 256 px on each side.
-    Uses round() so we match the reference app exactly.
-    """
     iw, ih = pil_img.size
     aspect = iw / ih
 
-    if aspect >= 1:          # landscape / square
+    if aspect >= 1:
         new_width  = 1024
         new_height = int(round(1024 / aspect))
-    else:                    # portrait
+    else:
         new_height = 1024
         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_height = max(256, min(1024, round(new_height / 8) * 8))
     return new_width, new_height
 
-
-def update_dimensions_from_image(image_list):
-    """
-    Called by the gallery .upload() event.
-    Returns updated slider values for width and height.
-    """
-    if not image_list:
-        return 1024, 1024
-
-    # gallery items arrive as PIL images when type="pil"
-    item = image_list[0]
-    img  = item[0] if isinstance(item, tuple) else item
-
-    if isinstance(img, str):
-        img = Image.open(img).convert("RGB")
-    elif not isinstance(img, Image.Image):
-        return 1024, 1024
-
-    return calc_dimensions(img)
-
-def parse_and_resize_images(input_images, width: int, height: int):
-    """
-    Parse the gallery input and resize every frame to (width, height).
-    Returns a list[PIL.Image] or None.
-    """
-    if input_images is None:
+def parse_and_resize_images(image_paths: List[str], width: int, height: int):
+    if not image_paths:
         return None
-
-    raw_list = []
-
-    if isinstance(input_images, str):
-        if os.path.exists(input_images):
-            raw_list = [Image.open(input_images).convert("RGB")]
-    elif isinstance(input_images, Image.Image):
-        raw_list = [input_images.convert("RGB")]
-    elif isinstance(input_images, list):
-        for item in input_images:
-            try:
-                src = item[0] if isinstance(item, tuple) else item
-                if isinstance(src, str):
-                    raw_list.append(Image.open(src).convert("RGB"))
-                elif isinstance(src, Image.Image):
-                    raw_list.append(src.convert("RGB"))
-                elif hasattr(src, "name"):
-                    raw_list.append(Image.open(src.name).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
+    
+    resized = []
+    for path in image_paths:
+        try:
+            img = Image.open(path).convert("RGB")
+            resized.append(img.resize((width, height), Image.LANCZOS))
+        except Exception as e:
+            print(f"Skipping invalid image: {e}")
+            
+    return resized if resized else None
 
 def run_pipeline(pipe, lock, kwargs, seed):
     with lock:
-        gen    = torch.Generator(device="cpu").manual_seed(seed)
+        gen = torch.Generator(device="cpu").manual_seed(seed)
         result = pipe(**kwargs, generator=gen).images[0]
     return result
 
+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
+
+# --- Inference Function ---
 @spaces.GPU(duration=120)
 def infer(
-    prompt,
-    input_images=None,
-    seed=42,
-    randomize_seed=False,
-    width=1024,
-    height=1024,
-    num_inference_steps=4,
-    guidance_scale=1.0,
-    progress=gr.Progress(track_tqdm=True),
+    prompt: str,
+    image_paths: List[str] = None,
+    seed: int = 42,
+    randomize_seed: bool = False,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 4,
+    guidance_scale: float = 1.0,
 ):
     gc.collect()
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
     if not prompt or not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
+        raise ValueError("Please enter a prompt.")
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # ── width / height: derive from the first uploaded image if present ──
     image_list = None
-    if input_images:
-        # Re-derive dimensions from the actual first image so they are
-        # always consistent with what the pipeline will receive.
-        item = (
-            input_images[0][0]
-            if isinstance(input_images[0], tuple)
-            else input_images[0]
-        )
-        if isinstance(item, str):
-            first_pil = Image.open(item).convert("RGB")
-        elif isinstance(item, Image.Image):
-            first_pil = item.convert("RGB")
-        else:
-            first_pil = None
-
-        if first_pil is not None:
+    if image_paths and len(image_paths) > 0:
+        try:
+            first_pil = Image.open(image_paths[0]).convert("RGB")
             width, height = calc_dimensions(first_pil)
+            image_list = parse_and_resize_images(image_paths, width, height)
+        except Exception as e:
+            print(f"Error processing upload: {e}")
 
-        # parse + resize all images to the final (width, height)
-        image_list = parse_and_resize_images(input_images, width, height)
-
-    # ensure dims are multiples of 8 even for text-only runs
     width  = max(256, min(MAX_IMAGE_SIZE, round(int(width)  / 8) * 8))
     height = max(256, min(MAX_IMAGE_SIZE, round(int(height) / 8) * 8))
 
@@ -256,230 +156,830 @@ def infer(
     if image_list is not None:
         shared_kwargs["image"] = image_list
 
-    progress(0.30, desc="Launching both pipelines simultaneously...")
-
     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
-        )
+        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,
         )
 
-    progress(0.80, desc="✅ Both pipelines done!")
-
     out_standard = future_std.result()
     out_small    = future_small.result()
 
     gc.collect()
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
     return out_standard, out_small, seed
 
 
-@spaces.GPU(duration=120)
-def infer_example(prompt):
-    out_std, out_small, seed_used = infer(
-        prompt=prompt,
-        input_images=None,
-        seed=0,
-        randomize_seed=True,
-        width=1024,
-        height=1024,
-        num_inference_steps=4,
-        guidance_scale=1.0,
-    )
-    return out_std, out_small, seed_used
-
-
+# --- FastAPI 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,
-                    "path":   str(EXAMPLES_DIR / name),
-                    "prompt": example_prompts.get(
-                        name, "Edit this image while preserving composition."
-                    ),
-                })
-    return items
-
-EXAMPLE_ITEMS = get_example_items()
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1100px;
-}
-#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;
-}
-"""
-
-with gr.Blocks() as demo:
-
-    with gr.Column(elem_id="col-container"):
-
-        gr.Markdown(
-            "# **Flux.2-4B-Decoder-Comparator**",
-            elem_id="main-title",
-        )
-        gr.Markdown(
-            "Compare **FLUX.2-klein-4B** side-by-side with "
-            "[small decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)."
-        )
+    return [
+        {
+            "urls": ["/example-file/I1.jpg", "/example-file/I2.jpg"],
+            "prompt": "Make her wear these glasses in Image 2."
+        },
+        {
+            "urls": ["/example-file/1.jpg"],
+            "prompt": "Change the weather to stormy."
+        },
+        {
+            "urls": ["/example-file/2.jpg"],
+            "prompt": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition."
+        },
+        {
+            "urls": ["/example-file/3.jpg"],
+            "prompt": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent."
+        },
+        {
+            "urls": ["/example-file/4.jpg"],
+            "prompt": "Make the texture high-resolution."
+        }
+    ]
 
-        with gr.Row(equal_height=True):
-
-            with gr.Column():
-                input_images = gr.Gallery(
-                    label="Input Images",
-                    type="pil",
-                    columns=2,
-                    rows=1,
-                    height=300,
-                    allow_preview=True,
-                )
-
-                prompt = gr.Text(
-                    label="Prompt",
-                    max_lines=1,
-                    show_label=True,
-                    placeholder="e.g., A black cat holding a sign that says hello world...",
-                )
-
-                run_button = gr.Button("Run Comparison", variant="primary")
-
-            with gr.Column():
-                with gr.Row():
-                    with gr.Column():
-                        result_standard = gr.Image(
-                            label="Standard Decoder",
-                            show_label=True,
-                            interactive=False,
-                            format="png",
-                            height=250,
-                        )
-                    with gr.Column():
-                        result_small = gr.Image(
-                            label="Small Decoder",
-                            show_label=True,
-                            interactive=False,
-                            format="png",
-                            height=250,
-                        )
-
-                seed_output = gr.Number(label="Seed Used", precision=0, visible=False)
-
-                with gr.Accordion("Advanced Settings", open=False, visible=False):
-                    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,
-                        )
-
-        gr.Examples(
-            examples=[
-                [["examples/I1.jpg", "examples/I2.jpg"], "Make her wear these glasses in Image 2."],
-                [["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",
-        )
+@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
+    if not path.exists():
+        return JSONResponse({"error": "File not found"}, status_code=404)
+    return FileResponse(path, filename=filename, media_type="image/png")
+
+@app.get("/api/download-zip")
+async def download_zip(std: str, small: str):
+    """Packages both generated images into a single ZIP file and streams it."""
+    std_name = Path(std).name
+    small_name = Path(small).name
+    
+    std_path = OUTPUT_DIR / std_name
+    small_path = OUTPUT_DIR / small_name
+    
+    if not std_path.exists() or not small_path.exists():
+        return JSONResponse({"error": "Generated files not found"}, status_code=404)
+        
+    memory_file = io.BytesIO()
+    with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
+        zf.write(std_path, arcname=f"Standard_Decoder_{std_name}")
+        zf.write(small_path, arcname=f"Small_Decoder_{small_name}")
+        
+    memory_file.seek(0)
+    
+    return StreamingResponse(
+        memory_file,
+        media_type="application/zip",
+        headers={"Content-Disposition": f"attachment; filename=Flux2_Comparison_{uuid.uuid4().hex[:6]}.zip"}
+    )
 
-        gr.Markdown(
-            "[*](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) "
-            "Experimental Space — FLUX.2 [klein] 4B VAE Decoder Comparison."
+@app.post("/api/compare")
+async def compare_images(
+    prompt: str = Form(...),
+    seed: str = Form("0"),
+    randomize_seed: str = Form("true"),
+    width: str = Form("1024"),
+    height: str = Form("1024"),
+    steps: str = Form("4"),
+    guidance: str = Form("1.0"),
+    images: Optional[List[UploadFile]] = File(None),
+):
+    temp_paths = []
+    try:
+        image_paths = []
+        if images:
+            for upload in images:
+                if not upload.filename: continue
+                suffix = Path(upload.filename).suffix or ".png"
+                temp_path = OUTPUT_DIR / f"upload_{uuid.uuid4().hex}{suffix}"
+                content = await upload.read()
+                with open(temp_path, "wb") as f:
+                    f.write(content)
+                temp_paths.append(str(temp_path))
+                image_paths.append(str(temp_path))
+
+        result_std, result_small, used_seed = infer(
+            prompt=prompt,
+            image_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),
         )
 
-    input_images.upload(
-        fn=update_dimensions_from_image,
-        inputs=[input_images],
-        outputs=[width, height_slider],
-    )
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            input_images,
-            seed,
-            randomize_seed,
-            width,
-            height_slider,
-            num_inference_steps,
-            guidance_scale,
-        ],
-        outputs=[result_standard, result_small, seed_output],
-    )
+        std_filename = save_image(result_std, prefix="std")
+        small_filename = save_image(result_small, prefix="small")
+
+        return JSONResponse({
+            "success": True,
+            "seed": used_seed,
+            "std_url": f"/download/{std_filename}",
+            "small_url": f"/download/{small_filename}",
+            "std_filename": std_filename,
+            "small_filename": small_filename,
+            "device": DEVICE_LABEL,
+        })
+
+    except Exception as e:
+        return JSONResponse({"success": False, "error": str(e)}, status_code=500)
+    finally:
+        for p in temp_paths:
+            if os.path.exists(p):
+                os.remove(p)
+
+# --- Frontend ---
+@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>
+  <link href="https://fonts.googleapis.com/css2?family=Ubuntu:wght@300;400;500;700&display=swap" rel="stylesheet">
+  <style>
+    :root {{
+      --ub-aubergine: #2C001E;
+      --ub-aubergine-dark: #1f0015;
+      --ub-orange: #E95420;
+      --ub-orange-hover: #c4461a;
+      --ub-panel: #3D3D3D;
+      --ub-panel-light: #4f4f4f;
+      --ub-border: rgba(255,255,255,0.1);
+      --ub-text: #FFFFFF;
+      --ub-muted: #b0b0b0;
+      --ub-input: #2b2b2b;
+      --panel-radius: 8px;
+    }}
+
+    * {{ box-sizing: border-box; font-family: 'Ubuntu', sans-serif; }}
+
+    body {{
+      margin: 0; padding: 0;
+      background: var(--ub-aubergine);
+      color: var(--ub-text);
+      min-height: 100vh;
+      display: flex;
+      flex-direction: column;
+    }}
+
+    .topbar {{
+      background: var(--ub-aubergine-dark);
+      padding: 16px 24px;
+      border-bottom: 1px solid var(--ub-border);
+      text-align: center;
+      font-weight: 700;
+      letter-spacing: 0.5px;
+      color: var(--ub-orange);
+    }}
+
+    .container {{
+      max-width: 1300px;
+      margin: 0 auto;
+      padding: 30px 20px;
+      flex: 1;
+      width: 100%;
+    }}
+
+    .header-text {{
+      text-align: center;
+      margin-bottom: 30px;
+    }}
+    .header-text h1 {{
+      margin: 0 0 10px 0;
+      font-size: 2.2rem;
+    }}
+    .header-text p {{
+      color: var(--ub-muted);
+      margin: 0;
+    }}
+
+    /* FIXED LAYOUT GRID */
+    .layout {{
+      display: grid;
+      grid-template-columns: 420px 1fr;
+      gap: 24px;
+      align-items: stretch;
+      height: 650px; 
+    }}
+
+    .panel {{
+      background: var(--ub-panel);
+      border-radius: var(--panel-radius);
+      box-shadow: 0 8px 24px rgba(0,0,0,0.2);
+      display: flex;
+      flex-direction: column;
+      overflow: hidden;
+      height: 100%;
+    }}
+
+    .panel-header {{
+      padding: 16px 20px;
+      background: rgba(0,0,0,0.2);
+      border-bottom: 1px solid var(--ub-border);
+      font-weight: 500;
+      font-size: 1.1rem;
+      flex-shrink: 0;
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+    }}
+
+    .panel-body-scroll {{
+      flex: 1;
+      padding: 20px;
+      overflow-y: auto;
+      display: flex;
+      flex-direction: column;
+    }}
+
+    /* Input Forms */
+    .form-group {{ margin-bottom: 20px; flex-shrink: 0; }}
+    .label {{
+      display: block; font-weight: 500; font-size: 14px;
+      color: var(--ub-muted); margin-bottom: 8px;
+    }}
+
+    .textarea, .input {{
+      width: 100%;
+      background: var(--ub-input);
+      border: 1px solid var(--ub-border);
+      color: var(--ub-text);
+      padding: 12px;
+      border-radius: 4px;
+      outline: none;
+      font-size: 14px;
+    }}
+    .textarea:focus, .input:focus {{ border-color: var(--ub-orange); }}
+    .textarea {{ min-height: 100px; resize: vertical; }}
+
+    /* Upload Zone */
+    .upload-zone {{
+      background: var(--ub-input);
+      border: 1px dashed var(--ub-muted);
+      border-radius: 4px;
+      padding: 15px;
+      text-align: center;
+      cursor: pointer;
+      transition: background 0.2s, border-color 0.2s;
+      min-height: 100px;
+      display: flex;
+      flex-direction: column;
+      justify-content: center;
+      align-items: center;
+    }}
+    .upload-zone:hover, .upload-zone.dragover {{
+      border-color: var(--ub-orange);
+      background: rgba(233,84,32,0.05);
+    }}
+    .upload-zone input[type="file"] {{ display: none; }}
+    .upload-text {{ pointer-events: none; color: var(--ub-muted); }}
+    
+    .preview-grid {{
+      display: none;
+      grid-template-columns: repeat(auto-fill, minmax(70px, 1fr));
+      gap: 10px;
+      width: 100%;
+    }}
+    .thumb {{
+      position: relative; aspect-ratio: 1;
+      border-radius: 4px; overflow: hidden;
+      border: 1px solid var(--ub-border);
+    }}
+    .thumb img {{ width: 100%; height: 100%; object-fit: cover; display: block; }}
+    .thumb-remove {{
+      position: absolute; top: 4px; right: 4px;
+      background: rgba(0,0,0,0.7); color: white;
+      border: none; border-radius: 50%; width: 20px; height: 20px;
+      display: flex; align-items: center; justify-content: center;
+      cursor: pointer; font-size: 12px;
+    }}
+    
+    .add-more-btn {{
+      display: flex; align-items: center; justify-content: center;
+      border: 2px dashed var(--ub-muted); border-radius: 4px;
+      color: var(--ub-muted); font-size: 26px; cursor: pointer;
+      aspect-ratio: 1; transition: all 0.2s; background: transparent;
+    }}
+    .add-more-btn:hover {{
+      border-color: var(--ub-orange); color: var(--ub-orange);
+      background: rgba(233,84,32,0.05);
+    }}
+
+    /* Advanced Accordion */
+    .advanced-toggle {{
+      width: 100%; background: none; border: none; color: var(--ub-orange);
+      text-align: left; padding: 10px 0; font-weight: 500; cursor: pointer;
+      display: flex; justify-content: space-between; align-items: center;
+      flex-shrink: 0;
+    }}
+    .advanced-icon {{ font-weight: bold; font-size: 18px; line-height: 1; }}
+    .advanced-body {{ display: none; padding-top: 10px; flex-shrink: 0; }}
+    .advanced-body.open {{ display: block; }}
+    .grid-2 {{ display: grid; grid-template-columns: 1fr 1fr; gap: 15px; }}
+
+    /* Status Container */
+    .status-container {{
+      margin-top: 20px; margin-bottom: 20px;
+      border: 1px solid var(--ub-border); border-radius: 4px;
+      background: #200014; display: flex; flex-direction: column;
+      flex: 1; min-height: 100px; max-height: 200px;
+    }}
+    .status-header {{
+      padding: 8px 12px; font-size: 11px; font-weight: 700;
+      color: var(--ub-muted); background: rgba(0,0,0,0.4);
+      border-bottom: 1px solid var(--ub-border); text-transform: uppercase;
+      letter-spacing: 0.5px; flex-shrink: 0;
+    }}
+    .status-log {{
+      padding: 10px; font-family: 'Courier New', Courier, monospace;
+      font-size: 12px; color: #eeeeee; overflow-y: auto;
+      flex: 1; display: flex; flex-direction: column; gap: 4px;
+    }}
+    .log-time {{ color: #777; margin-right: 8px; }}
+    .log-info {{ color: #5bc0eb; }}
+    .log-success {{ color: #9bc53d; }}
+    .log-error {{ color: #ff5e5b; }}
+
+    /* Buttons */
+    .btn {{
+      width: 100%; padding: 14px; border: none; border-radius: 4px;
+      font-size: 16px; font-weight: 700; cursor: pointer;
+      transition: opacity 0.2s, background 0.2s; flex-shrink: 0;
+    }}
+    .btn-primary {{
+      background: var(--ub-orange); color: white;
+      box-shadow: 0 4px 12px rgba(233,84,32,0.3);
+    }}
+    .btn-primary:hover {{ background: var(--ub-orange-hover); }}
+    .btn:disabled {{ opacity: 0.6; cursor: not-allowed; }}
+
+    /* Top-Right Download Icon */
+    .action-icon {{
+      display: none; background: none; border: none; color: var(--ub-muted);
+      cursor: pointer; padding: 4px; transition: color 0.2s;
+    }}
+    .action-icon:hover {{ color: var(--ub-orange); }}
+
+    /* SLIDER CONTAINER */
+    .panel-body-slider {{
+      flex: 1; display: flex; flex-direction: column;
+      padding: 0; position: relative;
+    }}
+    .slider-stage {{
+      position: absolute; top: 0; left: 0; right: 0; bottom: 0;
+      background: #111; overflow: hidden; display: flex;
+      align-items: center; justify-content: center;
+    }}
+    .slider-empty {{ color: var(--ub-muted); text-align: center; z-index: 1; }}
+    
+    .slider-img {{
+      position: absolute; top: 0; left: 0; width: 100%; height: 100%;
+      object-fit: contain; display: none; user-select: none; -webkit-user-drag: none;
+    }}
+    #imgSmall {{ clip-path: inset(0 50% 0 0); }}
+
+    .slider-handle {{
+      position: absolute; left: 50%; top: 0; bottom: 0;
+      width: 4px; background: var(--ub-orange); cursor: ew-resize; display: none; z-index: 10;
+    }}
+    .slider-handle::after {{
+      content: '◀ ▶'; position: absolute; top: 50%; left: 50%;
+      transform: translate(-50%, -50%); width: 40px; height: 30px;
+      background: var(--ub-orange); color: white; border-radius: 15px;
+      display: flex; align-items: center; justify-content: center;
+      font-size: 10px; font-weight: bold; box-shadow: 0 2px 6px rgba(0,0,0,0.5);
+    }}
+
+    .slider-labels {{
+      position: absolute; top: 15px; left: 15px; right: 15px;
+      display: none; justify-content: space-between;
+      pointer-events: none; z-index: 5;
+    }}
+    .badge {{
+      background: rgba(0,0,0,0.6); color: white; padding: 6px 12px;
+      border-radius: 20px; font-size: 13px; backdrop-filter: blur(4px);
+    }}
+
+    /* UPDATED LOADER ANIMATION (Minimalist Single Circle) */
+    .loader {{
+      position: absolute; inset: 0; 
+      background: rgba(20, 0, 10, 0.7); /* dark aubergine tint */
+      backdrop-filter: blur(6px);
+      display: none; flex-direction: column; 
+      align-items: center; justify-content: center; z-index: 20;
+    }}
+    .spinner-single {{
+      width: 55px; height: 55px;
+      border: 3px solid rgba(255, 255, 255, 0.1);
+      border-top-color: var(--ub-orange);
+      border-radius: 50%;
+      animation: spin 1s cubic-bezier(0.4, 0.0, 0.2, 1) infinite;
+      margin-bottom: 20px;
+    }}
+    .loader-text {{
+      font-weight: 500;
+      font-size: 15px;
+      color: #ffffff;
+      letter-spacing: 1px;
+      animation: pulse 1.5s ease-in-out infinite;
+    }}
+    @keyframes pulse {{
+      0%, 100% {{ opacity: 1; }}
+      50% {{ opacity: 0.5; }}
+    }}
+    @keyframes spin {{ 
+      to {{ transform: rotate(360deg); }} 
+    }}
+
+    /* Examples */
+    .examples-section {{ margin-top: 40px; }}
+    .examples-section h3 {{ border-bottom: 1px solid var(--ub-border); padding-bottom: 10px; }}
+    .examples-grid {{
+      display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 20px;
+    }}
+    .ex-card {{
+      background: var(--ub-panel); border-radius: 4px; overflow: hidden;
+      cursor: pointer; transition: transform 0.2s, box-shadow 0.2s;
+    }}
+    .ex-card:hover {{ transform: translateY(-3px); box-shadow: 0 6px 16px rgba(0,0,0,0.3); }}
+    .ex-card-img-wrap {{ width: 100%; aspect-ratio: 1; display: flex; background: #000; }}
+    .ex-card-img-wrap img {{ height: 100%; object-fit: cover; }}
+    .ex-card p {{ padding: 12px; margin: 0; font-size: 13px; color: var(--ub-muted); line-height: 1.4; }}
+
+    @media (max-width: 900px) {{
+      .layout {{ grid-template-columns: 1fr; height: auto; }}
+      .panel-body-slider {{ height: 450px; flex: none; }}
+      .slider-stage {{ position: relative; height: 100%; }}
+    }}
+  </style>
+</head>
+<body>
+
+  <div class="topbar">Flux.2-4B VAE Decoder Comparator</div>
+
+  <div class="container">
+    <div class="header-text">
+      <h1>Standard vs. Small Decoder</h1>
+      <p>Upload an image, enter a prompt, and use the slider to compare outputs in real-time.</p>
+    </div>
+
+    <div class="layout">
+      <div class="panel">
+        <div class="panel-header">Settings</div>
+        <div class="panel-body-scroll">
+          <div class="form-group">
+            <label class="label">Input Images (Optional)</label>
+            <div class="upload-zone" id="dropZone">
+              <input type="file" id="fileInput" multiple accept="image/*" />
+              <div class="upload-text" id="uploadText">Click or Drag & Drop images here</div>
+              <div class="preview-grid" id="previewGrid"></div>
+            </div>
+          </div>
+
+          <div class="form-group">
+            <label class="label">Prompt</label>
+            <textarea id="promptInput" class="textarea" placeholder="Describe the edit or generation..."></textarea>
+          </div>
+
+          <button class="advanced-toggle" id="advToggle">
+            <span>Advanced Settings</span> <span class="advanced-icon" id="advIcon">+</span>
+          </button>
+          
+          <div class="advanced-body" id="advBody">
+            <div class="grid-2">
+              <div class="form-group">
+                <label class="label">Seed</label>
+                <input type="number" id="seed" class="input" value="0">
+              </div>
+              <div class="form-group">
+                <label class="label">Steps</label>
+                <input type="number" id="steps" class="input" value="4">
+              </div>
+              <div class="form-group">
+                <label class="label">Width</label>
+                <input type="number" id="width" class="input" value="1024" step="8">
+              </div>
+              <div class="form-group">
+                <label class="label">Height</label>
+                <input type="number" id="height" class="input" value="1024" step="8">
+              </div>
+              <div class="form-group" style="grid-column: span 2;">
+                <label class="label">Guidance Scale</label>
+                <input type="number" id="guidance" class="input" value="1.0" step="0.1">
+              </div>
+              <div class="form-group" style="grid-column: span 2;">
+                <label style="display:flex; align-items:center; gap:8px; font-size:14px; color:var(--ub-text);">
+                  <input type="checkbox" id="randomize" checked> Randomize Seed
+                </label>
+              </div>
+            </div>
+          </div>
+
+          <div class="status-container">
+            <div class="status-header">Execution Log</div>
+            <div class="status-log" id="statusLog">
+              <div><span class="log-time">[{DEVICE_LABEL}]</span><span>System Ready...</span></div>
+            </div>
+          </div>
+
+          <button class="btn btn-primary" id="runBtn">Run Comparison</button>
+        </div>
+      </div>
+
+      <div class="panel">
+        <div class="panel-header">
+          <span>Comparison View</span>
+          <button id="downloadZipBtn" class="action-icon" title="Download Both Images (ZIP)">
+            <svg width="22" height="22" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" viewBox="0 0 24 24">
+              <path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"></path>
+              <polyline points="7 10 12 15 17 10"></polyline>
+              <line x1="12" y1="15" x2="12" y2="3"></line>
+            </svg>
+          </button>
+        </div>
+        <div class="panel-body-slider">
+          <div class="slider-stage" id="sliderStage">
+            <div class="slider-empty" id="sliderEmpty">
+              <svg width="48" height="48" fill="none" stroke="currentColor" stroke-width="1.5" viewBox="0 0 24 24" style="margin-bottom:10px; opacity:0.5;">
+                <path d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z"></path>
+              </svg>
+              <div>Results will appear here</div>
+            </div>
+
+            <img id="imgStd" class="slider-img" alt="Standard Decoder" />
+            <img id="imgSmall" class="slider-img" alt="Small Decoder" />
+            
+            <div class="slider-labels" id="sliderLabels">
+              <div class="badge">Standard Decoder</div>
+              <div class="badge">Small Decoder</div>
+            </div>
+
+            <div class="slider-handle" id="sliderHandle"></div>
+
+            <div class="loader" id="loader">
+              <div class="spinner-single"></div>
+              <div class="loader-text">Running both decoders...</div>
+            </div>
+          </div>
+        </div>
+      </div>
+    </div>
+
+    <div class="examples-section">
+      <h3>Examples</h3>
+      <div class="examples-grid" id="examplesGrid"></div>
+    </div>
+  </div>
+
+  <script>
+    const examples = {examples_json};
+    let filesState = [];
+    let currentStdFilename = "";
+    let currentSmallFilename = "";
+
+    // UI Elements
+    const dropZone = document.getElementById('dropZone');
+    const fileInput = document.getElementById('fileInput');
+    const previewGrid = document.getElementById('previewGrid');
+    const uploadText = document.getElementById('uploadText');
+    const promptInput = document.getElementById('promptInput');
+    const runBtn = document.getElementById('runBtn');
+    const downloadZipBtn = document.getElementById('downloadZipBtn');
+    
+    // Status Log
+    const statusLog = document.getElementById('statusLog');
+
+    function logMsg(msg, styleClass="") {{
+      const div = document.createElement('div');
+      const timeStr = new Date().toLocaleTimeString('en-US', {{hour12:false}});
+      div.innerHTML = `<span class="log-time">[${{timeStr}}]</span><span class="${{styleClass}}">${{msg}}</span>`;
+      statusLog.appendChild(div);
+      statusLog.scrollTop = statusLog.scrollHeight; // auto-scroll to bottom
+    }}
+
+    // Slider Elements
+    const sliderStage = document.getElementById('sliderStage');
+    const imgStd = document.getElementById('imgStd');
+    const imgSmall = document.getElementById('imgSmall');
+    const sliderHandle = document.getElementById('sliderHandle');
+    const sliderLabels = document.getElementById('sliderLabels');
+    const sliderEmpty = document.getElementById('sliderEmpty');
+    const loader = document.getElementById('loader');
+
+    // Advanced Toggle logic (+ / -)
+    document.getElementById('advToggle').onclick = function() {{
+      const body = document.getElementById('advBody');
+      body.classList.toggle('open');
+      document.getElementById('advIcon').innerText = body.classList.contains('open') ? '−' : '+';
+    }};
+
+    // --- File Upload Logic ---
+    function renderPreviews() {{
+      previewGrid.innerHTML = '';
+      if(filesState.length > 0) {{
+        uploadText.style.display = 'none';
+        previewGrid.style.display = 'grid';
+        
+        filesState.forEach((f, i) => {{
+          const div = document.createElement('div');
+          div.className = 'thumb';
+          const img = document.createElement('img');
+          img.src = URL.createObjectURL(f);
+          const btn = document.createElement('button');
+          btn.className = 'thumb-remove';
+          btn.innerText = '×';
+          btn.onclick = (e) => {{ e.stopPropagation(); filesState.splice(i, 1); renderPreviews(); }};
+          div.appendChild(img); div.appendChild(btn);
+          previewGrid.appendChild(div);
+        }});
+        
+        // Append dynamic + button
+        const addBtn = document.createElement('div');
+        addBtn.className = 'add-more-btn';
+        addBtn.innerHTML = '+';
+        addBtn.onclick = (e) => {{ e.stopPropagation(); fileInput.click(); }};
+        previewGrid.appendChild(addBtn);
+        
+      }} else {{
+        uploadText.style.display = 'block';
+        previewGrid.style.display = 'none';
+      }}
+    }}
+
+    dropZone.onclick = (e) => {{ if(e.target === dropZone || e.target === uploadText) fileInput.click(); }};
+    fileInput.onchange = (e) => {{ filesState.push(...Array.from(e.target.files)); renderPreviews(); fileInput.value=''; }};
+    dropZone.ondragover = (e) => {{ e.preventDefault(); dropZone.classList.add('dragover'); }};
+    dropZone.ondragleave = () => dropZone.classList.remove('dragover');
+    dropZone.ondrop = (e) => {{
+      e.preventDefault(); dropZone.classList.remove('dragover');
+      if(e.dataTransfer.files.length) {{ filesState.push(...Array.from(e.dataTransfer.files)); renderPreviews(); }}
+    }};
+
+    // --- Examples Logic ---
+    async function loadExample(urls, text) {{
+      filesState = [];
+      renderPreviews();
+      promptInput.value = text;
+      logMsg("Loading example: " + text, "log-info");
+      
+      try {{
+        for(let i=0; i<urls.length; i++) {{
+            const res = await fetch(urls[i]);
+            const blob = await res.blob();
+            const filename = urls[i].split('/').pop();
+            filesState.push(new File([blob], filename, {{type: blob.type}}));
+        }}
+        renderPreviews();
+        
+        window.scrollTo({{top: 0, behavior: 'smooth'}});
+        
+        setTimeout(() => {{
+           logMsg("Example loaded. Starting comparison...", "log-info");
+           runBtn.click();
+        }}, 500);
+
+      }} catch (e) {{ 
+        logMsg("Failed to load example images.", "log-error");
+        alert('Failed to load example image.'); 
+      }}
+    }}
+
+    const exGrid = document.getElementById('examplesGrid');
+    examples.forEach(ex => {{
+      const card = document.createElement('div');
+      card.className = 'ex-card';
+      
+      let imgHTML = '';
+      if(ex.urls.length > 1) {{
+        imgHTML = `
+          <div class="ex-card-img-wrap">
+             <img src="${{ex.urls[0]}}" style="width:50%; border-right:1px solid #000;">
+             <img src="${{ex.urls[1]}}" style="width:50%;">
+          </div>
+        `;
+      }} else {{
+        imgHTML = `<div class="ex-card-img-wrap"><img src="${{ex.urls[0]}}" style="width:100%;"></div>`;
+      }}
+      
+      card.innerHTML = `${{imgHTML}}<p>${{ex.prompt}}</p>`;
+      card.onclick = () => loadExample(ex.urls, ex.prompt);
+      exGrid.appendChild(card);
+    }});
+
+    // --- Image Slider Logic ---
+    let isDragging = false;
+    
+    function updateSlider(clientX) {{
+      const rect = sliderStage.getBoundingClientRect();
+      let pos = Math.max(0, Math.min(clientX - rect.left, rect.width));
+      let percent = (pos / rect.width) * 100;
+      
+      sliderHandle.style.left = percent + '%';
+      imgSmall.style.clipPath = `inset(0 ${{100 - percent}}% 0 0)`;
+    }}
+
+    sliderHandle.addEventListener('mousedown', () => isDragging = true);
+    window.addEventListener('mouseup', () => isDragging = false);
+    window.addEventListener('mousemove', (e) => {{
+      if (!isDragging) return;
+      updateSlider(e.clientX);
+    }});
+    
+    sliderHandle.addEventListener('touchstart', () => isDragging = true);
+    window.addEventListener('touchend', () => isDragging = false);
+    window.addEventListener('touchmove', (e) => {{
+      if (!isDragging) return;
+      updateSlider(e.touches[0].clientX);
+    }});
+
+    // --- Download Zip Logic ---
+    downloadZipBtn.onclick = () => {{
+      if(!currentStdFilename || !currentSmallFilename) return;
+      logMsg("Initiating ZIP download...", "log-info");
+      window.location.href = `/api/download-zip?std=${{currentStdFilename}}&small=${{currentSmallFilename}}`;
+    }};
+
+    // --- Form Submission ---
+    runBtn.onclick = async () => {{
+      const prompt = promptInput.value.trim();
+      if(!prompt) {{
+        logMsg("Validation failed: Prompt is empty.", "log-error");
+        return alert("Enter a prompt");
+      }}
+
+      logMsg("Initializing generation sequence...", "log-info");
+
+      const fd = new FormData();
+      fd.append('prompt', prompt);
+      fd.append('seed', document.getElementById('seed').value);
+      fd.append('randomize_seed', document.getElementById('randomize').checked);
+      fd.append('width', document.getElementById('width').value);
+      fd.append('height', document.getElementById('height').value);
+      fd.append('steps', document.getElementById('steps').value);
+      fd.append('guidance', document.getElementById('guidance').value);
+      
+      filesState.forEach(f => fd.append('images', f));
+
+      loader.style.display = 'flex';
+      runBtn.disabled = true;
+      downloadZipBtn.style.display = 'none';
+
+      logMsg("Sending request to backend. Running both VAE models...", "log-info");
+
+      try {{
+        const res = await fetch('/api/compare', {{ method: 'POST', body: fd }});
+        const data = await res.json();
+        
+        if(data.success) {{
+          logMsg(`Success! Inference completed. Used seed: ${{data.seed}}`, "log-success");
+          
+          currentStdFilename = data.std_filename;
+          currentSmallFilename = data.small_filename;
+
+          imgStd.src = data.std_url;
+          imgSmall.src = data.small_url;
+          
+          imgStd.onload = () => {{
+            sliderEmpty.style.display = 'none';
+            imgStd.style.display = 'block';
+            imgSmall.style.display = 'block';
+            sliderHandle.style.display = 'block';
+            sliderLabels.style.display = 'flex';
+            downloadZipBtn.style.display = 'block'; // Reveal download button
+            
+            // Reset slider to center
+            const rect = sliderStage.getBoundingClientRect();
+            updateSlider(rect.left + rect.width / 2);
+          }};
+        }} else {{
+          logMsg("Error processing request: " + data.error, "log-error");
+          alert('Error: ' + data.error);
+        }}
+      }} catch(e) {{
+        logMsg("Network or server connection failed.", "log-error");
+        alert('Failed to connect to server.');
+      }} finally {{
+        loader.style.display = 'none';
+        runBtn.disabled = false;
+        logMsg("Sequence finished. Ready for next input.", "");
+      }}
+    }};
+  </script>
+</body>
+</html>
+"""
 
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch(
-        theme=orange_red_theme, css=css,
-        mcp_server=True,
-        ssr_mode=False,
-        show_error=True,
-    )
+app.launch()