Update app.py

app.py

@@ -4,7 +4,6 @@ import numpy as np
 import random
 import tempfile
 import zipfile
-import threading
 
 import spaces
 import torch
@@ -25,55 +24,9 @@ login(token=os.environ.get("hf"))
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ----------------------------
-# Pipeline singleton (fast path)
-# ----------------------------
-_PIPELINE = None
-_PIPELINE_LOCK = threading.Lock()
-
-
-def _enable_fast_cuda_settings():
-    if not torch.cuda.is_available():
-        return
-    try:
-        torch.backends.cuda.matmul.allow_tf32 = True
-        torch.backends.cudnn.allow_tf32 = True
-        torch.backends.cudnn.benchmark = True
-        torch.set_float32_matmul_precision("high")
-        try:
-            torch.backends.cuda.enable_flash_sdp(True)
-            torch.backends.cuda.enable_mem_efficient_sdp(True)
-            torch.backends.cuda.enable_math_sdp(False)
-        except Exception:
-            pass
-    except Exception:
-        pass
-
-
-def get_pipeline():
-    global _PIPELINE
-    if _PIPELINE is not None:
-        return _PIPELINE
-
-    with _PIPELINE_LOCK:
-        if _PIPELINE is not None:
-            return _PIPELINE
-
-        _enable_fast_cuda_settings()
-
-        pipe = QwenImageLayeredPipeline.from_pretrained(
-            "Qwen/Qwen-Image-Layered",
-            torch_dtype=dtype,
-        )
-
-        # Fastest mode: keep weights on GPU if available
-        if device == "cuda":
-            pipe.to("cuda")
-        else:
-            pipe.to("cpu")
-
-        _PIPELINE = pipe
-        return _PIPELINE
+pipeline = QwenImageLayeredPipeline.from_pretrained(
+    "Qwen/Qwen-Image-Layered", torch_dtype=dtype
+).to(device)
 
 
 def ensure_dirname(path: str):
@@ -115,6 +68,29 @@ def imagelist_to_pptx(img_files):
         return tmp.name
 
 
+def export_zip_from_pil(images):
+    paths = []
+    for img in images:
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        img.save(tmp.name)
+        paths.append(tmp.name)
+
+    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(paths):
+                zipf.write(p, f"layer_{i+1}.png")
+        return tmpzip.name
+
+
+def export_pptx_from_pil(images):
+    paths = []
+    for img in images:
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        img.save(tmp.name)
+        paths.append(tmp.name)
+    return imagelist_to_pptx(paths)
+
+
 def _clamp_int(x, default: int, lo: int, hi: int) -> int:
     try:
         v = int(x)
@@ -123,72 +99,95 @@ def _clamp_int(x, default: int, lo: int, hi: int) -> int:
     return max(lo, min(hi, v))
 
 
-def _normalize_rgba(pil: Image.Image) -> Image.Image:
-    if pil.mode != "RGBA":
-        pil = pil.convert("RGB").convert("RGBA")
-    return pil
+def _norm_resolution(x):
+    x = _clamp_int(x, default=640, lo=640, hi=1024)
+    if x not in (640, 1024):
+        return 640
+    return x
+
+
+def _norm_image(input_image):
+    if isinstance(input_image, list):
+        input_image = input_image[0]
+
+    if isinstance(input_image, str):
+        return Image.open(input_image).convert("RGB").convert("RGBA")
+    if isinstance(input_image, Image.Image):
+        return input_image.convert("RGB").convert("RGBA")
+    if isinstance(input_image, np.ndarray):
+        return Image.fromarray(input_image).convert("RGB").convert("RGBA")
+
+    raise ValueError(f"Unsupported input_image type: {type(input_image)}")
+
+
+def _make_node(
+    name,
+    parent_id,
+    images,
+    params,
+    refine_meta=None,
+):
+    node_id = random_str(10)
+    return {
+        "id": node_id,
+        "name": name,
+        "parent": parent_id,
+        "children": [],
+        "images": images,  # list[PIL.Image]
+        "params": params,  # dict
+        "refine_meta": refine_meta,  # dict | None
+    }
 
 
 def _history_choices(history):
-    """
-    history: list[dict] with keys:
-      id, parent, title, layers(list[PIL]), meta(optional)
-    """
+    # Dropdown choices: list of (label, value)
+    nodes = history.get("nodes", {})
+    order = history.get("order", [])
     choices = []
-    by_id = {n["id"]: n for n in history}
-    for i, node in enumerate(history):
-        n_layers = len(node.get("layers", []) or [])
-        parent = node.get("parent")
-
-        depth = 0
-        pid = parent
-        seen = set()
-        while pid and pid in by_id and pid not in seen:
-            seen.add(pid)
-            depth += 1
-            pid = by_id[pid].get("parent")
-
-        prefix = "  " * min(depth, 6)
-        label = f"{prefix}{i+1}. {node.get('title','Node')}  (layers={n_layers})"
-        choices.append((label, node["id"]))
+    for nid in order:
+        n = nodes.get(nid)
+        if not n:
+            continue
+        cnt = len(n.get("images") or [])
+        label = f"{n.get('name','node')} · {cnt} layers · {nid}"
+        choices.append((label, nid))
     return choices
 
 
-def _find_node(history, node_id):
-    for n in history:
-        if n.get("id") == node_id:
-            return n
-    return None
+def _chips_for_node(history, node_id):
+    nodes = history.get("nodes", {})
+    if node_id not in nodes:
+        return ""
 
+    n = nodes[node_id]
+    parent = n.get("parent")
+    children = n.get("children") or []
+    root = history.get("root")
 
