Update app.py

app.py

@@ -1,96 +1,63 @@
 import os
 import uuid
+import numpy as np
 import random
 import tempfile
 import zipfile
-import numpy as np
+import threading
+from datetime import datetime
 
 import spaces
 import torch
 import gradio as gr
-
 from PIL import Image
-from pptx import Presentation
+
 from diffusers import QwenImageLayeredPipeline
+from pptx import Presentation
+
 
 LOG_DIR = "/tmp/local"
 MAX_SEED = np.iinfo(np.int32).max
 
-# Reduce allocator fragmentation (new name; old PYTORCH_CUDA_ALLOC_CONF is deprecated)
-os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
-
 # Optional HF login (Spaces secret env var "hf")
 from huggingface_hub import login
-login(token=os.environ.get("hf"))
-
-# ----------------------------
-# Device / dtype (memory-safe)
-# ----------------------------
-has_cuda = torch.cuda.is_available()
-device = "cuda" if has_cuda else ("mps" if torch.backends.mps.is_available() else "cpu")
+_HF_TOKEN = os.environ.get("hf")
+if _HF_TOKEN:
+    login(token=_HF_TOKEN)
 
-# fp16 is typically best for VRAM; CPU uses fp32
-torch_dtype = torch.float16 if device in ("cuda", "mps") else torch.float32
 
 # ----------------------------
-# Load pipeline (avoid CPU RAM spikes)
+# Helpers
 # ----------------------------
-pipeline = QwenImageLayeredPipeline.from_pretrained(
-    "Qwen/Qwen-Image-Layered",
-    torch_dtype=torch_dtype,
-    low_cpu_mem_usage=True,
-)
-
-# Memory helpers (guarded)
-if hasattr(pipeline, "enable_attention_slicing"):
-    pipeline.enable_attention_slicing()
-
-# This pipeline may NOT expose enable_vae_slicing(), so guard both ways
-if hasattr(pipeline, "enable_vae_slicing"):
-    pipeline.enable_vae_slicing()
-elif hasattr(pipeline, "vae") and hasattr(pipeline.vae, "enable_slicing"):
-    pipeline.vae.enable_slicing()
-
-if device == "cuda":
-    # Best for Spaces: keep CPU RAM lower and avoid huge peak VRAM at startup
-    # (requires accelerate, usually present in Spaces)
-    try:
-        pipeline.enable_model_cpu_offload()
-    except Exception:
-        pipeline.to("cuda")
-elif device == "mps":
-    pipeline.to("mps")
-else:
-    pipeline.to("cpu")
-
-
 def ensure_dirname(path: str):
     if path and not os.path.exists(path):
         os.makedirs(path, exist_ok=True)
 
 
+def px_to_emu(px, dpi=96):
+    inch = px / dpi
+    return int(inch * 914400)
+
+
 def imagelist_to_pptx(img_files):
     with Image.open(img_files[0]) as img:
-        img_width_px, img_height_px = img.size
-
-    def px_to_emu(px, dpi=96):
-        inch = px / dpi
-        return int(inch * 914400)
+        w, h = img.size
 
     prs = Presentation()
-    prs.slide_width = px_to_emu(img_width_px)
-    prs.slide_height = px_to_emu(img_height_px)
+    prs.slide_width = px_to_emu(w)
+    prs.slide_height = px_to_emu(h)
 
     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(img_width_px),
-            height=px_to_emu(img_height_px),
+            width=px_to_emu(w),
+            height=px_to_emu(h),
         )
 
     with tempfile.NamedTemporaryFile(suffix=".pptx", delete=False) as tmp:
@@ -98,7 +65,15 @@ def imagelist_to_pptx(img_files):
         return tmp.name
 
 
-def _clamp_int(x, default: int, lo: int, hi: int) -> int:
+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
+
+
+def clamp_int(x, default: int, lo: int, hi: int) -> int:
     try:
         v = int(x)
     except Exception:
@@ -106,139 +81,389 @@ def _clamp_int(x, default: int, lo: int, hi: int) -> int:
     return max(lo, min(hi, v))
 
 
