Update app.py

app.py

@@ -1,77 +1,151 @@
+# app.py — Qwen/Qwen-Image-Layered (ZeroGPU-friendly) with:
+# - Decompose
+# - Refine selected layer (recursive / hierarchical)
+# - History "nodes" (tree-like lineage via parent)
+# - Switch to any node
+# - Export ZIP/PPTX for selected node
+# - Layer selection UX: dropdown + 1-row thumbnail gallery (Photoshop-like)
+# - Auto-open "Refined layers" accordion after refine
+
 import os
 import uuid
-import numpy as np
+import json
+import time
 import random
+import shutil
 import tempfile
 import zipfile
 import threading
-from datetime import datetime
+from typing import Dict, Any, List, Tuple, Optional
 
+import numpy as np
 import spaces
 import torch
 import gradio as gr
 from PIL import Image
-
-from diffusers import QwenImageLayeredPipeline
 from pptx import Presentation
+from huggingface_hub import login
+from diffusers import QwenImageLayeredPipeline
 
-
-LOG_DIR = "/tmp/local"
+# -------------------------
+# Config
+# -------------------------
+APP_TMP_ROOT = "/tmp/qwen_image_layered"
 MAX_SEED = np.iinfo(np.int32).max
 
-# Optional HF login (Spaces secret env var "hf")
-from huggingface_hub import login
-_HF_TOKEN = os.environ.get("hf")
-if _HF_TOKEN:
-    login(token=_HF_TOKEN)
+# (Optional) HF token as Spaces secret env var "hf"
+HF_TOKEN = os.environ.get("hf")
+if HF_TOKEN:
+    try:
+        login(token=HF_TOKEN)
+    except Exception as e:
+        print("HF login failed (continuing):", e)
 
+# Reduce allocator fragmentation (new name; old PYTORCH_CUDA_ALLOC_CONF is deprecated)
+os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
 
-# ----------------------------
-# Helpers
-# ----------------------------
-def ensure_dirname(path: str):
-    if path and not os.path.exists(path):
-        os.makedirs(path, exist_ok=True)
+SESSION_ID = uuid.uuid4().hex[:10]
+SESSION_DIR = os.path.join(APP_TMP_ROOT, f"session_{SESSION_ID}")
+os.makedirs(SESSION_DIR, exist_ok=True)
 
+# -------------------------
+# Pipeline (lazy + cached)
+# -------------------------
+_PIPELINE = None
+_PIPELINE_LOCK = threading.Lock()
 
-def px_to_emu(px, dpi=96):
-    inch = px / dpi
-    return int(inch * 914400)
+def _has_cuda() -> bool:
+    return torch.cuda.is_available()
 
+def get_pipeline() -> QwenImageLayeredPipeline:
+    """
+    Lazy-load the pipeline the first time we actually run on GPU (ZeroGPU attaches GPU only inside @spaces.GPU).
+    We keep the pipeline cached for the lifetime of the Space process.
+    """
+    global _PIPELINE
+    if _PIPELINE is not None:
+        return _PIPELINE
 
-def imagelist_to_pptx(img_files):
-    with Image.open(img_files[0]) as img:
-        w, h = img.size
+    with _PIPELINE_LOCK:
+        if _PIPELINE is not None:
+            return _PIPELINE
 
-    prs = Presentation()
-    prs.slide_width = px_to_emu(w)
-    prs.slide_height = px_to_emu(h)
+        cuda = _has_cuda()
+        # Prefer bf16 on NVIDIA; fallback to fp16 (some GPUs/drivers might not like bf16)
+        preferred_dtype = torch.bfloat16 if cuda else torch.float32
 