-def _layers_to_temp_pngs(layers):
-    temp_files = []
-    for img in layers:
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-        _normalize_rgba(img).save(tmp.name)
-        temp_files.append(tmp.name)
-    return temp_files
+    tags = []
+    if node_id == root:
+        tags.append("[root]")
+    if parent:
+        tags.append(f"[parent: {parent}]")
+    else:
+        tags.append("[parent: —]")
+    tags.append(f"[children: {len(children)}]")
+    return " ".join(tags)
 
 
-def _export_zip_from_layers(layers):
-    temp_files = _layers_to_temp_pngs(layers)
-    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(temp_files):
-                zipf.write(p, f"layer_{i+1}.png")
-        return tmpzip.name
+def _get_current_node(history, node_id):
+    nodes = history.get("nodes", {})
+    return nodes.get(node_id)
 
 
-def _export_pptx_from_layers(layers):
-    temp_files = _layers_to_temp_pngs(layers)
-    return imagelist_to_pptx(temp_files)
+def _generator_for_seed(seed):
+    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
+    return torch.Generator(device=gen_device).manual_seed(seed)
 
 
-# ----------------------------
-# ZeroGPU duration helper
-# ----------------------------
-def get_duration(
-    input_image=None,
+# Dynamic duration callable: must accept the same args as on_decompose_click(). It returns seconds.
+def get_duration_decompose(
+    input_image,
     seed=777,
     randomize_seed=False,
     prompt=None,
@@ -200,16 +199,23 @@ def get_duration(
     use_en_prompt=True,
     resolution=640,
     gpu_duration=1000,
-    **kwargs,
 ):
     return _clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
 
 
-# ----------------------------
-# GPU ops
-# ----------------------------
-@spaces.GPU(duration=get_duration)
-def run_decompose_gpu(
+# Dynamic duration callable for refine (same args + refine-specific)
+def get_duration_refine(
+    history,
+    current_node_id,
+    refine_layer_index=0,
+    sub_layers=3,
+    gpu_duration=1000,
+):
+    return _clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
+
+
+@spaces.GPU(duration=get_duration_decompose)
+def on_decompose_click(
     input_image,
     seed=777,
     randomize_seed=False,
@@ -223,35 +229,28 @@ def run_decompose_gpu(
     resolution=640,
     gpu_duration=1000,
 ):
+    # Seed
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    resolution = _clamp_int(resolution, default=640, lo=640, hi=1024)
-    if resolution not in (640, 1024):
-        resolution = 640
+    resolution = _norm_resolution(resolution)
+    pil_image = _norm_image(input_image)
 
-    if isinstance(input_image, list):
-        input_image = input_image[0]
-
-    if isinstance(input_image, str):
-        pil_image = Image.open(input_image)
-    elif isinstance(input_image, Image.Image):
-        pil_image = input_image
-    elif isinstance(input_image, np.ndarray):
-        pil_image = Image.fromarray(input_image)
-    else:
-        raise ValueError(f"Unsupported input_image type: {type(input_image)}")
-
-    pil_image = _normalize_rgba(pil_image)
-
-    pipe = get_pipeline()
-
-    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
-    generator = torch.Generator(device=gen_device).manual_seed(int(seed))
+    params = {
+        "seed": seed,
+        "prompt": prompt,
+        "negative_prompt": neg_prompt,
+        "true_cfg_scale": true_guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "layers": layer,
+        "resolution": resolution,
+        "cfg_normalize": cfg_norm,
+        "use_en_prompt": use_en_prompt,
+    }
 
     inputs = {
         "image": pil_image,
-        "generator": generator,
+        "generator": _generator_for_seed(seed),
         "true_cfg_scale": true_guidance_scale,
         "prompt": prompt,
         "negative_prompt": neg_prompt,
@@ -263,471 +262,349 @@ def run_decompose_gpu(
         "use_en_prompt": use_en_prompt,
     }
 
+    print("DECOMPOSE INPUTS:", inputs)
+    print("REQUESTED GPU DURATION:", gpu_duration)
+
     with torch.inference_mode():
-        if torch.cuda.is_available():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                out = pipe(**inputs)
-        else:
-            out = pipe(**inputs)
+        out = pipeline(**inputs)
+        output_images = out.images[0]  # list of PIL
+
+    # New history (reset)
+    history = {"nodes": {}, "order": [], "root": None}
+
+    root_node = _make_node(
+        name="Decompose (root)",
+        parent_id=None,
+        images=output_images,
+        params=params,
+        refine_meta=None,
+    )
+    history["nodes"][root_node["id"]] = root_node
+    history["order"].append(root_node["id"])
+    history["root"] = root_node["id"]
 
-    layers_out = out.images[0]
-    layers_out = [_normalize_rgba(x) for x in layers_out]
-    return layers_out
+    current_node_id = root_node["id"]
 
+    # History UI
+    choices = _history_choices(history)
+    chips = _chips_for_node(history, current_node_id)
 
-@spaces.GPU(duration=get_duration)
-def run_refine_gpu(
-    base_layers,
-    selected_index: int,
-    seed=777,
-    randomize_seed=False,
-    prompt=None,
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
+    # Layer selection defaults
+    refine_layer_index = 0
+    refine_layer_dropdown_choices = [f"Layer {i+1}" for i in range(len(output_images))]
+    refine_layer_dropdown_value = (
+        refine_layer_dropdown_choices[0] if refine_layer_dropdown_choices else None
+    )
+
+    # Clear exports on new run
+    export_pptx = None
+    export_zip = None
+
+    # Refined output empty
+    refined_gallery = []
+
+    return (
+        history,
+        current_node_id,
+        output_images,  # decomposed gallery
+        output_images,  # picker gallery (1 row)
+        gr.update(choices=choices, value=current_node_id),  # history dropdown
+        gr.update(value=refine_layer_index),  # refine layer index state
+        gr.update(choices=refine_layer_dropdown_choices, value=refine_layer_dropdown_value),
+        chips,
+        refined_gallery,
+        export_pptx,
+        export_zip,
+        gr.update(open=False),  # refined accordion closed
+    )
+
+
+@spaces.GPU(duration=get_duration_refine)
+def on_refine_click(
+    history,
+    current_node_id,
+    refine_layer_index=0,
     sub_layers=3,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=640,
     gpu_duration=1000,
 ):
-    if not base_layers or not isinstance(base_layers, list):
-        raise ValueError("No base layers to refine. Run Decompose first.")
+    if not history or not current_node_id:
+        raise gr.Error("No active decomposition yet. Run Decompose first.")
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    node = _get_current_node(history, current_node_id)
+    if not node:
+        raise gr.Error("Current node not found in history.")
 
-    resolution = _clamp_int(resolution, default=640, lo=640, hi=1024)
-    if resolution not in (640, 1024):
-        resolution = 640
+    images = node.get("images") or []
+    if not images:
+        raise gr.Error("Current node has no images to refine.")
 
-    sub_layers = _clamp_int(sub_layers, default=3, lo=2, hi=10)
+    idx = _clamp_int(refine_layer_index, default=0, lo=0, hi=max(0, len(images) - 1))
+    if idx >= len(images):
+        idx = 0
 
-    idx = _clamp_int(selected_index, default=0, lo=0, hi=len(base_layers) - 1)
-    selected_layer = _normalize_rgba(base_layers[idx])
+    selected_layer = images[idx]
 
-    pipe = get_pipeline()
+    # Reuse params from this node (no separate refine steps/resolution/cfg)
+    p = node.get("params") or {}
+    seed = p.get("seed", 777)
+    prompt = p.get("prompt", None)
+    neg_prompt = p.get("negative_prompt", " ")
+    true_guidance_scale = p.get("true_cfg_scale", 4.0)
+    num_inference_steps = p.get("num_inference_steps", 50)
+    resolution = p.get("resolution", 640)
+    cfg_norm = p.get("cfg_normalize", True)
+    use_en_prompt = p.get("use_en_prompt", True)
 
-    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
-    generator = torch.Generator(device=gen_device).manual_seed(int(seed))
+    sub_layers = _clamp_int(sub_layers, default=3, lo=2, hi=10)
 
     inputs = {
         "image": selected_layer,
-        "generator": generator,
+        "generator": _generator_for_seed(seed),
         "true_cfg_scale": true_guidance_scale,
         "prompt": prompt,
         "negative_prompt": neg_prompt,
         "num_inference_steps": num_inference_steps,
         "num_images_per_prompt": 1,
-        "layers": sub_layers,
-        "resolution": resolution,
-        "cfg_normalize": cfg_norm,
+        "layers": sub_layers,       # <-- sub-layers
+        "resolution": resolution,    # <-- reuse
+        "cfg_normalize": cfg_norm,   # <-- reuse
         "use_en_prompt": use_en_prompt,
     }
 
-    with torch.inference_mode():
-        if torch.cuda.is_available():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                out = pipe(**inputs)
-        else:
-            out = pipe(**inputs)
-
-    refined = out.images[0]
-    refined = [_normalize_rgba(x) for x in refined]
-    return refined
+    print("REFINE INPUTS:", inputs)
+    print("REQUESTED GPU DURATION:", gpu_duration)
 
+    with torch.inference_mode():
+        out = pipeline(**inputs)
+        refined_images = out.images[0]
 
-# ----------------------------
-# Gradio glue (history + UX)
-# ----------------------------
-def _init_state():
-    return {
-        "history": [],
-        "active_node_id": None,
-        "selected_layer_idx": 0,
+    refine_meta = {
+        "from_node": current_node_id,
+        "layer_index": idx,
+        "sub_layers": sub_layers,
     }
 
-
-def _set_active_node(state, node_id):
-    state["active_node_id"] = node_id
-    state["selected_layer_idx"] = 0
-    return state
-
-
-def _node_layers_and_picker_updates(node):
-    layers_out = node.get("layers") or []
-    layer_choices = [(f"Layer {i+1}", i) for i in range(len(layers_out))]
-    return layers_out, layer_choices
-
-
-def on_decompose_click(
-    input_image,
-    seed,
-    randomize_seed,
-    prompt,
-    neg_prompt,
-    true_guidance_scale,
-    num_inference_steps,
-    layer,
-    cfg_norm,
-    use_en_prompt,
-    resolution,
-    gpu_duration,
-    state,
-):
-    if state is None:
-        state = _init_state()
-
-    layers_out = run_decompose_gpu(
-        input_image=input_image,
-        seed=seed,
-        randomize_seed=randomize_seed,
-        prompt=prompt,
-        neg_prompt=neg_prompt,
-        true_guidance_scale=true_guidance_scale,
-        num_inference_steps=num_inference_steps,
-        layer=layer,
-        cfg_norm=cfg_norm,
-        use_en_prompt=use_en_prompt,
-        resolution=resolution,
-        gpu_duration=gpu_duration,
+    child = _make_node(
+        name=f"Refine L{idx+1} → {sub_layers}",
+        parent_id=current_node_id,
+        images=refined_images,
+        params=p,
+        refine_meta=refine_meta,
     )
 
-    node_id = random_str(10)
-    node = {
-        "id": node_id,
-        "parent": None,
-        "title": "Decompose",
-        "layers": layers_out,
-        "meta": {"type": "decompose"},
-    }
+    # Insert child into history
+    history["nodes"][child["id"]] = child
+    history["order"].append(child["id"])
+    history["nodes"][current_node_id].setdefault("children", []).append(child["id"])
 
-    state["history"].append(node)
-    _set_active_node(state, node_id)
+    # Move current to child
+    current_node_id = child["id"]
 
-    choices = _history_choices(state["history"])
-    _, layer_choices = _node_layers_and_picker_updates(node)
+    # Update history dropdown
+    choices = _history_choices(history)
+    chips = _chips_for_node(history, current_node_id)
 
-    return (
-        state,
-        choices,
-        node_id,          # selected history node
-        layers_out,       # base gallery
-        layers_out,       # picker gallery
-        layer_choices,    # dropdown choices
-        0,                # dropdown selected index
-        gr.Accordion.update(open=False),
-        [],               # refined gallery cleared
-        node.get("title", ""),
+    # Update layer pickers for new current node
+    refine_layer_index = 0
+    refine_layer_dropdown_choices = [f"Layer {i+1}" for i in range(len(refined_images))]
+    refine_layer_dropdown_value = (
+        refine_layer_dropdown_choices[0] if refine_layer_dropdown_choices else None
     )
 
-
-def on_history_change(node_id, state):
-    if state is None:
-        state = _init_state()
-
-    node = _find_node(state["history"], node_id)
-    if not node:
-        return (
-            state,
-            [],
-            [],
-            [],
-            0,
-            gr.Accordion.update(open=False),
-            [],
-            "",
-        )
-
-    _set_active_node(state, node_id)
-    layers_out, layer_choices = _node_layers_and_picker_updates(node)
-
+    # Auto open refined accordion (and collapse refined selection is handled via updates)
     return (
-        state,
-        layers_out,
-        layers_out,
-        layer_choices,
-        0,
-        gr.Accordion.update(open=False),
-        [],
-        node.get("title", ""),
+        history,
+        current_node_id,
+        refined_images,  # decomposed gallery now shows current node
+        refined_images,  # picker gallery
+        gr.update(choices=choices, value=current_node_id),
+        gr.update(value=refine_layer_index),
+        gr.update(choices=refine_layer_dropdown_choices, value=refine_layer_dropdown_value),
+        chips,
+        refined_images,  # refined gallery
+        None,  # export pptx reset
+        None,  # export zip reset
+        gr.update(open=True),  # refined accordion open  ✅ (replaced Accordion.update)
     )
 
 
-def on_picker_select(evt: gr.SelectData, state):
-    if state is None:
-        state = _init_state()
-    idx = int(evt.index) if evt and evt.index is not None else 0
-    state["selected_layer_idx"] = idx
-    return state, idx
+def on_picker_select(evt: gr.SelectData):
+    # evt.index for Gallery is int when selecting an item
+    try:
+        return int(evt.index)
+    except Exception:
+        return 0
 
 
-def on_layer_dropdown_change(layer_idx, state):
-    if state is None:
-        state = _init_state()
+def on_refine_layer_dropdown_change(label):
+    # label is "Layer K"
+    if not label:
+        return 0
     try:
-        idx = int(layer_idx)
+        k = int(str(label).split()[-1])
+        return max(0, k - 1)
     except Exception:
-        idx = 0
-    state["selected_layer_idx"] = idx
-    return state
-
-
-def _append_refine_node(state, parent_node, selected_idx, sub_layers_value, refined_layers):
-    new_id = random_str(10)
-    new_node = {
-        "id": new_id,
-        "parent": parent_node["id"],
-        "title": f"Refine: Layer {selected_idx+1}",
-        "layers": refined_layers,
-        "meta": {
-            "type": "refine",
-            "refine_from": parent_node["id"],
-            "refine_layer_idx": int(selected_idx),
-            "sub_layers": int(sub_layers_value),
-        },
-    }
-    state["history"].append(new_node)
-    _set_active_node(state, new_id)
-    return new_node
+        return 0
 
 
-def on_refine_click(
-    seed,
-    randomize_seed,
-    prompt,
-    neg_prompt,
-    true_guidance_scale,
-    num_inference_steps,
-    cfg_norm,
-    use_en_prompt,
-    resolution,
-    gpu_duration,
-    sub_layers,
-    state,
-    history_node_id,
-    layer_dropdown_idx,
-):
-    if state is None:
-        state = _init_state()
+def on_history_change(history, node_id):
+    if not history or not node_id:
+        return (
+            None,
+            [],
+            [],
+            gr.update(),
+            gr.update(value=0),
+            gr.update(choices=[], value=None),
+            "",
+            [],
+            None,
+            None,
+            gr.update(open=False),  # refined accordion closed
+        )
 
-    node = _find_node(state["history"], history_node_id)
+    node = _get_current_node(history, node_id)
     if not node:
-        raise gr.Error("No active node selected. Run Decompose first.")
-
-    base_layers = node.get("layers") or []
-    if not base_layers:
-        raise gr.Error("Selected node has no layers to refine.")
+        return (
+            node_id,
+            [],
+            [],
+            gr.update(),
+            gr.update(value=0),
+            gr.update(choices=[], value=None),
+            "",
+            [],
+            None,
+            None,
+            gr.update(open=False),
+        )
 
-    try:
-        selected_idx = int(layer_dropdown_idx)
-    except Exception:
-        selected_idx = int(state.get("selected_layer_idx", 0) or 0)
-
-    refined_layers = run_refine_gpu(
-        base_layers=base_layers,
-        selected_index=selected_idx,
-        seed=seed,
-        randomize_seed=randomize_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,
-    )
+    imgs = node.get("images") or []
+    chips = _chips_for_node(history, node_id)
 
-    new_node = _append_refine_node(
-        state=state,
-        parent_node=node,
-        selected_idx=selected_idx,
-        sub_layers_value=sub_layers,
-        refined_layers=refined_layers,
+    refine_layer_index = 0
+    refine_layer_dropdown_choices = [f"Layer {i+1}" for i in range(len(imgs))]
+    refine_layer_dropdown_value = (
+        refine_layer_dropdown_choices[0] if refine_layer_dropdown_choices else None
     )
 
-    choices = _history_choices(state["history"])
-    _, layer_choices = _node_layers_and_picker_updates(new_node)
-
+    # Keep refined panel closed when user jumps around history
     return (
-        state,
-        choices,
-        new_node["id"],
-        refined_layers,                 # base gallery shows refined node
-        refined_layers,                 # picker shows refined node
-        layer_choices,
-        0,
-        gr.Accordion.update(open=True),
-        refined_layers,                 # refined gallery
-        new_node.get("title", ""),
+        node_id,
+        imgs,
+        imgs,
+        gr.update(choices=_history_choices(history), value=node_id),
+        gr.update(value=refine_layer_index),
+        gr.update(choices=refine_layer_dropdown_choices, value=refine_layer_dropdown_value),
+        chips,
+        [],
+        None,
+        None,
+        gr.update(open=False),  # ✅ replaced Accordion.update
     )
 
 
-def on_back_to_parent_click(state, history_node_id):
-    if state is None:
-        state = _init_state()
-
-    node = _find_node(state["history"], history_node_id)
-    if not node:
-        raise gr.Error("Select a node in History.")
-
-    parent_id = node.get("parent")
-    if not parent_id:
-        # already root
-        layers_out, layer_choices = _node_layers_and_picker_updates(node)
+def on_back_to_parent(history, current_node_id):
+    if not history or not current_node_id:
         return (
-            state,
-            history_node_id,
-            layers_out,
-            layers_out,
-            layer_choices,
-            0,
-            gr.Accordion.update(open=False),
+            current_node_id,
+            [],
             [],
-            node.get("title", ""),
+            gr.update(),
+            gr.update(value=0),
+            gr.update(choices=[], value=None),
+            "",
+            [],
+            None,
+            None,
+            gr.update(open=False),
         )
 
-    parent = _find_node(state["history"], parent_id)
+    node = _get_current_node(history, current_node_id)
+    parent = node.get("parent") if node else None
     if not parent:
-        raise gr.Error("Parent not found in history (corrupted history).")
+        # already at root or missing parent
+        parent = current_node_id
 
-    _set_active_node(state, parent_id)
-    layers_out, layer_choices = _node_layers_and_picker_updates(parent)
+    return on_history_change(history, parent)
 
-    return (
-        state,
-        parent_id,
-        layers_out,
-        layers_out,
-        layer_choices,
-        0,
-        gr.Accordion.update(open=False),
-        [],
-        parent.get("title", ""),
-    )
 
-
-def on_redo_refine_click(
-    seed,
-    randomize_seed,
-    prompt,
-    neg_prompt,
-    true_guidance_scale,
-    num_inference_steps,
-    cfg_norm,
-    use_en_prompt,
-    resolution,
-    gpu_duration,
-    state,
-    history_node_id,
-):
-    if state is None:
-        state = _init_state()
-
-    node = _find_node(state["history"], history_node_id)
+def on_redo_refine(history, current_node_id, gpu_duration=1000):
+    # If current node is a refined node, redo the same refine from its parent with same meta
+    if not history or not current_node_id:
+        raise gr.Error("No active node.")
+    node = _get_current_node(history, current_node_id)
     if not node:
-        raise gr.Error("Select a node in History.")
+        raise gr.Error("Node not found.")
+    meta = node.get("refine_meta")
+    if not meta:
+        raise gr.Error("This node has no refine metadata to redo (not a refined node).")
 
-    meta = node.get("meta") or {}
-    if meta.get("type") != "refine":
-        raise gr.Error("Redo refine работает только для refine-узлов (не для Decompose).")
+    parent_id = meta.get("from_node")
+    layer_index = meta.get("layer_index", 0)
+    sub_layers = meta.get("sub_layers", 3)
 
-    parent_id = meta.get("refine_from") or node.get("parent")
-    if not parent_id:
-        raise gr.Error("Refine node has no parent info.")
+    # Temporarily switch to parent for redo logic
+    return on_refine_click(history, parent_id, layer_index, sub_layers, gpu_duration)
 
-    parent = _find_node(state["history"], parent_id)
-    if not parent:
-        raise gr.Error("Parent not found in history.")
-
-    base_layers = parent.get("layers") or []
-    if not base_layers:
-        raise gr.Error("Parent node has no layers.")
-
-    selected_idx = int(meta.get("refine_layer_idx", 0))
-    sub_layers_value = int(meta.get("sub_layers", 3))
-
-    refined_layers = run_refine_gpu(
-        base_layers=base_layers,
-        selected_index=selected_idx,
-        seed=seed,
-        randomize_seed=randomize_seed,
-        prompt=prompt,
-        neg_prompt=neg_prompt,
-        true_guidance_scale=true_guidance_scale,
-        num_inference_steps=num_inference_steps,
-        sub_layers=sub_layers_value,
-        cfg_norm=cfg_norm,
-        use_en_prompt=use_en_prompt,
-        resolution=resolution,
-        gpu_duration=gpu_duration,
-    )
 
-    new_node = _append_refine_node(
-        state=state,
-        parent_node=parent,
-        selected_idx=selected_idx,
-        sub_layers_value=sub_layers_value,
-        refined_layers=refined_layers,
+def on_duplicate_node(history, current_node_id):
+    if not history or not current_node_id:
+        raise gr.Error("No active node to duplicate.")
+
+    node = _get_current_node(history, current_node_id)
+    if not node:
+        raise gr.Error("Node not found.")
+
+    dup = _make_node(
+        name=f"{node.get('name','node')} (copy)",
+        parent_id=node.get("parent"),
+        images=node.get("images") or [],
+        params=node.get("params") or {},
+        refine_meta=node.get("refine_meta"),
     )
 
-    choices = _history_choices(state["history"])
-    _, layer_choices = _node_layers_and_picker_updates(new_node)
+    history["nodes"][dup["id"]] = dup
+    history["order"].append(dup["id"])
 
-    return (
-        state,
-        choices,
-        new_node["id"],
-        refined_layers,
-        refined_layers,
-        layer_choices,
-        0,
-        gr.Accordion.update(open=True),
-        refined_layers,
-        new_node.get("title", ""),
-    )
+    # Attach to same parent if any
+    parent = dup.get("parent")
+    if parent and parent in history["nodes"]:
+        history["nodes"][parent].setdefault("children", []).append(dup["id"])
 
+    # Jump to duplicated node
+    return on_history_change(history, dup["id"])
 
-def on_rename_node_click(state, history_node_id, new_name):
-    if state is None:
-        state = _init_state()
 
-    node = _find_node(state["history"], history_node_id)
+def on_rename_node(history, current_node_id, new_name):
+    if not history or not current_node_id:
+        raise gr.Error("No active node.")
+    node = _get_current_node(history, current_node_id)
     if not node:
-        raise gr.Error("Select a node in History.")
-
+        raise gr.Error("Node not found.")
     new_name = (new_name or "").strip()
     if not new_name:
-        # no-op
-        choices = _history_choices(state["history"])
-        return state, choices, history_node_id, node.get("title", "")
+        raise gr.Error("Name cannot be empty.")
+    node["name"] = new_name
+    # Update dropdown label list
+    return gr.update(choices=_history_choices(history), value=current_node_id)
 
-    node["title"] = new_name
-    choices = _history_choices(state["history"])
-    return state, choices, history_node_id, node.get("title", "")
 
-
-def on_export_click(state, node_id, export_kind: str):
-    if state is None:
-        state = _init_state()
-    node = _find_node(state["history"], node_id)
+def on_export_current(history, current_node_id):
+    if not history or not current_node_id:
+        raise gr.Error("No active node.")
+    node = _get_current_node(history, current_node_id)
     if not node:
-        raise gr.Error("Select a node in History to export.")
-    layers = node.get("layers") or []
-    if not layers:
-        raise gr.Error("Selected node has no layers to export.")
-    if export_kind == "pptx":
-        return _export_pptx_from_layers(layers)
-    if export_kind == "zip":
-        return _export_zip_from_layers(layers)
-    raise gr.Error("Unknown export kind.")
-
-
-# ----------------------------
-# UI
-# ----------------------------
+        raise gr.Error("Node not found.")
+    imgs = node.get("images") or []
+    if not imgs:
+        raise gr.Error("Node has no images to export.")
+
+    pptx_path = export_pptx_from_pil(imgs)
+    zip_path = export_zip_from_pil(imgs)
+    return pptx_path, zip_path
+
+
 ensure_dirname(LOG_DIR)
 
 examples = [
@@ -747,13 +624,17 @@ examples = [
 ]
 
 with gr.Blocks() as demo:
-    state = gr.State(_init_state())
+    # Server-side state
+    history_state = gr.State(None)
+    current_node_id_state = gr.State(None)
+    refine_layer_index_state = gr.State(0)
 
     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;">'
         )
+
         gr.Markdown(
             """
 The text prompt is intended to describe the overall content of the input image—including elements that may be partially occluded (e.g., you may specify the text hidden behind a foreground object). It is not designed to control the semantic content of individual layers explicitly.
@@ -764,7 +645,7 @@ The text prompt is intended to describe the overall content of the input image
             with gr.Column(scale=1):
                 input_image = gr.Image(label="Input Image", image_mode="RGBA")
 
-                with gr.Accordion("Advanced Settings", open=False):
+                with gr.Accordion("Settings", open=False):
                     prompt = gr.Textbox(
                         label="Prompt (Optional)",
                         placeholder="Please enter the prompt to descibe the image. (Optional)",
@@ -778,7 +659,13 @@ The text prompt is intended to describe the overall content of the input image
                         lines=2,
                     )
 
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                    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(
@@ -811,7 +698,9 @@ The text prompt is intended to describe the overall content of the input image
                         value=640,
                     )
 
-                    cfg_norm = gr.Checkbox(label="Whether enable CFG normalization", value=True)
+                    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,
@@ -824,52 +713,23 @@ The text prompt is intended to describe the overall content of the input image
                         placeholder="e.g. 60, 120, 300, 1000, 1500",
                     )
 
-                decompose_btn = gr.Button("Decompose!", variant="primary")
-
-                with gr.Accordion("History", open=True):
-                    history_dropdown = gr.Dropdown(
-                        label="Nodes",
-                        choices=[],
-                        value=None,
-                        interactive=True,
-                    )
-
-                    with gr.Row():
-                        back_parent_btn = gr.Button("← Back to parent")
-                        redo_refine_btn = gr.Button("↺ Redo refine")
-
-                    branch_name = gr.Textbox(
-                        label="Branch name",
-                        value="",
-                        lines=1,
-                        placeholder="Rename selected node...",
-                    )
-                    rename_btn = gr.Button("Rename selected node")
-
-                    with gr.Row():
-                        export_pptx_btn = gr.Button("Export PPTX (selected node)")
-                        export_zip_btn = gr.Button("Export ZIP (selected node)")
+                run_button = gr.Button("Decompose!", variant="primary")
 
-                    export_pptx_file = gr.File(label="Download PPTX")
-                    export_zip_file = gr.File(label="Download ZIP")
-
-                with gr.Accordion("Refine layer", open=True):
+                with gr.Accordion("Refine (Recursive)", open=True):
                     gr.Markdown("Pick a layer visually (like Photoshop), then refine it into sub-layers.")
 
-                    layer_picker = gr.Gallery(
-                        label="Layer Picker (click a thumbnail)",
+                    # One-row “Photoshop-like” picker gallery
+                    picker_gallery = gr.Gallery(
+                        label="Layer picker (click a layer)",
                         columns=8,
                         rows=1,
-                        height="auto",
                         format="png",
-                        show_label=True,
                     )
 
-                    layer_idx_dropdown = gr.Dropdown(
-                        label="Refine layer index",
+                    refine_layer_dropdown = gr.Dropdown(
+                        label="Refine layer (fallback)",
                         choices=[],
-                        value=0,
-                        interactive=True,
+                        value=None,
                     )
 
                     sub_layers = gr.Slider(
@@ -880,23 +740,69 @@ The text prompt is intended to describe the overall content of the input image
                         value=3,
                     )
 
-                    refine_btn = gr.Button("Refine selected layer", variant="secondary")
+                    refine_button = gr.Button("Refine selected layer", variant="secondary")
 
             with gr.Column(scale=2):
-                base_gallery = gr.Gallery(label="Current node layers", columns=4, rows=1, format="png")
+                # History / navigation
+                with gr.Accordion("History", open=True):
+                    history_dropdown = gr.Dropdown(
+                        label="Nodes",
+                        choices=[],
+                        value=None,
+                    )
+                    chips_md = gr.Markdown("")
+
+                    with gr.Row():
+                        back_button = gr.Button("← back to parent")
+                        redo_button = gr.Button("↺ redo refine")
+                        dup_button = gr.Button("Duplicate node (branch)")
 
-                refined_accordion = gr.Accordion("Refined layers", open=False)
+                    with gr.Row():
+                        rename_text = gr.Textbox(label="Branch name", value="", lines=1)
+                        rename_button = gr.Button("Rename")
+
+                # Main outputs
+                decomp_accordion = gr.Accordion("Layers (Current node)", open=True)
+                with decomp_accordion:
+                    gallery = gr.Gallery(label="Layers", columns=4, rows=1, format="png")
+
+                refined_accordion = gr.Accordion("Refined layers (Latest refine)", open=False)
                 with refined_accordion:
-                    refined_gallery = gr.Gallery(label="Refined layers output", columns=4, rows=1, format="png")
+                    refined_gallery = gr.Gallery(label="Refined", columns=4, rows=1, format="png")
+
+                with gr.Row():
+                    export_button = gr.Button("Export ZIP/PPTX (current node)")
+                with gr.Row():
+                    export_file = gr.File(label="Download PPTX")
+                    export_zip_file = gr.File(label="Download ZIP")
 
+    # Examples (run decompose)
     gr.Examples(
         examples=examples,
         inputs=[input_image],
+        outputs=[gallery, export_file, export_zip_file],
+        fn=lambda img: ([], None, None),  # keep examples UI; actual run via click
+        examples_per_page=14,
         cache_examples=False,
+        run_on_click=False,
+    )
+
+    # Picker selection -> refine_layer_index_state
+    picker_gallery.select(
+        fn=on_picker_select,
+        inputs=None,
+        outputs=refine_layer_index_state,
+    )
+
+    # Dropdown selection -> refine_layer_index_state
+    refine_layer_dropdown.change(
+        fn=on_refine_layer_dropdown_change,
+        inputs=refine_layer_dropdown,
+        outputs=refine_layer_index_state,
     )
 
     # Decompose
-    decompose_btn.click(
+    run_button.click(
         fn=on_decompose_click,
         inputs=[
             input_image,
@@ -911,151 +817,138 @@ The text prompt is intended to describe the overall content of the input image
             use_en_prompt,
             resolution,
             gpu_duration,
-            state,
         ],
         outputs=[
-            state,
-            history_dropdown,
+            history_state,
+            current_node_id_state,
+            gallery,
+            picker_gallery,
             history_dropdown,
-            base_gallery,
-            layer_picker,
-            layer_idx_dropdown,
-            layer_idx_dropdown,
-            refined_accordion,
-            refined_gallery,
-            branch_name,
-        ],
-    )
-
-    # History change
-    history_dropdown.change(
-        fn=on_history_change,
-        inputs=[history_dropdown, state],
-        outputs=[
-            state,
-            base_gallery,
-            layer_picker,
-            layer_idx_dropdown,
-            layer_idx_dropdown,
-            refined_accordion,
+            refine_layer_index_state,
+            refine_layer_dropdown,
+            chips_md,
             refined_gallery,
-            branch_name,
+            export_file,
+            export_zip_file,
+            refined_accordion,  # gr.update(open=...) returned
         ],
     )
 
-    # Picker click
-    layer_picker.select(
-        fn=on_picker_select,
-        inputs=[state],
-        outputs=[state, layer_idx_dropdown],
-    )
-
-    # Dropdown change -> state sync
-    layer_idx_dropdown.change(
-        fn=on_layer_dropdown_change,
-        inputs=[layer_idx_dropdown, state],
-        outputs=[state],
-    )
-
     # Refine
-    refine_btn.click(
+    refine_button.click(
         fn=on_refine_click,
         inputs=[
-            seed,
-            randomize_seed,
-            prompt,
-            neg_prompt,
-            true_guidance_scale,
-            num_inference_steps,
-            cfg_norm,
-            use_en_prompt,
-            resolution,
-            gpu_duration,
+            history_state,
+            current_node_id_state,
+            refine_layer_index_state,
             sub_layers,
-            state,
-            history_dropdown,
-            layer_idx_dropdown,
+            gpu_duration,
         ],
         outputs=[
-            state,
+            history_state,
+            current_node_id_state,
+            gallery,
+            picker_gallery,
             history_dropdown,
-            history_dropdown,
-            base_gallery,
-            layer_picker,
-            layer_idx_dropdown,
-            layer_idx_dropdown,
-            refined_accordion,
+            refine_layer_index_state,
+            refine_layer_dropdown,
+            chips_md,
             refined_gallery,
-            branch_name,
+            export_file,
+            export_zip_file,
+            refined_accordion,  # ✅ uses gr.update(open=True)
         ],
     )
 
-    # Back to parent
-    back_parent_btn.click(
-        fn=on_back_to_parent_click,
-        inputs=[state, history_dropdown],
+    # History jump
+    history_dropdown.change(
+        fn=on_history_change,
+        inputs=[history_state, history_dropdown],
         outputs=[
-            state,
+            current_node_id_state,
+            gallery,
+            picker_gallery,
             history_dropdown,
-            base_gallery,
-            layer_picker,
-            layer_idx_dropdown,
-            layer_idx_dropdown,
-            refined_accordion,
+            refine_layer_index_state,
+            refine_layer_dropdown,
+            chips_md,
             refined_gallery,
-            branch_name,
+            export_file,
+            export_zip_file,
+            refined_accordion,  # ✅ uses gr.update(open=False)
         ],
     )
 
-    # Redo refine (same parent/index/sub_layers as the selected refine node)
-    redo_refine_btn.click(
-        fn=on_redo_refine_click,
-        inputs=[
-            seed,
-            randomize_seed,
-            prompt,
-            neg_prompt,
-            true_guidance_scale,
-            num_inference_steps,
-            cfg_norm,
-            use_en_prompt,
-            resolution,
-            gpu_duration,
-            state,
+    # Back to parent
+    back_button.click(
+        fn=on_back_to_parent,
+        inputs=[history_state, current_node_id_state],
+        outputs=[
+            current_node_id_state,
+            gallery,
+            picker_gallery,
             history_dropdown,
+            refine_layer_index_state,
+            refine_layer_dropdown,
+            chips_md,
+            refined_gallery,
+            export_file,
+            export_zip_file,
+            refined_accordion,
         ],
+    )
+
+    # Redo refine
+    redo_button.click(
+        fn=on_redo_refine,
+        inputs=[history_state, current_node_id_state, gpu_duration],
         outputs=[
-            state,
+            history_state,
+            current_node_id_state,
+            gallery,
+            picker_gallery,
             history_dropdown,
-            history_dropdown,
-            base_gallery,
-            layer_picker,
-            layer_idx_dropdown,
-            layer_idx_dropdown,
-            refined_accordion,
+            refine_layer_index_state,
+            refine_layer_dropdown,
+            chips_md,
             refined_gallery,
-            branch_name,
+            export_file,
+            export_zip_file,
+            refined_accordion,
         ],
     )
 
-    # Rename selected node
-    rename_btn.click(
-        fn=on_rename_node_click,
-        inputs=[state, history_dropdown, branch_name],
-        outputs=[state, history_dropdown, history_dropdown, branch_name],
+    # Duplicate node (branch)
+    dup_button.click(
+        fn=on_duplicate_node,
+        inputs=[history_state, current_node_id_state],
+        outputs=[
+            current_node_id_state,
+            gallery,
+            picker_gallery,
+            history_dropdown,
+            refine_layer_index_state,
+            refine_layer_dropdown,
+            chips_md,
+            refined_gallery,
+            export_file,
+            export_zip_file,
+            refined_accordion,
+        ],
     )
 
-    # Export selected node
-    export_pptx_btn.click(
-        fn=lambda st, node_id: on_export_click(st, node_id, "pptx"),
-        inputs=[state, history_dropdown],
-        outputs=[export_pptx_file],
+    # Rename
+    rename_button.click(
+        fn=on_rename_node,
+        inputs=[history_state, current_node_id_state, rename_text],
+        outputs=[history_dropdown],
     )
 
-    export_zip_btn.click(
-        fn=lambda st, node_id: on_export_click(st, node_id, "zip"),
-        inputs=[state, history_dropdown],
-        outputs=[export_zip_file],
+    # Export
+    export_button.click(
+        fn=on_export_current,
+        inputs=[history_state, current_node_id_state],
+        outputs=[export_file, export_zip_file],
     )
 
 if __name__ == "__main__":