-def _safe_open_rgba(img_like):
-    if isinstance(img_like, list):
-        img_like = img_like[0]
-    if isinstance(img_like, str):
-        return Image.open(img_like).convert("RGB").convert("RGBA")
-    if isinstance(img_like, Image.Image):
-        return img_like.convert("RGB").convert("RGBA")
-    if isinstance(img_like, np.ndarray):
-        return Image.fromarray(img_like).convert("RGB").convert("RGBA")
-    raise ValueError(f"Unsupported input_image type: {type(img_like)}")
+def norm_resolution(x):
+    v = clamp_int(x, default=1024, lo=640, hi=1024)
+    return v if v in (640, 1024) else 1024
 
 
-def _update_refine_index_ui(n_layers: int, current_idx: int | None = None):
-    n_layers = max(1, int(n_layers))
-    if current_idx is None:
-        current_idx = 1
-    current_idx = max(1, min(int(current_idx), n_layers))
-    return gr.update(minimum=1, maximum=n_layers, value=current_idx)
+def load_rgba(path: str) -> Image.Image:
+    return Image.open(path).convert("RGBA")
 
 
-# Dynamic duration callable: must accept same args as decompose() and refine()
-def get_duration(
-    input_image,
-    seed=0,
-    randomize_seed=True,
-    prompt="",
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    layer=7,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=1024,
-    gpu_duration="1000",
-    refine_layer_index=1,
-    refine_sub_layers=3,
-):
-    return _clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
+def labels_for_layers(n: int):
+    return [f"Layer {i}" for i in range(1, n + 1)]
 
 
-@spaces.GPU(duration=get_duration)
-def decompose(
-    input_image,
-    seed=0,
-    randomize_seed=True,
-    prompt="",
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    layer=7,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=1024,
-    gpu_duration="1000",
-    refine_layer_index=1,   # passed in (so we can "clamp" it красиво)
-    refine_sub_layers=3,    # unused here, but kept for duration signature parity
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+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 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
+
+
+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
+
+
+def render_breadcrumb(path_list):
+    if not path_list:
+        return "—"
+    return " → ".join(path_list)
+
+
+# ----------------------------
+# ZeroGPU-friendly pipeline (lazy init)
+# ----------------------------
+PIPELINE = None
+PIPELINE_LOCK = threading.Lock()
+
+TORCH_DTYPE = torch.float16  # important for RAM/VRAM
+MODEL_ID = "Qwen/Qwen-Image-Layered"
+
+
+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
+
+        pipe = QwenImageLayeredPipeline.from_pretrained(
+            MODEL_ID,
+            torch_dtype=TORCH_DTYPE,
+            low_cpu_mem_usage=True,
+        )
+
+        # 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()
+
+        PIPELINE = pipe
+        return PIPELINE
+
 
-    resolution = _clamp_int(resolution, default=1024, lo=640, hi=1024)
-    if resolution not in (640, 1024):
-        resolution = 1024
+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
+
+
+# ----------------------------
+# Dynamic GPU duration (ZeroGPU)
+# ----------------------------
+def get_duration(*args, **kwargs):
+    return clamp_int(kwargs.get("gpu_duration", 1000), default=1000, lo=20, hi=1500)
+
+
+# ----------------------------
+# Node creation utilities
+# ----------------------------
+def add_node(history, parent_id, title, layer_paths, pptx_path, zip_path, meta: dict):
+    node_id = short_id()
+    node = {
+        "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": [],
+    }
+    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
+
+
+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:** —",
+        )
+
+    paths = node["layer_paths"]
+    images = [load_rgba(p) for p in paths]  # small N <= 10
+    strip = images
+
+    labels = labels_for_layers(len(paths))
+    dd = gr.update(choices=labels, value=(labels[0] if labels else None))
+
+    selected_idx0 = 0
+    preview = load_rgba(paths[0]) if paths else None
+
+    breadcrumb = f"**Node path:** {render_breadcrumb(node.get('path', []))}"
+
+    return (
+        images, strip,
+        node["pptx_path"], node["zip_path"],
+        dd,
+        selected_idx0,
+        preview,
+        breadcrumb,
+    )
+
+
+# ----------------------------
+# 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:
+        return (
+            gr.update(choices=[], value=None),
+            [],  # current_layer_paths_state
+            [], [], None, None,
+            gr.update(choices=[], value=None),
+            0,
+            None,
+            "**Node path:** —",
+        )
 
-    pil_image = _safe_open_rgba(input_image)
+    node = find_node(history, node_id)
+    if not node:
+        raise gr.Error("History node not found (state mismatch).")
 