-    slide = prs.slides.add_slide(prs.slide_layouts[6])
-    left = top = 0
-
-    # Stack all images on top of each other (layers)
-    for img_path in img_files:
-        slide.shapes.add_picture(
-            img_path,
-            left,
-            top,
-            width=px_to_emu(w),
-            height=px_to_emu(h),
+        load_kwargs = dict(
+            torch_dtype=preferred_dtype,
+            low_cpu_mem_usage=True,
         )
 
-    with tempfile.NamedTemporaryFile(suffix=".pptx", delete=False) as tmp:
-        prs.save(tmp.name)
-        return tmp.name
-
-
-def make_zip(paths, prefix="layer"):
-    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmpzip:
-        with zipfile.ZipFile(tmpzip.name, "w", zipfile.ZIP_DEFLATED) as z:
-            for i, p in enumerate(paths):
-                z.write(p, f"{prefix}_{i+1}.png")
-        return tmpzip.name
-
+        # Some diffusers builds/pipelines accept device_map; some don't.
+        # Try to use it to reduce peak RAM when possible.
+        try:
+            _PIPELINE = QwenImageLayeredPipeline.from_pretrained(
+                "Qwen/Qwen-Image-Layered",
+                **load_kwargs,
+                device_map="auto",
+            )
+            print("Loaded pipeline with device_map=auto")
+        except TypeError:
+            _PIPELINE = QwenImageLayeredPipeline.from_pretrained(
+                "Qwen/Qwen-Image-Layered",
+                **load_kwargs,
+            )
+            print("Loaded pipeline without device_map")
+
+        # Memory savers (guarded)
+        if hasattr(_PIPELINE, "enable_attention_slicing"):
+            try:
+                _PIPELINE.enable_attention_slicing()
+            except Exception:
+                pass
+
+        # VAE slicing exists on some pipelines; on your logs it may be missing -> guard
+        if hasattr(_PIPELINE, "enable_vae_slicing"):
+            try:
+                _PIPELINE.enable_vae_slicing()
+            except Exception:
+                pass
+        elif hasattr(_PIPELINE, "vae") and hasattr(_PIPELINE.vae, "enable_slicing"):
+            try:
+                _PIPELINE.vae.enable_slicing()
+            except Exception:
+                pass
+
+        # Place pipeline
+        if cuda:
+            # If enable_model_cpu_offload exists, it helps VRAM usage a lot
+            if hasattr(_PIPELINE, "enable_model_cpu_offload"):
+                try:
+                    _PIPELINE.enable_model_cpu_offload()
+                    print("Enabled model CPU offload")
+                except Exception as e:
+                    print("CPU offload failed, falling back to .to('cuda'):", e)
+                    _PIPELINE.to("cuda")
+            else:
+                _PIPELINE.to("cuda")
+        else:
+            _PIPELINE.to("cpu")
+
+        return _PIPELINE
+
+
+# -------------------------
+# Helpers: filesystem nodes
+# -------------------------
+def _now_id() -> str:
+    return time.strftime("%H%M%S") + "_" + uuid.uuid4().hex[:6]
+
+def ensure_dir(path: str):
+    os.makedirs(path, exist_ok=True)
+
+def sanitize_prompt(x):
+    if x is None:
+        return None
+    s = str(x).strip()
+    return s if s else None
 
 def clamp_int(x, default: int, lo: int, hi: int) -> int:
     try:
@@ -80,708 +154,672 @@ def clamp_int(x, default: int, lo: int, hi: int) -> int:
         v = default
     return max(lo, min(hi, v))
 
+def px_to_emu(px, dpi=96):
+    inch = px / dpi
+    return int(inch * 914400)
 
-def norm_resolution(x):
-    v = clamp_int(x, default=1024, lo=640, hi=1024)
-    return v if v in (640, 1024) else 1024
-
-
-def load_rgba(path: str) -> Image.Image:
-    return Image.open(path).convert("RGBA")
-
-
-def labels_for_layers(n: int):
-    return [f"Layer {i}" for i in range(1, n + 1)]
-
-
-def parse_layer_label(label: str, default_idx0: int = 0) -> int:
-    # "Layer 3" -> 2
-    if not label:
-        return default_idx0
-    try:
-        num = int(label.strip().split()[-1])
-        return max(0, num - 1)
-    except Exception:
-        return default_idx0
-
-
-def now_str():
-    return datetime.utcnow().strftime("%H:%M:%S")
-
-
-def short_id():
-    return uuid.uuid4().hex[:8]
-
-
-def find_node(history, node_id: str):
-    for n in history:
-        if n["id"] == node_id:
-            return n
-    return None
+def imagelist_to_pptx(img_paths: List[str]) -> str:
+    # Slide size = first image size
+    with Image.open(img_paths[0]) as img:
+        w, h = img.size
 
+    prs = Presentation()
+    prs.slide_width = px_to_emu(w)
+    prs.slide_height = px_to_emu(h)
+    slide = prs.slides.add_slide(prs.slide_layouts[6])
 
-def compute_depth_and_path(history, node_id: str):
-    n = find_node(history, node_id)
-    if not n:
-        return 0, []
-    depth = 0
-    path = [n["title"]]
-    cur = n
-    while cur.get("parent_id"):
-        parent = find_node(history, cur["parent_id"])
-        if not parent:
-            break
-        depth += 1
-        path.append(parent["title"])
-        cur = parent
-    path.reverse()
-    return depth, path
+    for p in img_paths:
+        slide.shapes.add_picture(p, 0, 0, width=px_to_emu(w), height=px_to_emu(h))
+
+    out = tempfile.NamedTemporaryFile(suffix=".pptx", delete=False)
+    out.close()
+    prs.save(out.name)
+    return out.name
+
+def imagelist_to_zip(img_paths: List[str]) -> str:
+    out = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
+    out.close()
+    with zipfile.ZipFile(out.name, "w", zipfile.ZIP_DEFLATED) as zf:
+        for i, p in enumerate(img_paths):
+            zf.write(p, f"layer_{i+1}.png")
+    return out.name
+
+def write_images_to_node_dir(node_dir: str, images: List[Image.Image]) -> List[str]:
+    ensure_dir(node_dir)
+    paths = []
+    for i, im in enumerate(images):
+        p = os.path.join(node_dir, f"layer_{i+1:02d}.png")
+        im.save(p)
+        paths.append(p)
+    return paths
+
+def layer_choices_from_paths(paths: List[str]) -> List[Tuple[str, int]]:
+    # return list of (label, index)
+    out = []
+    for i in range(len(paths)):
+        out.append((f"Layer {i+1}", i))
+    return out
+
+def pretty_history_list(nodes: Dict[str, Any], selected_id: Optional[str]) -> str:
+    # Simple readable text block
+    # Show newest first
+    items = sorted(nodes.values(), key=lambda n: n.get("created_at", 0), reverse=True)
+    lines = []
+    for n in items[:50]:
+        mark = "👉 " if n["id"] == selected_id else "   "
+        parent = n.get("parent_id")
+        parent_txt = f" (parent: {parent})" if parent else ""
+        lines.append(f"{mark}{n['id']} — {n['name']} [{len(n['images'])} imgs]{parent_txt}")
+    if not lines:
+        return "_No history yet. Click **Decompose** to create the first node._"
+    return "```\n" + "\n".join(lines) + "\n```"
+
+
+# -------------------------
+# GPU duration control (ZeroGPU)
+# -------------------------
+def get_duration(
+    *args,
+    gpu_duration=1000,
+    **kwargs
+):
+    return clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
 
 
-def history_choices(history):
-    # Pretty dropdown labels with indentation + id
-    choices = []
-    for n in history:
-        depth = n.get("depth", 0)
-        indent = "  " * depth + ("↳ " if depth > 0 else "")
-        choices.append((f"{indent}{n['title']}  [{n['id']}]", n["id"]))
-    return choices
+# -------------------------
+# Core GPU runs
+# -------------------------
+@spaces.GPU(duration=get_duration)
+def run_decompose_gpu(
+    input_image,
+    seed=0,
+    prompt=None,
+    neg_prompt=" ",
+    true_guidance_scale=4.0,
+    num_inference_steps=50,
+    layers=7,
+    cfg_norm=True,
+    use_en_prompt=True,
+    resolution=1024,
+    gpu_duration=1000,
+):
+    pipe = get_pipeline()
 
+    # Normalize
+    resolution = int(resolution) if str(resolution).isdigit() else 1024
+    resolution = 1024 if resolution == 1024 else 640
+    prompt = sanitize_prompt(prompt)
+    neg_prompt = str(neg_prompt) if neg_prompt is not None else " "
 
-def render_breadcrumb(path_list):
-    if not path_list:
-        return "—"
-    return " → ".join(path_list)
+    # Image
+    if isinstance(input_image, list):
+        input_image = input_image[0]
+    if isinstance(input_image, str):
+        img = Image.open(input_image).convert("RGB").convert("RGBA")
+    elif isinstance(input_image, Image.Image):
+        img = input_image.convert("RGB").convert("RGBA")
+    elif isinstance(input_image, np.ndarray):
+        img = Image.fromarray(input_image).convert("RGB").convert("RGBA")
+    else:
+        raise ValueError(f"Unsupported input_image type: {type(input_image)}")
 
+    # Generator
+    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
+    generator = torch.Generator(device=gen_device).manual_seed(int(seed))
+
+    inputs = dict(
+        image=img,
+        generator=generator,
+        true_cfg_scale=float(true_guidance_scale),
+        prompt=prompt,
+        negative_prompt=neg_prompt,
+        num_inference_steps=int(num_inference_steps),
+        num_images_per_prompt=1,
+        layers=int(layers),
+        resolution=int(resolution),
+        cfg_normalize=bool(cfg_norm),
+        use_en_prompt=bool(use_en_prompt),
+    )
 
-# ----------------------------
-# ZeroGPU-friendly pipeline (lazy init)
-# ----------------------------
-PIPELINE = None
-PIPELINE_LOCK = threading.Lock()
+    print("DECOMPOSE INPUTS:", {k: (v if k != "image" else "<PIL>") for k, v in inputs.items()})
 
-TORCH_DTYPE = torch.float16  # important for RAM/VRAM
-MODEL_ID = "Qwen/Qwen-Image-Layered"
+    with torch.inference_mode():
+        out = pipe(**inputs)
+        imgs = out.images[0]  # list[PIL]
 
+    return imgs
 
-def get_pipeline():
-    """
-    Load ONLY inside GPU functions.
-    This avoids cold-start CPU-only load that can blow 30GB.
-    """
-    global PIPELINE
-    if PIPELINE is not None:
-        return PIPELINE
 
-    with PIPELINE_LOCK:
-        if PIPELINE is not None:
-            return PIPELINE
+@spaces.GPU(duration=get_duration)
+def run_refine_gpu(
+    selected_layer_path: str,
+    seed=0,
+    prompt=None,
+    neg_prompt=" ",
+    true_guidance_scale=4.0,
+    num_inference_steps=50,
+    sub_layers=3,
+    cfg_norm=True,
+    use_en_prompt=True,
+    resolution=1024,
+    gpu_duration=1000,
+):
+    pipe = get_pipeline()
 
-        pipe = QwenImageLayeredPipeline.from_pretrained(
-            MODEL_ID,
-            torch_dtype=TORCH_DTYPE,
-            low_cpu_mem_usage=True,
-        )
+    resolution = int(resolution) if str(resolution).isdigit() else 1024
+    resolution = 1024 if resolution == 1024 else 640
+    prompt = sanitize_prompt(prompt)
+    neg_prompt = str(neg_prompt) if neg_prompt is not None else " "
 
-        # memory helpers (guarded)
-        if hasattr(pipe, "enable_attention_slicing"):
-            pipe.enable_attention_slicing()
-        if hasattr(pipe, "enable_vae_slicing"):
-            pipe.enable_vae_slicing()
-        elif hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_slicing"):
-            pipe.vae.enable_slicing()
+    img = Image.open(selected_layer_path).convert("RGB").convert("RGBA")
 
-        PIPELINE = pipe
-        return PIPELINE
+    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
+    generator = torch.Generator(device=gen_device).manual_seed(int(seed))
+
+    inputs = dict(
+        image=img,
+        generator=generator,
+        true_cfg_scale=float(true_guidance_scale),
+        prompt=prompt,
+        negative_prompt=neg_prompt,
+        num_inference_steps=int(num_inference_steps),
+        num_images_per_prompt=1,
+        layers=int(sub_layers),
+        resolution=int(resolution),
+        cfg_normalize=bool(cfg_norm),
+        use_en_prompt=bool(use_en_prompt),
+    )
 
+    print("REFINE INPUTS:", {k: (v if k != "image" else "<PIL>") for k, v in inputs.items()})
 
-def ensure_device_strategy(pipe):
-    # Prefer CPU offload on CUDA to keep peak VRAM lower
-    if torch.cuda.is_available() and hasattr(pipe, "enable_model_cpu_offload"):
-        pipe.enable_model_cpu_offload()
-    elif torch.cuda.is_available():
-        try:
-            pipe.to("cuda")
-        except Exception:
-            pass
+    with torch.inference_mode():
+        out = pipe(**inputs)
+        imgs = out.images[0]
 
+    return imgs
 
-# ----------------------------
-# Dynamic GPU duration (ZeroGPU)
-# ----------------------------
-def get_duration(*args, **kwargs):
-    return clamp_int(kwargs.get("gpu_duration", 1000), default=1000, lo=20, hi=1500)
 
+# -------------------------
+# App state structure
+# -------------------------
+def new_state() -> Dict[str, Any]:
+    return {
+        "nodes": {},         # node_id -> node dict
+        "current_id": None,  # selected node id
+        "root_dir": SESSION_DIR,
+    }
 
-# ----------------------------
-# Node creation utilities
-# ----------------------------
-def add_node(history, parent_id, title, layer_paths, pptx_path, zip_path, meta: dict):
-    node_id = short_id()
-    node = {
+def add_node(
+    state: Dict[str, Any],
+    name: str,
+    image_paths: List[str],
+    parent_id: Optional[str] = None,
+) -> Dict[str, Any]:
+    node_id = _now_id()
+    state["nodes"][node_id] = {
         "id": node_id,
-        "parent_id": parent_id,           # None for root
-        "title": title,                   # shown in history
-        "layer_paths": layer_paths,       # list[str]
-        "pptx_path": pptx_path,           # str
-        "zip_path": zip_path,             # str
-        "n_layers": len(layer_paths),
-        "created_at": now_str(),
-        "meta": meta or {},
-        "depth": 0,
-        "path": [],
+        "name": name,
+        "parent_id": parent_id,
+        "images": image_paths,
+        "created_at": time.time(),
     }
-    history = list(history) if history else []
-    history.append(node)
-
-    # update depth/path for all nodes (simple, history small)
-    for n in history:
-        d, p = compute_depth_and_path(history, n["id"])
-        n["depth"] = d
-        n["path"] = p
-
-    return history, node_id
-
+    state["current_id"] = node_id
+    return state
+
+def get_current_node(state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    cid = state.get("current_id")
+    if not cid:
+        return None
+    return state["nodes"].get(cid)
+
+def get_node(state: Dict[str, Any], node_id: str) -> Optional[Dict[str, Any]]:
+    return state["nodes"].get(node_id)
+
+def history_dropdown_choices(state: Dict[str, Any]) -> List[Tuple[str, str]]:
+    # label, value=node_id
+    items = sorted(state["nodes"].values(), key=lambda n: n.get("created_at", 0), reverse=True)
+    return [(f"{n['name']} — {n['id']}", n["id"]) for n in items]
+
+def update_layer_ui_from_node(node: Dict[str, Any]):
+    paths = node["images"]
+    dd_choices = [label for label, _ in layer_choices_from_paths(paths)]
+    # Dropdown values are labels, but we need mapping label->index; we’ll use index stored separately via State.
+    return paths, dd_choices
+
+
+# -------------------------
+# Gradio callbacks (non-GPU)
+# -------------------------
+def on_decompose_click(
+    state: Dict[str, Any],
+    input_image,
+    seed,
+    randomize_seed,
+    prompt,
+    neg_prompt,
+    true_guidance_scale,
+    num_inference_steps,
+    layers,
+    cfg_norm,
+    use_en_prompt,
+    resolution,
+    gpu_duration,
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-def node_to_ui(history, node_id):
-    """
-    Convert node -> UI outputs (gallery/strip, exports, dropdown choices, preview, breadcrumb).
-    """
-    node = find_node(history, node_id)
-    if not node:
-        empty = []
-        return (
-            empty, empty,
-            None, None,
-            gr.update(choices=[], value=None),
-            0,
-            None,
-            f"**Node path:** —",
-        )
+    # GPU run
+    imgs = run_decompose_gpu(
+        input_image=input_image,
+        seed=int(seed),
+        prompt=prompt,
+        neg_prompt=neg_prompt,
+        true_guidance_scale=true_guidance_scale,
+        num_inference_steps=num_inference_steps,
+        layers=layers,
+        cfg_norm=cfg_norm,
+        use_en_prompt=use_en_prompt,
+        resolution=resolution,
+        gpu_duration=gpu_duration,
+    )
 
-    paths = node["layer_paths"]
-    images = [load_rgba(p) for p in paths]  # small N <= 10
-    strip = images
+    # Save node
+    node_dir = os.path.join(state["root_dir"], "nodes", f"node_{_now_id()}_decompose")
+    ensure_dir(node_dir)
+    paths = write_images_to_node_dir(node_dir, imgs)
 
-    labels = labels_for_layers(len(paths))
-    dd = gr.update(choices=labels, value=(labels[0] if labels else None))
+    state = add_node(
+        state=state,
+        name=f"Decompose ({len(paths)} layers)",
+        image_paths=paths,
+        parent_id=None,
+    )
 
-    selected_idx0 = 0
-    preview = load_rgba(paths[0]) if paths else None
+    node = get_current_node(state)
+    # Prepare UI
+    layer_paths = node["images"]
+    dd_labels = [f"Layer {i+1}" for i in range(len(layer_paths))]
 
-    breadcrumb = f"**Node path:** {render_breadcrumb(node.get('path', []))}"
+    # Default selection = first layer
+    selected_idx = 0
+    selected_label = dd_labels[selected_idx]
+    selected_path = layer_paths[selected_idx]
 
     return (
-        images, strip,
-        node["pptx_path"], node["zip_path"],
-        dd,
-        selected_idx0,
-        preview,
-        breadcrumb,
+        state,
+        gr.update(value=state["current_id"], choices=history_dropdown_choices(state)),  # history dropdown
+        gr.update(value=pretty_history_list(state["nodes"], state["current_id"])),     # history text
+        layer_paths,                                                                   # main gallery
+        gr.update(choices=dd_labels, value=selected_label),                            # refine dropdown
+        layer_paths,                                                                   # refine thumb gallery
+        gr.update(value=selected_path),                                                # refine preview
+        gr.update(open=False),                                                         # refined accordion closed (fresh start)
+        [],                                                                            # refined gallery cleared
     )
 
+def _label_to_index(label: str) -> int:
+    try:
+        # "Layer 3" -> 2
+        n = int(label.strip().split()[-1])
+        return max(0, n - 1)
+    except Exception:
+        return 0
 
-# ----------------------------
-# Selection handlers
-# ----------------------------
-def on_layer_dropdown_change(layer_label, current_layer_paths):
-    if not current_layer_paths:
-        return 0, None
-    idx0 = parse_layer_label(layer_label, 0)
-    idx0 = max(0, min(idx0, len(current_layer_paths) - 1))
-    return idx0, load_rgba(current_layer_paths[idx0])
-
-
-def on_gallery_select(current_layer_paths, evt: gr.SelectData):
-    if not current_layer_paths:
-        return 0, None, gr.update()
-    idx = evt.index
-    if isinstance(idx, (tuple, list)):
-        idx0 = int(idx[-1])
-    else:
-        idx0 = int(idx)
-    idx0 = max(0, min(idx0, len(current_layer_paths) - 1))
-    label = f"Layer {idx0 + 1}"
-    return idx0, load_rgba(current_layer_paths[idx0]), gr.update(value=label)
-
-
-def on_history_select(history, node_id):
-    if not node_id:
+def on_history_change(state: Dict[str, Any], node_id: str):
+    node = get_node(state, node_id) if node_id else None
+    if not node:
         return (
+            state,
+            gr.update(value=pretty_history_list(state["nodes"], state.get("current_id"))),
             gr.update(choices=[], value=None),
-            [],  # current_layer_paths_state
-            [], [], None, None,
-            gr.update(choices=[], value=None),
-            0,
-            None,
-            "**Node path:** —",
+            [],
+            gr.update(value=None),
+            [],
+            gr.update(open=False),
+            [],
         )
 
-    node = find_node(history, node_id)
-    if not node:
-        raise gr.Error("History node not found (state mismatch).")
-
-    # Build UI for selected node
-    images, strip, pptx_path, zip_path, layer_dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
+    state["current_id"] = node_id
+    paths = node["images"]
+    dd_labels = [f"Layer {i+1}" for i in range(len(paths))]
+    selected_idx = 0
+    selected_label = dd_labels[selected_idx] if dd_labels else None
+    selected_path = paths[selected_idx] if paths else None
 
     return (
-        layer_dd,
-        node["layer_paths"],       # current_layer_paths_state
-        images,
-        strip,
-        pptx_path,
-        zip_path,
-        layer_dd,
-        sel_idx0,
-        preview,
-        breadcrumb,
+        state,
+        gr.update(value=pretty_history_list(state["nodes"], state["current_id"])),
+        gr.update(choices=dd_labels, value=selected_label),
+        paths,
+        gr.update(value=selected_path),
+        paths,
+        gr.update(open=False),
+        [],
     )
 
+def on_refine_thumb_select(state: Dict[str, Any], evt: gr.SelectData):
+    node = get_current_node(state)
+    if not node:
+        return gr.update(), gr.update()
 
-# ----------------------------
-# Inference core
-# ----------------------------
-def run_pipeline_decompose(pipe, image_rgba, seed, randomize_seed, prompt, neg_prompt,
-                          true_guidance_scale, num_inference_steps, layers, cfg_norm,
-                          use_en_prompt, resolution):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
-    gen = torch.Generator(device=gen_device).manual_seed(int(seed))
-
-    inputs = {
-        "image": image_rgba,
-        "generator": gen,
-        "true_cfg_scale": float(true_guidance_scale),
-        "prompt": prompt if prompt is not None else "",
-        "negative_prompt": neg_prompt,
-        "num_inference_steps": int(num_inference_steps),
-        "num_images_per_prompt": 1,
-        "layers": int(layers),
-        "resolution": int(resolution),
-        "cfg_normalize": bool(cfg_norm),
-        "use_en_prompt": bool(use_en_prompt),
-    }
-
-    with torch.inference_mode():
-        out = pipe(**inputs)
-        return out.images[0]  # list[PIL]
-
+    idx = int(evt.index) if evt and evt.index is not None else 0
+    idx = max(0, min(idx, len(node["images"]) - 1))
+    label = f"Layer {idx+1}"
+    path = node["images"][idx]
+    return gr.update(value=label), gr.update(value=path)
 
-# ----------------------------
-# GPU functions
-# ----------------------------
-@spaces.GPU(duration=get_duration)
-def do_decompose(
-    input_image,
-    seed=0,
-    randomize_seed=True,
-    prompt="",
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    layers=7,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=1024,
-    gpu_duration=1000,
-    history=None,
+def on_refine_dropdown_change(state: Dict[str, Any], label: str):
+    node = get_current_node(state)
+    if not node:
+        return gr.update(value=None)
+    idx = _label_to_index(label)
+    idx = max(0, min(idx, len(node["images"]) - 1))
+    return gr.update(value=node["images"][idx])
+
+def on_refine_click(
+    state: Dict[str, Any],
+    refine_layer_label: str,
+    seed,
+    randomize_seed,
+    prompt,
+    neg_prompt,
+    true_guidance_scale,
+    num_inference_steps,
+    sub_layers,
+    cfg_norm,
+    use_en_prompt,
+    resolution,
+    gpu_duration,
 ):
-    if isinstance(input_image, list):
-        input_image = input_image[0]
-
-    if isinstance(input_image, str):
-        pil_image = Image.open(input_image).convert("RGBA")
-    elif isinstance(input_image, Image.Image):
-        pil_image = input_image.convert("RGBA")
-    elif isinstance(input_image, np.ndarray):
-        pil_image = Image.fromarray(input_image).convert("RGBA")
-    else:
-        raise ValueError(f"Unsupported input_image type: {type(input_image)}")
-
-    resolution = norm_resolution(resolution)
-    layers = clamp_int(layers, default=7, lo=2, hi=10)
+    node = get_current_node(state)
+    if not node:
+        raise gr.Error("No active node. Run Decompose first.")
 
-    pipe = get_pipeline()
-    ensure_device_strategy(pipe)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-    imgs = run_pipeline_decompose(
-        pipe, pil_image, seed, randomize_seed, prompt, neg_prompt,
-        true_guidance_scale, num_inference_steps, layers, cfg_norm, use_en_prompt, resolution
+    idx = _label_to_index(refine_layer_label)
+    idx = max(0, min(idx, len(node["images"]) - 1))
+    selected_path = node["images"][idx]
+
+    # GPU run refine
+    imgs = run_refine_gpu(
+        selected_layer_path=selected_path,
+        seed=int(seed),
+        prompt=prompt,
+        neg_prompt=neg_prompt,
+        true_guidance_scale=true_guidance_scale,
+        num_inference_steps=num_inference_steps,
+        sub_layers=sub_layers,
+        cfg_norm=cfg_norm,
+        use_en_prompt=use_en_prompt,
+        resolution=resolution,
+        gpu_duration=gpu_duration,
     )
 
-    # Save images to temp
-    layer_paths = []
-    gallery = []
-    for img in imgs:
-        gallery.append(img)
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-        img.save(tmp.name)
-        layer_paths.append(tmp.name)
-
-    pptx_path = imagelist_to_pptx(layer_paths)
-    zip_path = make_zip(layer_paths, prefix="layer")
-
-    meta = {
-        "kind": "decompose",
-        "resolution": resolution,
-        "layers": layers,
-        "steps": int(num_inference_steps),
-    }
-    title = f"Decompose ({len(layer_paths)} layers) @ {resolution}"
+    # Save new node (child of current)
+    node_dir = os.path.join(state["root_dir"], "nodes", f"node_{_now_id()}_refine")
+    ensure_dir(node_dir)
+    paths = write_images_to_node_dir(node_dir, imgs)
 
-    history = history or []
-    history, node_id = add_node(history, parent_id=None, title=title,
-                               layer_paths=layer_paths, pptx_path=pptx_path, zip_path=zip_path, meta=meta)
-
-    # Update history dropdown
-    hist_dd = gr.update(choices=history_choices(history), value=node_id)
+    state = add_node(
+        state=state,
+        name=f"Refine {refine_layer_label} → ({len(paths)} sub-layers)",
+        image_paths=paths,
+        parent_id=node["id"],
+    )
 
-    # Set current node UI
-    images, strip, pptx, zipp, layer_dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
+    # Now current node is the refined node (so "Refine again" is just clicking Refine again)
+    new_node = get_current_node(state)
+    new_paths = new_node["images"]
+    dd_labels = [f"Layer {i+1}" for i in range(len(new_paths))]
+    selected_idx2 = 0
+    selected_label2 = dd_labels[selected_idx2] if dd_labels else None
+    selected_path2 = new_paths[selected_idx2] if new_paths else None
 
     return (
-        history,
-        node_id,
-        hist_dd,
-        layer_paths,     # current_layer_paths_state
-        images,
-        strip,
-        pptx,
-        zipp,
-        layer_dd,
-        sel_idx0,
-        preview,
-        breadcrumb,
+        state,
+        gr.update(value=state["current_id"], choices=history_dropdown_choices(state)),
+        gr.update(value=pretty_history_list(state["nodes"], state["current_id"])),
+        new_paths,                                 # main gallery becomes the refined node
+        gr.update(choices=dd_labels, value=selected_label2),
+        new_paths,                                 # refine thumbnails now reflect refined node
+        gr.update(value=selected_path2),
+        gr.update(open=True),                      # auto-open refined accordion
+        new_paths,                                 # refined gallery shows refined node too
     )
 
+def on_export_click(state: Dict[str, Any], node_id: str):
+    node = get_node(state, node_id) if node_id else get_current_node(state)
+    if not node:
+        raise gr.Error("No node selected to export.")
+    paths = node["images"]
+    if not paths:
+        raise gr.Error("Selected node has no images.")
+    pptx = imagelist_to_pptx(paths)
+    zipp = imagelist_to_zip(paths)
+    return pptx, zipp
 
-@spaces.GPU(duration=get_duration)
-def do_refine(
-    history,
-    current_node_id,
-    current_layer_paths,
-    selected_layer_idx0,
-    sub_layers=3,
-    seed=0,
-    randomize_seed=True,
-    prompt="",
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=1024,
-    gpu_duration=1000,
-):
-    if not history or not current_node_id:
-        raise gr.Error("Сначала сделай Decompose (создай root-узел).")
-    if not current_layer_paths:
-        raise gr.Error("Нет слоёв в текущем узле (state).")
-
-    parent = find_node(history, current_node_id)
-    if not parent:
-        raise gr.Error("Текущий узел не найден в history.")
-
-    resolution = norm_resolution(resolution)
-    sub_layers = clamp_int(sub_layers, default=3, lo=2, hi=10)
-
-    idx0 = clamp_int(selected_layer_idx0, default=0, lo=0, hi=len(current_layer_paths) - 1)
-    selected_img = load_rgba(current_layer_paths[idx0])
 
-    pipe = get_pipeline()
-    ensure_device_strategy(pipe)
+# -------------------------
+# UI
+# -------------------------
+state = gr.State(new_state())
 
-    imgs = run_pipeline_decompose(
-        pipe, selected_img, seed, randomize_seed, prompt, neg_prompt,
-        true_guidance_scale, num_inference_steps, sub_layers, cfg_norm, use_en_prompt, resolution
+with gr.Blocks() as demo:
+    gr.HTML(
+        '<img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/layered/qwen-image-layered-logo.png" '
+        'alt="Qwen-Image-Layered Logo" width="520" style="display:block;margin:0 auto;">'
     )
-
-    # Save images to temp
-    layer_paths = []
-    gallery = []
-    for img in imgs:
-        gallery.append(img)
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-        img.save(tmp.name)
-        layer_paths.append(tmp.name)
-
-    pptx_path = imagelist_to_pptx(layer_paths)
-    zip_path = make_zip(layer_paths, prefix="refined")
-
-    meta = {
-        "kind": "refine",
-        "resolution": resolution,
-        "sub_layers": sub_layers,
-        "steps": int(num_inference_steps),
-        "refined_from": {"node_id": current_node_id, "layer_index": idx0},
-    }
-    title = f"Refine L{idx0+1} → {len(layer_paths)} sub @ {resolution}"
-
-    history, node_id = add_node(history, parent_id=current_node_id, title=title,
-                               layer_paths=layer_paths, pptx_path=pptx_path, zip_path=zip_path, meta=meta)
-
-    # Update history dropdown and set current node to the new child
-    hist_dd = gr.update(choices=history_choices(history), value=node_id)
-
-    images, strip, pptx, zipp, layer_dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
-
-    return (
-        history,
-        node_id,
-        hist_dd,
-        layer_paths,  # current_layer_paths_state
-        images,
-        strip,
-        pptx,
-        zipp,
-        layer_dd,
-        sel_idx0,
-        preview,
-        breadcrumb,
+    gr.Markdown(
+        "Decompose → choose a layer → Refine → build a recursive history tree. Export any node to ZIP/PPTX."
     )
 
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(label="Input Image", image_mode="RGBA")
 
-# ----------------------------
-# UI
-# ----------------------------
-ensure_dirname(LOG_DIR)
-examples = [f"assets/test_images/{i}.png" for i in range(1, 14)]
+            with gr.Accordion("Settings", open=False):
+                prompt = gr.Textbox(
+                    label="Prompt (Optional)",
+                    placeholder="Describe the image (overall content, including occluded parts).",
+                    value="",
+                    lines=2,
+                )
+                neg_prompt = gr.Textbox(
+                    label="Negative Prompt (Optional)",
+                    value=" ",
+                    lines=2,
+                )
 
-with gr.Blocks() as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.HTML(
-            '<img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/layered/qwen-image-layered-logo.png" '
-            'alt="Qwen-Image-Layered Logo" width="600" style="display: block; margin: 0 auto;">'
-        )
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-        gr.Markdown(
-            """
-Prompt описывает изображение в целом (включая частично закрытые элементы).
-Refine делает рекурсивную декомпозицию выбранного слоя текущего узла (узлы сохраняются в History).
-            """
-        )
+                true_guidance_scale = gr.Slider(
+                    label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0
+                )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps", minimum=1, maximum=100, step=1, value=50
+                )
 
-        # States
-        history_state = gr.State([])
-        current_node_id_state = gr.State(None)
-        current_layer_paths_state = gr.State([])
-        selected_layer_idx0_state = gr.State(0)
+                resolution = gr.Radio(
+                    label="Processing resolution",
+                    choices=[640, 1024],
+                    value=1024,  # default 1024
+                )
 
-        with gr.Row():
-            with gr.Column(scale=1):
-                input_image = gr.Image(label="Input Image", image_mode="RGBA")
+                cfg_norm = gr.Checkbox(label="Whether enable CFG normalization", value=True)
+                use_en_prompt = gr.Checkbox(
+                    label="Automatic caption language if no prompt (True=EN, False=ZH)",
+                    value=True,
+                )
 
-                with gr.Accordion("Settings", open=False):
-                    prompt = gr.Textbox(label="Prompt (Optional)", value="", lines=2)
-                    neg_prompt = gr.Textbox(label="Negative Prompt (Optional)", value=" ", lines=2)
+                gpu_duration = gr.Textbox(
+                    label="GPU duration override (seconds, 20..1500)",
+                    value="1000",
+                    lines=1,
+                    placeholder="e.g. 120 / 300 / 1000 / 1500",
+                )
 
-                    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():
+                layers = gr.Slider(label="Layers (Decompose)", minimum=2, maximum=10, step=1, value=7)
 
-                    true_guidance_scale = gr.Slider(label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0)
-                    num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=50)
+            decompose_btn = gr.Button("Decompose", variant="primary")
 
-                    layers = gr.Slider(label="Layers (Decompose)", minimum=2, maximum=10, step=1, value=7)
-                    sub_layers = gr.Slider(label="Sub-layers (Refine)", minimum=2, maximum=10, step=1, value=3)
+            gr.Markdown("### History")
+            history_dd = gr.Dropdown(
+                label="Select node",
+                choices=[],
+                value=None,
+                interactive=True,
+            )
+            history_text = gr.Markdown("_No history yet._")
 
-                    resolution = gr.Radio(label="Resolution", choices=[640, 1024], value=1024)
+            with gr.Row():
+                export_pptx = gr.File(label="Download PPTX")
+                export_zip = gr.File(label="Download ZIP")
+            export_btn = gr.Button("Export selected node")
 
-                    cfg_norm = gr.Checkbox(label="CFG normalize", value=True)
-                    use_en_prompt = gr.Checkbox(label="Auto caption language (EN=True / ZH=False)", value=True)
+        with gr.Column(scale=2):
+            gr.Markdown("### Layers (Current node)")
+            main_gallery = gr.Gallery(label="", columns=4, rows=2, height=420, show_label=False)
 
-                    gpu_duration = gr.Textbox(
-                        label="GPU duration override (seconds, 20..1500)",
-                        value="1000",
-                        lines=1,
+            with gr.Accordion("Refine", open=True):
+                gr.Markdown("Pick a layer like in Photoshop (thumbnails) or via dropdown, then refine it.")
+                with gr.Row():
+                    refine_layer_dd = gr.Dropdown(
+                        label="Refine layer",
+                        choices=[],
+                        value=None,
+                        interactive=True,
+                    )
+                    sub_layers = gr.Slider(
+                        label="Sub-layers (Refine)",
+                        minimum=2, maximum=10, step=1, value=3
                     )
 
-                decompose_btn = gr.Button("Decompose (new root node)", variant="primary")
-                refine_btn = gr.Button("Refine selected layer (create child node)", variant="secondary")
-
-                gr.Markdown("### History (nodes)")
-                history_dd = gr.Dropdown(label="Select node", choices=[], value=None, interactive=True)
+                # Photoshop-like 1-row thumbnail strip
+                refine_thumb_gallery = gr.Gallery(
+                    label="Layer thumbnails (click to select)",
+                    columns=8,
+                    rows=1,
+                    height=120,
+                )
 
-                breadcrumb_md = gr.Markdown("**Node path:** —")
+                refine_preview = gr.Image(label="Selected layer preview", image_mode="RGBA")
 
-                gr.Markdown("### Layer selection (current node)")
-                layer_dd = gr.Dropdown(label="Select layer", choices=[], value=None, interactive=True)
-                layer_preview = gr.Image(label="Selected layer preview", image_mode="RGBA", interactive=False)
+                refine_btn = gr.Button("Refine selected layer", variant="secondary")
 
-            with gr.Column(scale=2):
-                current_gallery = gr.Gallery(label="Current node layers (click to select)", columns=4, rows=1, format="png")
-                current_strip = gr.Gallery(label="Layer strip (1 row)", columns=8, rows=1, format="png", height=120)
+            refined_acc = gr.Accordion("Refined layers", open=False)
+            with refined_acc:
+                refined_gallery = gr.Gallery(label="", columns=4, rows=1, height=220, show_label=False)
 
-                with gr.Row():
-                    export_pptx = gr.File(label="Download PPTX (selected node)")
-                    export_zip = gr.File(label="Download ZIP (selected node)")
-
-        # Examples run -> Decompose
-        gr.Examples(
-            examples=examples,
-            inputs=[input_image],
-            outputs=[
-                history_state, current_node_id_state, history_dd,
-                current_layer_paths_state, current_gallery, current_strip,
-                export_pptx, export_zip,
-                layer_dd, selected_layer_idx0_state, layer_preview,
-                breadcrumb_md,
-            ],
-            fn=do_decompose,
-            examples_per_page=14,
-            cache_examples=False,
-            run_on_click=True,
-        )
+    # Examples (optional)
+    examples = [f"assets/test_images/{i}.png" for i in range(1, 14)]
+    gr.Examples(
+        examples=examples,
+        inputs=[input_image],
+        fn=None,
+        cache_examples=False,
+        label="Examples",
+    )
 
-        # Decompose button
-        decompose_btn.click(
-            fn=do_decompose,
-            inputs=[
-                input_image,
-                seed,
-                randomize_seed,
-                prompt,
-                neg_prompt,
-                true_guidance_scale,
-                num_inference_steps,
-                layers,
-                cfg_norm,
-                use_en_prompt,
-                resolution,
-                gpu_duration,
-                history_state,
-            ],
-            outputs=[
-                history_state,
-                current_node_id_state,
-                history_dd,
-                current_layer_paths_state,
-                current_gallery,
-                current_strip,
-                export_pptx,
-                export_zip,
-                layer_dd,
-                selected_layer_idx0_state,
-                layer_preview,
-                breadcrumb_md,
-            ],
-        )
+    # -------------------------
+    # Wiring
+    # -------------------------
+    decompose_btn.click(
+        fn=on_decompose_click,
+        inputs=[
+            state,
+            input_image,
+            seed,
+            randomize_seed,
+            prompt,
+            neg_prompt,
+            true_guidance_scale,
+            num_inference_steps,
+            layers,
+            cfg_norm,
+            use_en_prompt,
+            resolution,
+            gpu_duration,
+        ],
+        outputs=[
+            state,
+            history_dd,
+            history_text,
+            main_gallery,
+            refine_layer_dd,
+            refine_thumb_gallery,
+            refine_preview,
+            refined_acc,
+            refined_gallery,
+        ],
+    )
 
-        # Refine button
-        refine_btn.click(
-            fn=do_refine,
-            inputs=[
-                history_state,
-                current_node_id_state,
-                current_layer_paths_state,
-                selected_layer_idx0_state,
-                sub_layers,
-                seed,
-                randomize_seed,
-                prompt,
-                neg_prompt,
-                true_guidance_scale,
-                num_inference_steps,
-                cfg_norm,
-                use_en_prompt,
-                resolution,
-                gpu_duration,
-            ],
-            outputs=[
-                history_state,
-                current_node_id_state,
-                history_dd,
-                current_layer_paths_state,
-                current_gallery,
-                current_strip,
-                export_pptx,
-                export_zip,
-                layer_dd,
-                selected_layer_idx0_state,
-                layer_preview,
-                breadcrumb_md,
-            ],
-        )
+    history_dd.change(
+        fn=on_history_change,
+        inputs=[state, history_dd],
+        outputs=[
+            state,
+            history_text,
+            refine_layer_dd,
+            main_gallery,
+            refine_preview,
+            refine_thumb_gallery,
+            refined_acc,
+            refined_gallery,
+        ],
+    )
 
-        # History selection -> load any node
-        def _history_change(history, node_id):
-            # returns:
-            # layer_dd_update,
-            # current_layer_paths_state,
-            # current_gallery,
-            # current_strip,
-            # export_pptx,
-            # export_zip,
-            # layer_dd,
-            # selected_layer_idx0_state,
-            # layer_preview,
-            # breadcrumb
-            node = find_node(history, node_id)
-            if not node:
-                return (
-                    gr.update(choices=[], value=None),
-                    [],
-                    [], [],
-                    None, None,
-                    gr.update(choices=[], value=None),
-                    0,
-                    None,
-                    "**Node path:** —",
-                )
-            images, strip, pptx, zipp, dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
-            return (
-                dd,
-                node["layer_paths"],
-                images,
-                strip,
-                pptx,
-                zipp,
-                dd,
-                sel_idx0,
-                preview,
-                breadcrumb,
-            )
+    refine_thumb_gallery.select(
+        fn=on_refine_thumb_select,
+        inputs=[state],
+        outputs=[refine_layer_dd, refine_preview],
+    )
 
-        history_dd.change(
-            fn=_history_change,
-            inputs=[history_state, history_dd],
-            outputs=[
-                layer_dd,
-                current_layer_paths_state,
-                current_gallery,
-                current_strip,
-                export_pptx,
-                export_zip,
-                layer_dd,
-                selected_layer_idx0_state,
-                layer_preview,
-                breadcrumb_md,
-            ],
-        )
+    refine_layer_dd.change(
+        fn=on_refine_dropdown_change,
+        inputs=[state, refine_layer_dd],
+        outputs=[refine_preview],
+    )
 
-        # Layer dropdown -> preview
-        layer_dd.change(
-            fn=on_layer_dropdown_change,
-            inputs=[layer_dd, current_layer_paths_state],
-            outputs=[selected_layer_idx0_state, layer_preview],
-        )
+    refine_btn.click(
+        fn=on_refine_click,
+        inputs=[
+            state,
+            refine_layer_dd,
+            seed,
+            randomize_seed,
+            prompt,
+            neg_prompt,
+            true_guidance_scale,
+            num_inference_steps,
+            sub_layers,
+            cfg_norm,
+            use_en_prompt,
+            resolution,
+            gpu_duration,
+        ],
+        outputs=[
+            state,
+            history_dd,
+            history_text,
+            main_gallery,
+            refine_layer_dd,
+            refine_thumb_gallery,
+            refine_preview,
+            refined_acc,
+            refined_gallery,
+        ],
+    )
 
-        # Click on gallery/strip -> select layer
-        current_gallery.select(
-            fn=on_gallery_select,
-            inputs=[current_layer_paths_state],
-            outputs=[selected_layer_idx0_state, layer_preview, layer_dd],
-        )
-        current_strip.select(
-            fn=on_gallery_select,
-            inputs=[current_layer_paths_state],
-            outputs=[selected_layer_idx0_state, layer_preview, layer_dd],
-        )
+    export_btn.click(
+        fn=on_export_click,
+        inputs=[state, history_dd],
+        outputs=[export_pptx, export_zip],
+    )
 
 if __name__ == "__main__":
-    demo.queue()
-    try:
-        demo.launch(ssr_mode=False)
-    except TypeError:
-        demo.launch()
+    # ssr_mode=False helps avoid those asyncio "Invalid file descriptor" messages in some envs
+    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)