-    # Generator on CPU works well with CPU offload too
-    gen = torch.Generator(device="cpu").manual_seed(seed)
+    # Build UI for selected node
+    images, strip, pptx_path, zip_path, layer_dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
+
+    return (
+        layer_dd,
+        node["layer_paths"],       # current_layer_paths_state
+        images,
+        strip,
+        pptx_path,
+        zip_path,
+        layer_dd,
+        sel_idx0,
+        preview,
+        breadcrumb,
+    )
+
+
+# ----------------------------
+# 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": pil_image,
+        "image": image_rgba,
         "generator": gen,
         "true_cfg_scale": float(true_guidance_scale),
-        "prompt": prompt if prompt else None,
+        "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(layer),
+        "layers": int(layers),
         "resolution": int(resolution),
         "cfg_normalize": bool(cfg_norm),
         "use_en_prompt": bool(use_en_prompt),
     }
 
-    print("DECOMPOSE INPUTS:", {k: v for k, v in inputs.items() if k != "image"})
-    print("REQUESTED GPU DURATION:", gpu_duration)
-
     with torch.inference_mode():
-        out = pipeline(**inputs)
-        output_images = out.images[0]  # list[PIL.Image]
+        out = pipe(**inputs)
+        return out.images[0]  # list[PIL]
 
-    # Save layers for exports + for refine stage
-    layer_paths = []
-    gallery_out = []
 
-    for img in output_images:
-        gallery_out.append(img)
+# ----------------------------
+# 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,
+):
+    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)
+
+    pipe = get_pipeline()
+    ensure_device_strategy(pipe)
+
+    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
+    )
+
+    # 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")
 
-    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmpzip:
-        with zipfile.ZipFile(tmpzip.name, "w", zipfile.ZIP_DEFLATED) as zipf:
-            for i, p in enumerate(layer_paths):
-                zipf.write(p, f"layer_{i+1}.png")
-        zip_path = tmpzip.name
+    meta = {
+        "kind": "decompose",
+        "resolution": resolution,
+        "layers": layers,
+        "steps": int(num_inference_steps),
+    }
+    title = f"Decompose ({len(layer_paths)} layers) @ {resolution}"
+
+    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)
 
-    # Reset refined outputs on new decompose
-    refined_gallery = []
-    refined_pptx = None
-    refined_zip = None
+    # Update history dropdown
+    hist_dd = gr.update(choices=history_choices(history), value=node_id)
 
-    # "совсем красиво": clamp current refine index to new [1..N]
-    refine_index_update = _update_refine_index_ui(len(layer_paths), refine_layer_index)
+    # Set current node UI
+    images, strip, pptx, zipp, layer_dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
 
     return (
-        gallery_out,
-        pptx_path,
-        zip_path,
-        layer_paths,            # gr.State
-        refined_gallery,
-        refined_pptx,
-        refined_zip,
-        refine_index_update,    # update refine slider bounds/value
+        history,
+        node_id,
+        hist_dd,
+        layer_paths,     # current_layer_paths_state
+        images,
+        strip,
+        pptx,
+        zipp,
+        layer_dd,
+        sel_idx0,
+        preview,
+        breadcrumb,
     )
 
 
 @spaces.GPU(duration=get_duration)
-def refine_selected_layer(
-    layer_paths,
-    refine_layer_index=1,
-    refine_sub_layers=3,
+def do_refine(
+    history,
+    current_node_id,
+    current_layer_paths,
+    selected_layer_idx0,
+    sub_layers=3,
     seed=0,
     randomize_seed=True,
     prompt="",
@@ -248,70 +473,79 @@ def refine_selected_layer(
     cfg_norm=True,
     use_en_prompt=True,
     resolution=1024,
-    gpu_duration="1000",
+    gpu_duration=1000,
 ):
-    if not layer_paths:
-        return [], None, None
+    if not history or not current_node_id:
+        raise gr.Error("Сначала сделай Decompose (создай root-узел).")
+    if not current_layer_paths:
+        raise gr.Error("Нет слоёв в текущем узле (state).")
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    parent = find_node(history, current_node_id)
+    if not parent:
+        raise gr.Error("Текущий узел не найден в history.")
 
-    # Clamp index into existing layers
-    n = len(layer_paths)
-    idx = _clamp_int(refine_layer_index, default=1, lo=1, hi=n) - 1
+    resolution = norm_resolution(resolution)
+    sub_layers = clamp_int(sub_layers, default=3, lo=2, hi=10)
 
-    sub_layers = _clamp_int(refine_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])
 
-    resolution = _clamp_int(resolution, default=1024, lo=640, hi=1024)
-    if resolution not in (640, 1024):
-        resolution = 1024
+    pipe = get_pipeline()
+    ensure_device_strategy(pipe)
 
-    selected_path = layer_paths[idx]
-    selected_layer_img = Image.open(selected_path).convert("RGBA")
+    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
+    )
 
-    gen = torch.Generator(device="cpu").manual_seed(seed)
+    # 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)
 
-    inputs = {
-        "image": selected_layer_img,
-        "generator": gen,
-        "true_cfg_scale": float(true_guidance_scale),
-        "prompt": prompt if prompt else None,
-        "negative_prompt": neg_prompt,
-        "num_inference_steps": int(num_inference_steps),
-        "num_images_per_prompt": 1,
-        "layers": int(sub_layers),      # <-- ключевой параметр рекурсивной декомпозиции
-        "resolution": int(resolution),  # тот же resolution (без отдельных опций для refine)
-        "cfg_normalize": bool(cfg_norm),
-        "use_en_prompt": bool(use_en_prompt),
+    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}"
 
-    print("REFINE INPUTS:", {k: v for k, v in inputs.items() if k != "image"})
-    print("REQUESTED GPU DURATION:", gpu_duration)
-    print(f"REFINE: base layer index={idx+1}/{n}, sub_layers={sub_layers}")
-
-    with torch.inference_mode():
-        out = pipeline(**inputs)
-        refined_images = out.images[0]
-
-    refined_paths = []
-    refined_gallery = []
-    for img in refined_images:
-        refined_gallery.append(img)
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-        img.save(tmp.name)
-        refined_paths.append(tmp.name)
+    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)
 
-    refined_pptx = imagelist_to_pptx(refined_paths)
+    # Update history dropdown and set current node to the new child
+    hist_dd = gr.update(choices=history_choices(history), value=node_id)
 
-    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmpzip:
-        with zipfile.ZipFile(tmpzip.name, "w", zipfile.ZIP_DEFLATED) as zipf:
-            for i, p in enumerate(refined_paths):
-                zipf.write(p, f"sub_layer_{i+1}.png")
-        refined_zip = tmpzip.name
+    images, strip, pptx, zipp, layer_dd, sel_idx0, preview, breadcrumb = node_to_ui(history, node_id)
 
-    return refined_gallery, refined_pptx, refined_zip
+    return (
+        history,
+        node_id,
+        hist_dd,
+        layer_paths,  # current_layer_paths_state
+        images,
+        strip,
+        pptx,
+        zipp,
+        layer_dd,
+        sel_idx0,
+        preview,
+        breadcrumb,
+    )
 
 
+# ----------------------------
+# UI
+# ----------------------------
 ensure_dirname(LOG_DIR)
 examples = [f"assets/test_images/{i}.png" for i in range(1, 14)]
 
@@ -321,169 +555,233 @@ with gr.Blocks() as demo:
             '<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;">'
         )
+
         gr.Markdown(
             """
-The text prompt is intended to describe the overall content of the input image—including elements that may be partially occluded.
-It is not designed to control the semantic content of individual layers explicitly.
+Prompt описывает изображение в целом (включая частично закрытые элементы).
+Refine делает рекурсивную декомпозицию выбранного слоя текущего узла (узлы сохраняются в History).
             """
         )
 
-        # State to store layer PNG paths from last Decompose
-        layer_paths_state = gr.State([])
+        # 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)
 
         with gr.Row():
             with gr.Column(scale=1):
                 input_image = gr.Image(label="Input Image", image_mode="RGBA")
 
-                with gr.Accordion("Advanced Settings", open=False):
-                    prompt = gr.Textbox(
-                        label="Prompt (Optional)",
-                        placeholder="Please enter the prompt to describe the image (optional)",
-                        value="",
-                        lines=2,
-                    )
-                    neg_prompt = gr.Textbox(
-                        label="Negative Prompt (Optional)",
-                        placeholder="Please enter the negative prompt",
-                        value=" ",
-                        lines=2,
-                    )
+                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)
 
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-                    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
-                    )
+                    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)
 
-                    layer = gr.Slider(label="Layers", minimum=2, maximum=10, step=1, value=7)
+                    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)
 
-                    # default 1024 as you asked earlier
-                    resolution = gr.Radio(
-                        label="Processing resolution",
-                        choices=[640, 1024],
-                        value=1024,
-                    )
+                    resolution = gr.Radio(label="Resolution", choices=[640, 1024], value=1024)
 
-                    cfg_norm = gr.Checkbox(label="Whether enable CFG normalization", value=True)
-                    use_en_prompt = gr.Checkbox(
-                        label="Automatic caption language if no prompt provided, True for EN, False for ZH",
-                        value=True,
-                    )
+                    cfg_norm = gr.Checkbox(label="CFG normalize", value=True)
+                    use_en_prompt = gr.Checkbox(label="Auto caption language (EN=True / ZH=False)", value=True)
 
                     gpu_duration = gr.Textbox(
                         label="GPU duration override (seconds, 20..1500)",
                         value="1000",
                         lines=1,
-                        placeholder="e.g. 60, 120, 300, 1000, 1500",
                     )
 
-                decompose_btn = gr.Button("Decompose!", variant="primary")
+                decompose_btn = gr.Button("Decompose (new root node)", variant="primary")
+                refine_btn = gr.Button("Refine selected layer (create child node)", variant="secondary")
 
-                with gr.Accordion("Refine layer (Recursive Decomposition)", open=False):
-                    refine_layer_index = gr.Slider(
-                        label="Refine layer index (1 = first layer)",
-                        minimum=1,
-                        maximum=7,
-                        step=1,
-                        value=1,
-                    )
-                    refine_sub_layers = gr.Slider(
-                        label="Sub-layers (how many to split selected layer into)",
-                        minimum=2,
-                        maximum=10,
-                        step=1,
-                        value=3,
-                    )
-                    refine_btn = gr.Button("Refine selected layer", variant="secondary")
+                gr.Markdown("### History (nodes)")
+                history_dd = gr.Dropdown(label="Select node", choices=[], value=None, interactive=True)
+
+                breadcrumb_md = gr.Markdown("**Node path:** —")
+
+                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)
 
             with gr.Column(scale=2):
-                gallery = gr.Gallery(label="Layers", columns=4, rows=1, format="png")
-                with gr.Row():
-                    export_file = gr.File(label="Download PPTX")
-                    export_zip_file = gr.File(label="Download ZIP")
+                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)
 
-                gr.Markdown("### Refined (sub-layers)")
-                refined_gallery = gr.Gallery(label="Sub-layers", columns=4, rows=1, format="png")
                 with gr.Row():
-                    refined_export_file = gr.File(label="Download refined PPTX")
-                    refined_export_zip_file = gr.File(label="Download refined ZIP")
-
-    # Examples run Decompose
-    gr.Examples(
-        examples=examples,
-        inputs=[input_image],
-        outputs=[
-            gallery,
-            export_file,
-            export_zip_file,
-            layer_paths_state,
-            refined_gallery,
-            refined_export_file,
-            refined_export_zip_file,
-            refine_layer_index,  # update slider bounds/value
-        ],
-        fn=decompose,
-        examples_per_page=14,
-        cache_examples=False,
-        run_on_click=True,
-    )
+                    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,
+        )
 
-    # Decompose button
-    decompose_btn.click(
-        fn=decompose,
-        inputs=[
-            input_image,
-            seed,
-            randomize_seed,
-            prompt,
-            neg_prompt,
-            true_guidance_scale,
-            num_inference_steps,
-            layer,
-            cfg_norm,
-            use_en_prompt,
-            resolution,
-            gpu_duration,
-            refine_layer_index,  # so we can clamp nicely after new decomposition
-            refine_sub_layers,   # for duration signature parity
-        ],
-        outputs=[
-            gallery,
-            export_file,
-            export_zip_file,
-            layer_paths_state,
-            refined_gallery,
-            refined_export_file,
-            refined_export_zip_file,
-            refine_layer_index,  # update slider bounds/value
-        ],
-    )
+        # 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,
+            ],
+        )
 
-    # Refine button
-    refine_btn.click(
-        fn=refine_selected_layer,
-        inputs=[
-            layer_paths_state,
-            refine_layer_index,
-            refine_sub_layers,
-            seed,
-            randomize_seed,
-            prompt,
-            neg_prompt,
-            true_guidance_scale,
-            num_inference_steps,
-            cfg_norm,
-            use_en_prompt,
-            resolution,
-            gpu_duration,
-        ],
-        outputs=[refined_gallery, refined_export_file, refined_export_zip_file],
-    )
+        # 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 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,
+            )
+
+        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,
+            ],
+        )
+
+        # 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],
+        )
+
+        # 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],
+        )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue()
+    try:
+        demo.launch(ssr_mode=False)
+    except TypeError:
+        demo.launch()