Update app.py

app.py

@@ -1,13 +1,15 @@
 import os
-import uuid
+import io
+import gc
 import json
 import time
-import shutil
+import uuid
 import numpy as np
 import random
 import tempfile
 import zipfile
-from typing import Any, Dict, List, Optional, Tuple
+from dataclasses import dataclass
+from typing import Dict, Any, List, Optional, Tuple
 
 import spaces
 import torch
@@ -17,14 +19,23 @@ from PIL import Image
 from diffusers import QwenImageLayeredPipeline
 from pptx import Presentation
 
+from huggingface_hub import login, HfApi, hf_hub_download
+
+
 LOG_DIR = "/tmp/local"
 MAX_SEED = np.iinfo(np.int32).max
 
 # Optional HF login (works in Spaces if you set HF token as secret env var "hf")
-from huggingface_hub import login, HfApi, hf_hub_download
-
 login(token=os.environ.get("hf"))
 
+# Dataset persistence (optional). Example: "username/qwen-layered-sessions"
+DATASET_REPO = os.environ.get("DATASET_REPO", "").strip()
+DATASET_BRANCH = os.environ.get("DATASET_BRANCH", "main").strip()
+
+# If you want to reduce allocator weirdness on some CUDA envs, you can set this.
+# Keep as "setdefault" so you can override in Space Variables.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", os.environ.get("PYTORCH_CUDA_ALLOC_CONF", ""))
+
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -32,56 +43,26 @@ pipeline = QwenImageLayeredPipeline.from_pretrained(
     "Qwen/Qwen-Image-Layered", torch_dtype=dtype
 ).to(device)
 
-# ----------------------------
-# Dataset repo persistence (no /data needed)
-# ----------------------------
-HF_TOKEN = os.environ.get("hf")  # secret
-HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO")  # e.g. "hexware/qwen-layered-sessions"
-_hf_api: Optional[HfApi] = None
-_persist_enabled = False
-
-def _init_dataset_repo() -> Tuple[bool, str]:
-    """
-    Returns (enabled, message)
-    """
-    global _hf_api, _persist_enabled
-    if not HF_TOKEN:
-        _persist_enabled = False
-        return False, "Persistence: disabled (no secret env var 'hf')."
-    if not HF_DATASET_REPO:
-        _persist_enabled = False
-        return False, "Persistence: disabled (set env var HF_DATASET_REPO to enable)."
+pipeline.set_progress_bar_config(disable=True)
 
-    try:
-        _hf_api = HfApi(token=HF_TOKEN)
-        # Create dataset repo if missing (private). If exists, this is no-op.
-        # NOTE: create_repo is available via HfApi in most versions.
-        _hf_api.create_repo(
-            repo_id=HF_DATASET_REPO,
-            repo_type="dataset",
-            private=True,
-            exist_ok=True,
-        )
-        _persist_enabled = True
-        return True, f"Persistence: enabled (dataset repo: {HF_DATASET_REPO})."
-    except Exception as e:
-        _persist_enabled = False
-        return False, f"Persistence: failed to init dataset repo: {type(e).__name__}: {e}"
+try:
+    torch.backends.cuda.matmul.allow_tf32 = True
+except Exception:
+    pass
 
-_enabled, _enabled_msg = _init_dataset_repo()
 
-# ----------------------------
-# Helpers
-# ----------------------------
 def ensure_dirname(path: str):
     if path and not os.path.exists(path):
         os.makedirs(path, exist_ok=True)
 
+
 def random_str(length=8):
     return uuid.uuid4().hex[:length]
 
-def _now_ts() -> float:
-    return time.time()
+
+def _now_ts() -> int:
+    return int(time.time())
+
 
 def _clamp_int(x, default: int, lo: int, hi: int) -> int:
     try:
@@ -90,7 +71,9 @@ def _clamp_int(x, default: int, lo: int, hi: int) -> int:
         v = default
     return max(lo, min(hi, v))
 
-def _pil_rgba(input_image) -> Image.Image:
+
+def _safe_img_rgba(input_image):
+    # Normalize image input
     if isinstance(input_image, list):
         input_image = input_image[0]
 
@@ -102,9 +85,11 @@ def _pil_rgba(input_image) -> Image.Image:
         pil_image = Image.fromarray(input_image).convert("RGB").convert("RGBA")
     else:
         raise ValueError(f"Unsupported input_image type: {type(input_image)}")
+
     return pil_image
 
-def imagelist_to_pptx(img_files: List[str]) -> str:
+
+def imagelist_to_pptx(img_files):
     with Image.open(img_files[0]) as img:
         img_width_px, img_height_px = img.size
 
@@ -133,29 +118,288 @@ def imagelist_to_pptx(img_files: List[str]) -> str:
         prs.save(tmp.name)
         return tmp.name
 
-def export_node_layers(layers: List[Image.Image]) -> Tuple[str, str]:
-    """
-    Returns (pptx_path, zip_path)
-    """
-    temp_files: List[str] = []
+
+def _write_layers_to_temp_pngs(layers: List[Image.Image]) -> List[str]:
+    temp_files = []
     for img in layers:
         tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
         img.save(tmp.name)
         temp_files.append(tmp.name)
+    return temp_files
 
-    pptx_path = imagelist_to_pptx(temp_files)
 
+def _build_zip_from_pngs(png_paths: List[str]) -> str:
     with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmpzip:
         with zipfile.ZipFile(tmpzip.name, "w", zipfile.ZIP_DEFLATED) as zipf:
-            for i, img_path in enumerate(temp_files):
+            for i, img_path in enumerate(png_paths):
                 zipf.write(img_path, f"layer_{i+1}.png")
-        zip_path = tmpzip.name
+        return tmpzip.name
+
+
+def _mk_node_name(kind: str, depth: int) -> str:
+    if kind == "root":
+        return "Root"
+    return f"Refine d{depth}"
+
+
+def _ds_enabled() -> bool:
+    return bool(DATASET_REPO) and bool(os.environ.get("hf"))
+
+
+def _ds_api() -> HfApi:
+    return HfApi(token=os.environ.get("hf"))
+
+
+def _ds_path(*parts: str) -> str:
+    return "/".join([p.strip("/") for p in parts if p is not None and p != ""])
+
+
+def ds_list_sessions() -> List[str]:
+    if not _ds_enabled():
+        return []
+    api = _ds_api()
+    files = api.list_repo_files(repo_id=DATASET_REPO, repo_type="dataset", revision=DATASET_BRANCH)
+    sessions = set()
+    for f in files:
+        if f.startswith("sessions/") and f.endswith("/index.json"):
+            # sessions/<sid>/index.json
+            parts = f.split("/")
+            if len(parts) >= 3:
+                sessions.add(parts[1])
+    return sorted(list(sessions))
+
+
+def ds_upload_bytes(path_in_repo: str, data: bytes, commit_message: str):
+    api = _ds_api()
+    with tempfile.NamedTemporaryFile(delete=False) as tmp:
+        tmp.write(data)
+        tmp.flush()
+        api.upload_file(
+            path_or_fileobj=tmp.name,
+            path_in_repo=path_in_repo,
+            repo_id=DATASET_REPO,
+            repo_type="dataset",
+            revision=DATASET_BRANCH,
+            commit_message=commit_message,
+        )
+
+
+def ds_upload_file(local_path: str, path_in_repo: str, commit_message: str):
+    api = _ds_api()
+    api.upload_file(
+        path_or_fileobj=local_path,
+        path_in_repo=path_in_repo,
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        revision=DATASET_BRANCH,
+        commit_message=commit_message,
+    )
+
+
+def ds_download_json(path_in_repo: str) -> Dict[str, Any]:
+    local = hf_hub_download(
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        revision=DATASET_BRANCH,
+        filename=path_in_repo,
+    )
+    with open(local, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def ds_download_file(path_in_repo: str) -> str:
+    return hf_hub_download(
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        revision=DATASET_BRANCH,
+        filename=path_in_repo,
+    )
 
-    return pptx_path, zip_path
 
-# ----------------------------
-# ZeroGPU duration
-# ----------------------------
+@dataclass
+class Node:
+    node_id: str
+    parent_id: Optional[str]
+    kind: str                  # "root" or "refine"
+    depth: int
+    name: str
+    created_at: int
+    params: Dict[str, Any]
+    refine_meta: Dict[str, Any]
+
+    layers: List[Image.Image]  # PIL layers (in-memory)
+    png_paths: List[str]       # local temp pngs
+    pptx_path: Optional[str]
+    zip_path: Optional[str]
+
+    children: List[str]
+
+
+def new_state() -> Dict[str, Any]:
+    return {
+        "session_id": random_str(12),
+        "nodes": {},              # node_id -> Node
+        "root_id": None,
+        "current_id": None,
+        "selected_layer_idx": 0,
+        "autosave": False,
+        "dataset_repo": DATASET_REPO,
+        "dataset_branch": DATASET_BRANCH,
+        "last_source_for_redo": None,  # (from_node_id, layer_idx, sub_layers, params)
+    }
+
+
+def _node_to_brief(n: Node) -> str:
+    short = n.node_id[:8]
+    return f"{n.name} · {short} · {n.kind} · depth {n.depth}"
+
+
+def _history_choices(state: Dict[str, Any]) -> List[Tuple[str, str]]:
+    nodes: Dict[str, Node] = state["nodes"]
+    if not nodes:
+        return []
+    # sort by created_at
+    ordered = sorted(nodes.values(), key=lambda x: x.created_at)
+    return [(_node_to_brief(n), n.node_id) for n in ordered]
+
+
+def _chips_html(state: Dict[str, Any], node_id: Optional[str]) -> str:
+    if not node_id or node_id not in state["nodes"]:
+        return ""
+    n: Node = state["nodes"][node_id]
+    root_id = state.get("root_id")
+    parent = n.parent_id
+    children = len(n.children or [])
+    chips = []
+    if n.node_id == root_id:
+        chips.append("[root]")
+    if parent:
+        chips.append("[parent]")
+    chips.append(f"[children:{children}]")
+    return " ".join(chips)
+
+
+def _ensure_exports(node: Node) -> Node:
+    if node.png_paths is None or len(node.png_paths) == 0:
+        node.png_paths = _write_layers_to_temp_pngs(node.layers)
+    if not node.pptx_path:
+        node.pptx_path = imagelist_to_pptx(node.png_paths)
+    if not node.zip_path:
+        node.zip_path = _build_zip_from_pngs(node.png_paths)
+    return node
+
+
+def _persist_node_to_dataset(state: Dict[str, Any], node: Node):
+    if not _ds_enabled():
+        return
+
+    sid = state["session_id"]
+    nid = node.node_id
+    base = _ds_path("sessions", sid, "nodes", nid)
+
+    # node meta
+    meta = {
+        "node_id": node.node_id,
+        "parent_id": node.parent_id,
+        "kind": node.kind,
+        "depth": node.depth,
+        "name": node.name,
+        "created_at": node.created_at,
+        "params": node.params,
+        "refine_meta": node.refine_meta,
+        "children": node.children,
+        "layer_count": len(node.layers),
+    }
+    ds_upload_bytes(_ds_path(base, "node.json"), json.dumps(meta, ensure_ascii=False, indent=2).encode("utf-8"),
+                    commit_message=f"save node meta {sid}/{nid}")
+
+    # ensure exports
+    node = _ensure_exports(node)
+
+    # pngs
+    for i, p in enumerate(node.png_paths):
+        ds_upload_file(p, _ds_path(base, f"layer_{i+1}.png"), commit_message=f"save layer png {sid}/{nid}")
+
+    # pptx/zip
+    ds_upload_file(node.pptx_path, _ds_path(base, "layers.pptx"), commit_message=f"save pptx {sid}/{nid}")
+    ds_upload_file(node.zip_path, _ds_path(base, "layers.zip"), commit_message=f"save zip {sid}/{nid}")
+
+    # session index
+    nodes: Dict[str, Node] = state["nodes"]
+    index = {
+        "session_id": sid,
+        "saved_at": _now_ts(),
+        "root_id": state["root_id"],
+        "current_id": state["current_id"],
+        "nodes": [
+            {
+                "node_id": x.node_id,
+                "parent_id": x.parent_id,
+                "kind": x.kind,
+                "depth": x.depth,
+                "name": x.name,
+                "created_at": x.created_at,
+                "layer_count": len(x.layers),
+            }
+            for x in sorted(nodes.values(), key=lambda z: z.created_at)
+        ],
+    }
+    ds_upload_bytes(_ds_path("sessions", sid, "index.json"),
+                    json.dumps(index, ensure_ascii=False, indent=2).encode("utf-8"),
+                    commit_message=f"save session index {sid}")
+
+
+def _load_session_from_dataset(session_id: str) -> Dict[str, Any]:
+    st = new_state()
+    st["session_id"] = session_id
+
+    index = ds_download_json(_ds_path("sessions", session_id, "index.json"))
+    st["root_id"] = index.get("root_id")
+    st["current_id"] = index.get("current_id")
+
+    nodes: Dict[str, Node] = {}
+
+    for item in index.get("nodes", []):
+        nid = item["node_id"]
+        meta = ds_download_json(_ds_path("sessions", session_id, "nodes", nid, "node.json"))
+
+        layer_count = int(meta.get("layer_count", 0))
+        layers = []
+        png_paths = []
+        for i in range(layer_count):
+            fp = ds_download_file(_ds_path("sessions", session_id, "nodes", nid, f"layer_{i+1}.png"))
+            layers.append(Image.open(fp).convert("RGBA"))
+            png_paths.append(fp)
+
+        pptx_path = ds_download_file(_ds_path("sessions", session_id, "nodes", nid, "layers.pptx"))
+        zip_path = ds_download_file(_ds_path("sessions", session_id, "nodes", nid, "layers.zip"))
+
+        n = Node(
+            node_id=nid,
+            parent_id=meta.get("parent_id"),
+            kind=meta.get("kind", "root"),
+            depth=int(meta.get("depth", 0)),
+            name=meta.get("name", _mk_node_name(meta.get("kind", "root"), int(meta.get("depth", 0)))),
+            created_at=int(meta.get("created_at", _now_ts())),
+            params=meta.get("params", {}),
+            refine_meta=meta.get("refine_meta", {}),
+            layers=layers,
+            png_paths=png_paths,
+            pptx_path=pptx_path,
+            zip_path=zip_path,
+            children=meta.get("children", []),
+        )
+        nodes[nid] = n
+
+    st["nodes"] = nodes
+
+    if st["current_id"] is None and st["root_id"] in nodes:
+        st["current_id"] = st["root_id"]
+
+    return st
+
+
+# Dynamic duration callable: must accept the same args as GPU function. It returns seconds.
 def get_duration(
     input_image,
     seed=777,
@@ -172,24 +416,22 @@ def get_duration(
 ):
     return _clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
 
-# ----------------------------
-# GPU pipeline runners
-# ----------------------------
+
 @spaces.GPU(duration=get_duration)
 def gpu_run_pipeline(
-    input_pil_image: Image.Image,
-    seed=777,
-    randomize_seed=False,
-    prompt=None,
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    layer=4,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=640,
-    gpu_duration=1000,
-) -> List[Image.Image]:
+    pil_image_rgba: Image.Image,
+    seed: int,
+    randomize_seed: bool,
+    prompt: str,
+    neg_prompt: str,
+    true_guidance_scale: float,
+    num_inference_steps: int,
+    layer: int,
+    cfg_norm: bool,
+    use_en_prompt: bool,
+    resolution: int,
+    gpu_duration: int,
+):
     # Seed
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -202,7 +444,7 @@ def gpu_run_pipeline(
     gen_device = "cuda" if torch.cuda.is_available() else "cpu"
 
     inputs = {
-        "image": input_pil_image,
+        "image": pil_image_rgba,
         "generator": torch.Generator(device=gen_device).manual_seed(seed),
         "true_cfg_scale": true_guidance_scale,
         "prompt": prompt,
@@ -215,510 +457,553 @@ def gpu_run_pipeline(
         "use_en_prompt": use_en_prompt,
     }
 
-    with torch.inference_mode():
-        out = pipeline(**inputs)
+    print("INFER INPUTS:", {k: (str(v)[:200] if isinstance(v, str) else v) for k, v in inputs.items()})
+    print("REQUESTED GPU DURATION:", gpu_duration)
 
-    # out.images[0] => list of PIL images
-    return out.images[0]
+    # Self-heal retry for rare CUDA/NVML allocator glitches on some envs
+    try:
+        with torch.inference_mode():
+            out = pipeline(**inputs)
+            output_images = out.images[0]  # list of PIL images (layers)
+    except RuntimeError as e:
+        msg = str(e)
+        if "NVML_SUCCESS" in msg or "CUDACachingAllocator" in msg:
+            print("Caught allocator/NVML error, retrying once after cache cleanup:", msg)
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+            gc.collect()
+            time.sleep(0.2)
+            with torch.inference_mode():
+                out = pipeline(**inputs)
+                output_images = out.images[0]
+        else:
+            raise
 
-# ----------------------------
-# Session / History model
-# ----------------------------
-def new_state() -> Dict[str, Any]:
-    sid = uuid.uuid4().hex
-    return {
-        "session_id": sid,
-        "nodes": {},          # node_id -> node dict
-        "root_id": None,
-        "current_id": None,
-        "selected_layer_idx": 0,
-        "last_refined_id": None,
-    }
+    # Ensure RGBA
+    fixed = []
+    for im in output_images:
+        if not isinstance(im, Image.Image):
+            im = Image.fromarray(np.array(im))
+        fixed.append(im.convert("RGBA"))
 
-def _node_label(node: Dict[str, Any]) -> str:
-    name = node.get("name") or node["id"][:8]
-    return f"{name} ({node['id'][:8]})"
-
-def _build_history_choices(st: Dict[str, Any]) -> List[Tuple[str, str]]:
-    # returns list of (label, value=node_id)
-    out = []
-    for nid, node in st["nodes"].items():
-        out.append((_node_label(node), nid))
-    # stable order by created
-    out.sort(key=lambda x: st["nodes"][x[1]].get("created_at", 0.0))
-    return out
-
-def _get_node(st: Dict[str, Any], node_id: Optional[str]) -> Optional[Dict[str, Any]]:
-    if not node_id:
-        return None
-    return st["nodes"].get(node_id)
-
-def _current_node(st: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-    return _get_node(st, st.get("current_id"))
-
-def _chips_text(st: Dict[str, Any], node_id: Optional[str]) -> str:
-    node = _get_node(st, node_id)
-    if not node:
-        return ""
-    chips = []
-    if node_id == st.get("root_id"):
-        chips.append("[root]")
-    if node.get("parent_id"):
-        chips.append("[parent]")
-    children = node.get("children_ids") or []
-    if children:
-        chips.append(f"[children:{len(children)}]")
-    return " ".join(chips)
+    try:
+        torch.cuda.empty_cache()
+    except Exception:
+        pass
 
-def _make_node(
-    st: Dict[str, Any],
+    return fixed, seed
+
+
+def _create_node(
+    state: Dict[str, Any],
     layers: List[Image.Image],
+    kind: str,
     parent_id: Optional[str],
-    name: Optional[str] = None,
-) -> str:
-    nid = uuid.uuid4().hex
-    node = {
-        "id": nid,
-        "name": name or ("root" if parent_id is None else "refine"),
-        "parent_id": parent_id,
-        "children_ids": [],
-        "created_at": _now_ts(),
-        "layers": layers,
-    }
-    st["nodes"][nid] = node
-    if parent_id:
-        parent = st["nodes"].get(parent_id)
-        if parent is not None:
-            parent.setdefault("children_ids", []).append(nid)
-    return nid
-
-def _set_current(st: Dict[str, Any], node_id: str):
-    st["current_id"] = node_id
-    st["selected_layer_idx"] = 0
-
-# ----------------------------
-# Persistence: save/load whole session as one zip in dataset repo
-# ----------------------------
-def _serialize_session_to_zip(st: Dict[str, Any]) -> str:
-    """
-    Create a zip file with:
-      session.json
-      nodes/<node_id>/layer_1.png ...
-    Returns local zip path.
-    """
-    tmpdir = tempfile.mkdtemp(prefix="sess_")
-    try:
-        sess_meta = {
-            "session_id": st["session_id"],
-            "root_id": st["root_id"],
-            "current_id": st["current_id"],
-            "selected_layer_idx": st.get("selected_layer_idx", 0),
-            "last_refined_id": st.get("last_refined_id"),
-            "nodes": {},
-        }
-
-        for nid, node in st["nodes"].items():
-            node_dir = os.path.join(tmpdir, "nodes", nid)
-            os.makedirs(node_dir, exist_ok=True)
-            layers: List[Image.Image] = node.get("layers") or []
-            for i, img in enumerate(layers):
-                img_path = os.path.join(node_dir, f"layer_{i+1}.png")
-                img.save(img_path)
-            sess_meta["nodes"][nid] = {
-                "id": nid,
-                "name": node.get("name"),
-                "parent_id": node.get("parent_id"),
-                "children_ids": node.get("children_ids") or [],
-                "created_at": node.get("created_at", 0.0),
-                "layer_count": len(layers),
-            }
-
-        meta_path = os.path.join(tmpdir, "session.json")
-        with open(meta_path, "w", encoding="utf-8") as f:
-            json.dump(sess_meta, f, ensure_ascii=False, indent=2)
-
-        out_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name
-        with zipfile.ZipFile(out_zip, "w", zipfile.ZIP_DEFLATED) as zf:
-            for root, _, files in os.walk(tmpdir):
-                for fn in files:
-                    abs_path = os.path.join(root, fn)
-                    rel_path = os.path.relpath(abs_path, tmpdir)
-                    zf.write(abs_path, rel_path)
-        return out_zip
-    finally:
-        shutil.rmtree(tmpdir, ignore_errors=True)
-
-def _deserialize_session_from_zip(zip_path: str) -> Dict[str, Any]:
-    tmpdir = tempfile.mkdtemp(prefix="sess_load_")
-    try:
-        with zipfile.ZipFile(zip_path, "r") as zf:
-            zf.extractall(tmpdir)
-
-        meta_path = os.path.join(tmpdir, "session.json")
-        with open(meta_path, "r", encoding="utf-8") as f:
-            meta = json.load(f)
-
-        st = new_state()
-        st["session_id"] = meta["session_id"]
-        st["root_id"] = meta.get("root_id")
-        st["current_id"] = meta.get("current_id")
-        st["selected_layer_idx"] = meta.get("selected_layer_idx", 0)
-        st["last_refined_id"] = meta.get("last_refined_id")
-
-        nodes_meta: Dict[str, Any] = meta.get("nodes", {})
-        # First pass: create node shells
-        for nid, nm in nodes_meta.items():
-            st["nodes"][nid] = {
-                "id": nid,
-                "name": nm.get("name"),
-                "parent_id": nm.get("parent_id"),
-                "children_ids": nm.get("children_ids") or [],
-                "created_at": nm.get("created_at", 0.0),
-                "layers": [],
-            }
-        # Second pass: load layers images
-        for nid, nm in nodes_meta.items():
-            layer_count = int(nm.get("layer_count", 0))
-            node_dir = os.path.join(tmpdir, "nodes", nid)
-            layers: List[Image.Image] = []
-            for i in range(layer_count):
-                p = os.path.join(node_dir, f"layer_{i+1}.png")
-                if os.path.exists(p):
-                    layers.append(Image.open(p).convert("RGBA"))
-            st["nodes"][nid]["layers"] = layers
-
-        return st
-    finally:
-        shutil.rmtree(tmpdir, ignore_errors=True)
-
-def save_session_to_hub(st: Dict[str, Any]) -> Tuple[str, str]:
-    """
-    Returns (status_text, session_id)
-    """
-    if not _persist_enabled or _hf_api is None:
-        return "Save: disabled (set HF_DATASET_REPO and secret hf write token).", st.get("session_id", "")
-    try:
-        zip_path = _serialize_session_to_zip(st)
-        path_in_repo = f"sessions/{st['session_id']}.zip"
-        _hf_api.upload_file(
-            path_or_fileobj=zip_path,
-            path_in_repo=path_in_repo,
-            repo_id=HF_DATASET_REPO,
-            repo_type="dataset",
-            commit_message=f"Save session {st['session_id']}",
-        )
-        return f"Saved to dataset repo: {path_in_repo}", st["session_id"]
-    except Exception as e:
-        return f"Save failed: {type(e).__name__}: {e}", st.get("session_id", "")
-    finally:
-        try:
-            if "zip_path" in locals() and os.path.exists(zip_path):
-                os.remove(zip_path)
-        except Exception:
-            pass
-
-def load_session_from_hub(session_id: str) -> Tuple[Optional[Dict[str, Any]], str]:
-    if not _persist_enabled:
-        return None, "Load: disabled (set HF_DATASET_REPO and secret hf write token)."
-    session_id = (session_id or "").strip()
-    if not session_id:
-        return None, "Load: please enter a Session ID."
-    try:
-        filename = f"sessions/{session_id}.zip"
-        local_zip = hf_hub_download(
-            repo_id=HF_DATASET_REPO,
-            repo_type="dataset",
-            filename=filename,
-            token=HF_TOKEN,
+    params: Dict[str, Any],
+    refine_meta: Dict[str, Any],
+) -> Node:
+    nid = random_str(24)
+    depth = 0
+    if parent_id and parent_id in state["nodes"]:
+        depth = state["nodes"][parent_id].depth + 1
+
+    n = Node(
+        node_id=nid,
+        parent_id=parent_id,
+        kind=kind,
+        depth=depth,
+        name=_mk_node_name(kind, depth),
+        created_at=_now_ts(),
+        params=params,
+        refine_meta=refine_meta,
+        layers=layers,
+        png_paths=[],
+        pptx_path=None,
+        zip_path=None,
+        children=[],
+    )
+    state["nodes"][nid] = n
+    if parent_id and parent_id in state["nodes"]:
+        state["nodes"][parent_id].children.append(nid)
+    if kind == "root" and state.get("root_id") is None:
+        state["root_id"] = nid
+    state["current_id"] = nid
+    state["selected_layer_idx"] = 0
+    return n
+
+
+def _current_node(state: Dict[str, Any]) -> Optional[Node]:
+    cid = state.get("current_id")
+    if cid and cid in state["nodes"]:
+        return state["nodes"][cid]
+    return None
+
+
+def _render_current(state: Dict[str, Any]):
+    n = _current_node(state)
+    if not n:
+        return (
+            [],  # main gallery
+            [],  # layer strip
+            gr.update(choices=[], value=None),  # layer dropdown
+            gr.update(choices=_history_choices(state), value=None),  # history dropdown
+            "",  # chips
+            None,  # pptx
+            None,  # zip
+            gr.update(open=False),  # refined accordion
+            [],  # refined gallery
         )
-        st = _deserialize_session_from_zip(local_zip)
-        return st, f"Loaded session: {session_id}"
-    except Exception as e:
-        return None, f"Load failed: {type(e).__name__}: {e}"
-
-# ----------------------------
-# UI Callbacks
-# ----------------------------
-def ui_boot() -> Tuple[str, Dict[str, Any]]:
-    ensure_dirname(LOG_DIR)
-    st = new_state()
-    return _enabled_msg, st
 
-def on_new_session(st: Dict[str, Any]) -> Tuple[Dict[str, Any], str, gr.Dropdown, List[Image.Image], List[Image.Image], str, str, str, Optional[str], Optional[str]]:
-    st = new_state()
+    # layer dropdown
+    layer_choices = [(f"Layer {i+1}", str(i)) for i in range(len(n.layers))]
+    dd = gr.update(choices=layer_choices, value=str(min(state.get("selected_layer_idx", 0), max(0, len(n.layers) - 1))))
+
+    # ensure exports available for selected node
+    n = _ensure_exports(n)
+
     return (
-        st,
-        st["session_id"],
-        gr.Dropdown(choices=[], value=None),
+        n.layers,                 # main gallery
+        n.layers,                 # strip gallery
+        dd,                       # layer dropdown
+        gr.update(choices=_history_choices(state), value=n.node_id),
+        _chips_html(state, n.node_id),
+        n.pptx_path,
+        n.zip_path,
+        gr.update(open=False),
         [],
-        [],
-        "",
-        "",
-        "",
-        None,
-        None,
     )
 
-def _render_from_state(st: Dict[str, Any]) -> Tuple[
-    gr.Dropdown,
-    List[Image.Image],
-    List[Image.Image],
-    gr.Number,
-    str
-]:
-    choices = _build_history_choices(st)
-    current = _current_node(st)
-    layers = current["layers"] if current else []
-    idx = st.get("selected_layer_idx", 0)
-    if layers:
-        idx = max(0, min(idx, len(layers) - 1))
-    st["selected_layer_idx"] = idx
-    chips = _chips_text(st, st.get("current_id"))
+
+def on_apply_fast_profile():
+    # Does not change defaults; only sets UI values when clicked.
     return (
-        gr.Dropdown(choices=choices, value=st.get("current_id")),
-        layers,
-        layers,  # mini gallery mirrors current layers
-        idx,
-        chips,
+        30,   # steps
+        5,    # layers
+        640,  # resolution
     )
 
-def on_history_select(node_id: str, st: Dict[str, Any]) -> Tuple[Dict[str, Any], gr.Dropdown, List[Image.Image], List[Image.Image], gr.Number, str]:
-    if node_id and node_id in st["nodes"]:
-        st["current_id"] = node_id
-        st["selected_layer_idx"] = 0
-    dd, layers, mini, idx, chips = _render_from_state(st)
-    return st, dd, layers, mini, idx, chips
 
-def on_layer_gallery_select(evt: gr.SelectData, st: Dict[str, Any]) -> Tuple[Dict[str, Any], gr.Number]:
-    # evt.index is int for Gallery
-    idx = int(evt.index) if evt and evt.index is not None else 0
-    current = _current_node(st)
-    if current:
-        layers = current.get("layers") or []
-        if layers:
-            idx = max(0, min(idx, len(layers) - 1))
-        else:
-            idx = 0
-    else:
-        idx = 0
-    st["selected_layer_idx"] = idx
-    return st, idx
-
-def on_back_to_parent(st: Dict[str, Any]) -> Tuple[Dict[str, Any], gr.Dropdown, List[Image.Image], List[Image.Image], gr.Number, str]:
-    cur = _current_node(st)
-    if cur and cur.get("parent_id"):
-        st["current_id"] = cur["parent_id"]
-        st["selected_layer_idx"] = 0
-    dd, layers, mini, idx, chips = _render_from_state(st)
-    return st, dd, layers, mini, idx, chips
-
-def on_duplicate_node(st: Dict[str, Any]) -> Tuple[Dict[str, Any], gr.Dropdown, List[Image.Image], List[Image.Image], gr.Number, str]:
-    cur = _current_node(st)
-    if cur:
-        # Duplicate current node as sibling (same parent)
-        layers = cur.get("layers") or []
-        parent_id = cur.get("parent_id")
-        name = (cur.get("name") or "node") + " copy"
-        new_id = _make_node(st, layers=layers, parent_id=parent_id, name=name)
-        _set_current(st, new_id)
-        if st.get("root_id") is None and parent_id is None:
-            st["root_id"] = new_id
-    dd, layers, mini, idx, chips = _render_from_state(st)
-    return st, dd, layers, mini, idx, chips
-
-def on_rename_node(new_name: str, st: Dict[str, Any]) -> Tuple[Dict[str, Any], gr.Dropdown]:
-    cur = _current_node(st)
-    if cur:
-        nn = (new_name or "").strip()
-        if nn:
-            cur["name"] = nn
-    dd, _, _, _, _ = _render_from_state(st)
-    return st, dd
-
-def on_export_selected(st: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]:
-    cur = _current_node(st)
-    if not cur:
-        return None, None
-    pptx_path, zip_path = export_node_layers(cur.get("layers") or [])
-    return pptx_path, zip_path
-
-def on_save_session(st: Dict[str, Any]) -> Tuple[str, str]:
-    status, sid = save_session_to_hub(st)
-    return status, sid
-
-def on_load_session(session_id: str, st: Dict[str, Any]) -> Tuple[
-    Dict[str, Any],
-    str,
-    gr.Dropdown,
-    List[Image.Image],
-    List[Image.Image],
-    gr.Number,
-    str
-]:
-    loaded, msg = load_session_from_hub(session_id)
-    if loaded is None:
-        dd, layers, mini, idx, chips = _render_from_state(st)
-        return st, msg, dd, layers, mini, idx, chips
-
-    st = loaded
-    dd, layers, mini, idx, chips = _render_from_state(st)
-    return st, msg, dd, layers, mini, idx, chips
-
-# GPU click handlers
+def on_toggle_autosave(val, state):
+    state["autosave"] = bool(val)
+    return state
+
+
 def on_decompose_click(
     input_image,
-    seed=777,
-    randomize_seed=False,
-    prompt=None,
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    layer=4,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=640,
-    gpu_duration=1000,
-    st: Optional[Dict[str, Any]] = None,
+    seed,
+    randomize_seed,
+    prompt,
+    neg_prompt,
+    true_guidance_scale,
+    num_inference_steps,
+    layer,
+    cfg_norm,
+    use_en_prompt,
+    resolution,
+    gpu_duration,
+    state,
 ):
-    if st is None:
-        st = new_state()
-
-    pil_image = _pil_rgba(input_image)
-    layers_out = gpu_run_pipeline(
-        pil_image,
-        seed=seed,
-        randomize_seed=randomize_seed,
+    if state is None:
+        state = new_state()
+
+    pil = _safe_img_rgba(input_image)
+
+    layers_out, used_seed = gpu_run_pipeline(
+        pil_image_rgba=pil,
+        seed=int(seed),
+        randomize_seed=bool(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,
+        true_guidance_scale=float(true_guidance_scale),
+        num_inference_steps=int(num_inference_steps),
+        layer=int(layer),
+        cfg_norm=bool(cfg_norm),
+        use_en_prompt=bool(use_en_prompt),
+        resolution=int(resolution),
+        gpu_duration=int(gpu_duration),
     )
 
-    # Reset session tree on Decompose (new root)
-    sid = st.get("session_id") or uuid.uuid4().hex
-    st = new_state()
-    st["session_id"] = sid
+    params = {
+        "seed": int(used_seed),
+        "randomize_seed": bool(randomize_seed),
+        "prompt": prompt,
+        "neg_prompt": neg_prompt,
+        "true_guidance_scale": float(true_guidance_scale),
+        "num_inference_steps": int(num_inference_steps),
+        "layers": int(layer),
+        "cfg_norm": bool(cfg_norm),
+        "use_en_prompt": bool(use_en_prompt),
+        "resolution": int(resolution),
+        "gpu_duration": int(gpu_duration),
+    }
 
-    root_id = _make_node(st, layers=layers_out, parent_id=None, name="root")
-    st["root_id"] = root_id
-    _set_current(st, root_id)
+    node = _create_node(
+        state=state,
+        layers=layers_out,
+        kind="root" if state.get("root_id") is None else "refine",
+        parent_id=None,
+        params=params,
+        refine_meta={"mode": "decompose"},
+    )
 
-    dd, layers, mini, idx, chips = _render_from_state(st)
+    if state.get("autosave"):
+        _persist_node_to_dataset(state, node)
+
+    # reset refined UI
+    main_gallery, strip, layer_dd, history_dd, chips, pptx, zzip, acc, refined = _render_current(state)
     return (
-        st,
-        dd,
-        layers,
-        mini,
-        idx,
+        state,
+        main_gallery,
+        strip,
+        layer_dd,
+        history_dd,
         chips,
-        gr.update(open=False),   # refined accordion closed
-        [],                      # refined gallery cleared
-        None,
-        None,
+        pptx,
+        zzip,
+        acc,
+        refined,
+        used_seed,
     )
 
+
+def on_layer_select_from_strip(evt: gr.SelectData, state):
+    # evt.index -> int
+    if state is None:
+        state = new_state()
+    idx = int(evt.index) if evt and evt.index is not None else 0
+    state["selected_layer_idx"] = idx
+    n = _current_node(state)
+    if not n:
+        return state, gr.update(), ""
+    dd = gr.update(value=str(min(idx, len(n.layers) - 1)))
+    return state, dd, f"Selected: Layer {idx+1}"
+
+
+def on_layer_dropdown_change(val, state):
+    if state is None:
+        state = new_state()
+    try:
+        idx = int(val)
+    except Exception:
+        idx = 0
+    state["selected_layer_idx"] = idx
+    return state, f"Selected: Layer {idx+1}"
+
+
 def on_refine_click(
-    sub_layers_count: int,
-    seed=777,
-    randomize_seed=False,
-    prompt=None,
-    neg_prompt=" ",
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    cfg_norm=True,
-    use_en_prompt=True,
-    resolution=640,
-    gpu_duration=1000,
-    st: Optional[Dict[str, Any]] = None,
+    refine_sub_layers,
+    seed,
+    randomize_seed,
+    prompt,
+    neg_prompt,
+    true_guidance_scale,
+    num_inference_steps,
+    cfg_norm,
+    use_en_prompt,
+    resolution,
+    gpu_duration,
+    state,
 ):
-    if st is None:
-        st = new_state()
-
-    cur = _current_node(st)
-    if not cur:
-        dd, layers, mini, idx, chips = _render_from_state(st)
+    if state is None:
+        state = new_state()
+    n = _current_node(state)
+    if not n:
         return (
-            st,
-            dd,
-            layers,
-            mini,
-            idx,
-            chips,
-            "Refine: no current node.",
-            gr.update(open=False),
+            state,
+            gr.update(open=True),
             [],
+            gr.update(choices=_history_choices(state)),
+            "",
             None,
             None,
         )
 
-    layers_list: List[Image.Image] = cur.get("layers") or []
-    if not layers_list:
-        dd, layers, mini, idx, chips = _render_from_state(st)
+    idx = int(state.get("selected_layer_idx", 0))
+    idx = max(0, min(idx, len(n.layers) - 1))
+    selected = n.layers[idx]
+
+    # refine creates new node under current as parent
+    layers_out, used_seed = gpu_run_pipeline(
+        pil_image_rgba=selected.convert("RGBA"),
+        seed=int(seed),
+        randomize_seed=bool(randomize_seed),
+        prompt=prompt,
+        neg_prompt=neg_prompt,
+        true_guidance_scale=float(true_guidance_scale),
+        num_inference_steps=int(num_inference_steps),
+        layer=int(refine_sub_layers),
+        cfg_norm=bool(cfg_norm),
+        use_en_prompt=bool(use_en_prompt),
+        resolution=int(resolution),
+        gpu_duration=int(gpu_duration),
+    )
+
+    params = {
+        "seed": int(used_seed),
+        "randomize_seed": bool(randomize_seed),
+        "prompt": prompt,
+        "neg_prompt": neg_prompt,
+        "true_guidance_scale": float(true_guidance_scale),
+        "num_inference_steps": int(num_inference_steps),
+        "layers": int(refine_sub_layers),
+        "cfg_norm": bool(cfg_norm),
+        "use_en_prompt": bool(use_en_prompt),
+        "resolution": int(resolution),
+        "gpu_duration": int(gpu_duration),
+    }
+    refine_meta = {
+        "mode": "refine",
+        "from_node_id": n.node_id,
+        "layer_idx": idx,
+        "sub_layers": int(refine_sub_layers),
+    }
+
+    node = _create_node(
+        state=state,
+        layers=layers_out,
+        kind="refine",
+        parent_id=n.node_id,
+        params=params,
+        refine_meta=refine_meta,
+    )
+
+    # remember for redo
+    state["last_source_for_redo"] = (n.node_id, idx, int(refine_sub_layers), params)
+
+    if state.get("autosave"):
+        _persist_node_to_dataset(state, node)
+
+    # Update UI: refined accordion opens; current node is refined node now
+    node = _ensure_exports(node)
+    chips = _chips_html(state, node.node_id)
+    history_dd = gr.update(choices=_history_choices(state), value=node.node_id)
+
+    return (
+        state,
+        gr.update(open=True),
+        node.layers,
+        history_dd,
+        chips,
+        node.pptx_path,
+        node.zip_path,
+    )
+
+
+def on_history_select(node_id, state):
+    if state is None:
+        state = new_state()
+    if node_id and node_id in state["nodes"]:
+        state["current_id"] = node_id
+        state["selected_layer_idx"] = 0
+
+    main_gallery, strip, layer_dd, history_dd, chips, pptx, zzip, acc, refined = _render_current(state)
+    return (
+        state,
+        main_gallery,
+        strip,
+        layer_dd,
+        history_dd,
+        chips,
+        pptx,
+        zzip,
+        acc,
+        refined,
+        "Selected: Layer 1",
+    )
+
+
+def on_back_to_parent(state):
+    if state is None:
+        state = new_state()
+    n = _current_node(state)
+    if n and n.parent_id and n.parent_id in state["nodes"]:
+        state["current_id"] = n.parent_id
+        state["selected_layer_idx"] = 0
+    return on_history_select(state.get("current_id"), state)
+
+
+def on_duplicate_node(state):
+    if state is None:
+        state = new_state()
+    n = _current_node(state)
+    if not n:
+        return on_history_select(state.get("current_id"), state)
+
+    # clone layers
+    cloned_layers = [im.copy() for im in n.layers]
+    params = dict(n.params)
+    refine_meta = {"mode": "duplicate", "from_node_id": n.node_id}
+    newn = _create_node(
+        state=state,
+        layers=cloned_layers,
+        kind="refine",
+        parent_id=n.parent_id,
+        params=params,
+        refine_meta=refine_meta,
+    )
+    newn.name = f"{n.name} (copy)"
+
+    if state.get("autosave"):
+        _persist_node_to_dataset(state, newn)
+
+    return on_history_select(newn.node_id, state)
+
+
+def on_rename_node(new_name, state):
+    if state is None:
+        state = new_state()
+    n = _current_node(state)
+    if n and new_name and isinstance(new_name, str):
+        n.name = new_name.strip()[:80] if new_name.strip() else n.name
+        state["nodes"][n.node_id] = n
+        if state.get("autosave"):
+            _persist_node_to_dataset(state, n)
+    return on_history_select(state.get("current_id"), state)
+
+
+def on_redo_refine(
+    seed,
+    randomize_seed,
+    prompt,
+    neg_prompt,
+    true_guidance_scale,
+    num_inference_steps,
+    cfg_norm,
+    use_en_prompt,
+    resolution,
+    gpu_duration,
+    state,
+):
+    if state is None:
+        state = new_state()
+
+    info = state.get("last_source_for_redo")
+    if not info:
         return (
-            st,
-            dd,
-            layers,
-            mini,
-            idx,
-            chips,
-            "Refine: current node has no layers.",
-            gr.update(open=False),
+            state,
+            gr.update(open=True),
             [],
+            gr.update(choices=_history_choices(state)),
+            _chips_html(state, state.get("current_id")),
             None,
             None,
         )
 
-    idx = int(st.get("selected_layer_idx", 0))
-    idx = max(0, min(idx, len(layers_list) - 1))
-    selected_layer = layers_list[idx].convert("RGBA")
+    from_node_id, layer_idx, sub_layers, _params = info
+    if from_node_id not in state["nodes"]:
+        return (
+            state,
+            gr.update(open=True),
+            [],
+            gr.update(choices=_history_choices(state)),
+            _chips_html(state, state.get("current_id")),
+            None,
+            None,
+        )
 
-    # Run pipeline again on selected layer, producing sub-layers
-    sub_layers_count = _clamp_int(sub_layers_count, default=3, lo=2, hi=10)
+    src = state["nodes"][from_node_id]
+    layer_idx = max(0, min(int(layer_idx), len(src.layers) - 1))
+    selected = src.layers[layer_idx]
 
-    sub_layers = gpu_run_pipeline(
-        selected_layer,
-        seed=seed,
-        randomize_seed=randomize_seed,
+    layers_out, used_seed = gpu_run_pipeline(
+        pil_image_rgba=selected.convert("RGBA"),
+        seed=int(seed),
+        randomize_seed=bool(randomize_seed),
         prompt=prompt,
         neg_prompt=neg_prompt,
-        true_guidance_scale=true_guidance_scale,
-        num_inference_steps=num_inference_steps,
-        layer=sub_layers_count,          # <-- only change: layers = sub_layers_count
-        cfg_norm=cfg_norm,
-        use_en_prompt=use_en_prompt,
-        resolution=resolution,
-        gpu_duration=gpu_duration,
+        true_guidance_scale=float(true_guidance_scale),
+        num_inference_steps=int(num_inference_steps),
+        layer=int(sub_layers),
+        cfg_norm=bool(cfg_norm),
+        use_en_prompt=bool(use_en_prompt),
+        resolution=int(resolution),
+        gpu_duration=int(gpu_duration),
     )
 
-    new_id = _make_node(st, layers=sub_layers, parent_id=cur["id"], name=f"refine L{idx+1}")
-    _set_current(st, new_id)
-    st["last_refined_id"] = new_id
+    params = {
+        "seed": int(used_seed),
+        "randomize_seed": bool(randomize_seed),
+        "prompt": prompt,
+        "neg_prompt": neg_prompt,
+        "true_guidance_scale": float(true_guidance_scale),
+        "num_inference_steps": int(num_inference_steps),
+        "layers": int(sub_layers),
+        "cfg_norm": bool(cfg_norm),
+        "use_en_prompt": bool(use_en_prompt),
+        "resolution": int(resolution),
+        "gpu_duration": int(gpu_duration),
+    }
+
+    refine_meta = {
+        "mode": "redo_refine",
+        "from_node_id": from_node_id,
+        "layer_idx": int(layer_idx),
+        "sub_layers": int(sub_layers),
+    }
 
-    dd, layers, mini, idx2, chips = _render_from_state(st)
+    node = _create_node(
+        state=state,
+        layers=layers_out,
+        kind="refine",
+        parent_id=from_node_id,
+        params=params,
+        refine_meta=refine_meta,
+    )
+    node.name = f"Redo refine d{node.depth}"
+
+    if state.get("autosave"):
+        _persist_node_to_dataset(state, node)
+
+    node = _ensure_exports(node)
+    chips = _chips_html(state, node.node_id)
+    history_dd = gr.update(choices=_history_choices(state), value=node.node_id)
 
-    # Export files for current node on-demand (not automatic)
     return (
-        st,
-        dd,
-        layers,
-        mini,
-        idx2,
+        state,
+        gr.update(open=True),
+        node.layers,
+        history_dd,
         chips,
-        f"Refined: created node {_node_label(st['nodes'][new_id])}",
-        gr.update(open=True),     # open refined accordion
-        sub_layers,               # show refined layers
-        None,
-        None,
+        node.pptx_path,
+        node.zip_path,
     )
 
-# ----------------------------
-# App UI
-# ----------------------------
+
+def on_save_current(state):
+    if state is None:
+        state = new_state()
+    n = _current_node(state)
+    if not n:
+        return "Nothing to save."
+    if not _ds_enabled():
+        return "Dataset persistence disabled. Set DATASET_REPO env var and provide hf token."
+    _persist_node_to_dataset(state, n)
+    return f"Saved node {n.node_id[:8]} to dataset session {state['session_id']}."
+
+
+def on_refresh_sessions():
+    if not _ds_enabled():
+        return gr.update(choices=[], value=None), "Dataset persistence disabled."
+    sessions = ds_list_sessions()
+    choices = [(s, s) for s in sessions]
+    return gr.update(choices=choices, value=(choices[-1][1] if choices else None)), f"Found {len(choices)} sessions."
+
+
+def on_load_session(session_id, state):
+    if state is None:
+        state = new_state()
+    if not session_id:
+        return on_history_select(state.get("current_id"), state)
+    if not _ds_enabled():
+        return on_history_select(state.get("current_id"), state)
+    st = _load_session_from_dataset(session_id)
+    # preserve autosave toggle choice from current UI state (if any)
+    st["autosave"] = bool(state.get("autosave", False))
+    return on_history_select(st.get("current_id"), st)
+
+
 ensure_dirname(LOG_DIR)
 
 examples = [
@@ -737,31 +1022,18 @@ examples = [
     "assets/test_images/13.png",
 ]
 
+
 with gr.Blocks() as demo:
-    st = gr.State(value=new_state())
+    state = gr.State(new_state())
 
     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;">'
         )
-
-        persist_status = gr.Markdown(_enabled_msg)
-
-        with gr.Row():
-            new_session_btn = gr.Button("New session", variant="secondary")
-            session_id_box = gr.Textbox(label="Session ID", value="", interactive=False)
-            save_btn = gr.Button("Save session to Dataset repo", variant="primary")
-            save_status = gr.Textbox(label="Save/Load status", value="", interactive=False)
-
-        with gr.Row():
-            load_session_id = gr.Textbox(label="Load Session ID", value="", placeholder="paste Session ID here")
-            load_btn = gr.Button("Load", variant="secondary")
-
         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.
+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.
             """
         )
 
@@ -837,80 +1109,105 @@ It is not designed to control the semantic content of individual layers explicit
                         placeholder="e.g. 60, 120, 300, 1000, 1500",
                     )
 
-                run_button = gr.Button("Decompose!", variant="primary")
+                with gr.Row():
+                    run_button = gr.Button("Decompose!", variant="primary")
+                    fast_button = gr.Button("Apply fast profile", variant="secondary")
 
-                gr.Markdown("### History")
-                history_dd = gr.Dropdown(
-                    label="Nodes",
-                    choices=[],
-                    value=None,
-                    interactive=True,
-                )
-                chips_md = gr.Markdown("")
+                with gr.Accordion("Refine (Recursive Decomposition)", open=True):
+                    refine_sub_layers = gr.Slider(
+                        label="Sub-layers (Refine)",
+                        minimum=2,
+                        maximum=10,
+                        step=1,
+                        value=3,
+                    )
 
-                with gr.Row():
-                    back_btn = gr.Button("← back to parent")
-                    dup_btn = gr.Button("Duplicate node (branch)")
+                    layer_pick_help = gr.Markdown("Pick a layer from the strip or dropdown below.")
 
-                with gr.Row():
-                    rename_inp = gr.Textbox(label="Branch name", value="", placeholder="rename current node")
-                    rename_btn = gr.Button("Rename")
+                    layer_strip = gr.Gallery(
+                        label="Layer strip (click to select)",
+                        columns=6,
+                        rows=1,
+                        height=110,
+                        format="png",
+                        show_label=True,
+                    )
 
-                with gr.Row():
-                    export_btn = gr.Button("Export selected node")
-                export_file = gr.File(label="Download PPTX")
-                export_zip_file = gr.File(label="Download ZIP")
+                    with gr.Row():
+                        layer_dropdown = gr.Dropdown(
+                            label="Selected layer",
+                            choices=[],
+                            value=None,
+                            interactive=True,
+                        )
 
-            with gr.Column(scale=2):
-                gr.Markdown("### Layers (current node)")
-                gallery = gr.Gallery(label="Layers", columns=4, rows=1, format="png")
+                    selected_layer_label = gr.Markdown("Selected: Layer 1")
 
-                with gr.Accordion("Layer picker (mini, click like Photoshop)", open=True):
-                    mini_gallery = gr.Gallery(label="Pick layer to refine", columns=7, rows=1, format="png")
-                    selected_layer_idx = gr.Number(label="Selected layer index (0-based)", value=0, interactive=False)
+                    refine_button = gr.Button("Refine selected layer", variant="primary")
+                    redo_button = gr.Button("↺ redo refine", variant="secondary")
 
-                with gr.Accordion("Refine selected layer", open=True):
-                    refine_info = gr.Textbox(label="Refine status", value="", interactive=False)
+                with gr.Accordion("History", open=True):
+                    history_chips = gr.Markdown("")
+                    history_dropdown = gr.Dropdown(
+                        label="Nodes",
+                        choices=[],
+                        value=None,
+                        interactive=True,
+                    )
                     with gr.Row():
-                        sub_layers_count = gr.Slider(
-                            label="Sub-layers (refine)",
-                            minimum=2,
-                            maximum=10,
-                            step=1,
-                            value=3,
-                        )
-                        refine_btn = gr.Button("Refine selected layer", variant="primary")
+                        back_parent_btn = gr.Button("← back to parent", variant="secondary")
+                        duplicate_btn = gr.Button("Duplicate node (branch)", variant="secondary")
+                    with gr.Row():
+                        rename_text = gr.Textbox(label="Branch name", value="", lines=1, placeholder="Rename current node")
+                        rename_btn = gr.Button("Rename", variant="secondary")
+
+                    autosave = gr.Checkbox(
+                        label=f"Auto-save to Dataset repo ({DATASET_REPO if DATASET_REPO else 'not set'})",
+                        value=False,
+                    )
+                    save_btn = gr.Button("Save current node now", variant="secondary")
+                    save_status = gr.Markdown("")
 
-                refined_acc = gr.Accordion("Refined layers (latest)", open=False)
-                with refined_acc:
-                    refined_gallery = gr.Gallery(label="Refined layers", columns=4, rows=1, format="png")
+                    with gr.Accordion("Load saved sessions (Dataset)", open=False):
+                        refresh_sessions_btn = gr.Button("Refresh sessions list", variant="secondary")
+                        sessions_dropdown = gr.Dropdown(label="Saved sessions", choices=[], value=None, interactive=True)
+                        load_session_btn = gr.Button("Load session", variant="primary")
+                        sessions_status = gr.Markdown("")
+
+            with gr.Column(scale=2):
+                gallery = gr.Gallery(label="Layers", columns=4, rows=2, format="png")
 
-    # Examples run Decompose
+                with gr.Row():
+                    export_file = gr.File(label="Download PPTX (selected node)")
+                    export_zip_file = gr.File(label="Download ZIP (selected node)")
+
+                refined_accordion = gr.Accordion("Refined layers", open=False)
+                with refined_accordion:
+                    refined_gallery = gr.Gallery(label="Refined layers", columns=4, rows=2, format="png")
+
+    # Examples
     gr.Examples(
         examples=examples,
         inputs=[input_image],
         outputs=[gallery, export_file, export_zip_file],
-        fn=lambda img: ([], None, None),
+        fn=lambda x: ([], None, None),
+        examples_per_page=14,
         cache_examples=False,
         run_on_click=False,
     )
 
-    # Boot / init
-    demo.load(
-        fn=ui_boot,
+    # Fast profile just updates UI fields
+    fast_button.click(
+        fn=on_apply_fast_profile,
         inputs=[],
-        outputs=[persist_status, st],
-    ).then(
-        fn=lambda st: st.get("session_id", ""),
-        inputs=[st],
-        outputs=[session_id_box],
+        outputs=[num_inference_steps, layer, resolution],
     )
 
-    # New session
-    new_session_btn.click(
-        fn=on_new_session,
-        inputs=[st],
-        outputs=[st, session_id_box, history_dd, gallery, mini_gallery, chips_md, refine_info, save_status, export_file, export_zip_file],
+    # Autosave toggle
+    autosave.change(
+        fn=on_toggle_autosave,
+        inputs=[autosave, state],
+        outputs=[state],
     )
 
     # Decompose
@@ -929,70 +1226,69 @@ It is not designed to control the semantic content of individual layers explicit
             use_en_prompt,
             resolution,
             gpu_duration,
-            st,
+            state,
         ],
         outputs=[
-            st,
-            history_dd,
+            state,
             gallery,
-            mini_gallery,
-            selected_layer_idx,
-            chips_md,
-            refined_acc,
-            refined_gallery,
+            layer_strip,
+            layer_dropdown,
+            history_dropdown,
+            history_chips,
             export_file,
             export_zip_file,
+            refined_accordion,
+            refined_gallery,
+            seed,
         ],
-    ).then(
-        fn=lambda st: st.get("session_id", ""),
-        inputs=[st],
-        outputs=[session_id_box],
     )
 
-    # History selection
-    history_dd.change(
-        fn=on_history_select,
-        inputs=[history_dd, st],
-        outputs=[st, history_dd, gallery, mini_gallery, selected_layer_idx, chips_md],
-    )
-
-    # Mini gallery click -> choose layer index
-    mini_gallery.select(
-        fn=on_layer_gallery_select,
-        inputs=[st],
-        outputs=[st, selected_layer_idx],
+    # Layer selection by clicking the strip
+    layer_strip.select(
+        fn=on_layer_select_from_strip,
+        inputs=[state],
+        outputs=[state, layer_dropdown, selected_layer_label],
     )
 
-    # Back to parent
-    back_btn.click(
-        fn=on_back_to_parent,
-        inputs=[st],
-        outputs=[st, history_dd, gallery, mini_gallery, selected_layer_idx, chips_md],
+    # Layer selection by dropdown
+    layer_dropdown.change(
+        fn=on_layer_dropdown_change,
+        inputs=[layer_dropdown, state],
+        outputs=[state, selected_layer_label],
     )
 
-    # Duplicate node
-    dup_btn.click(
-        fn=on_duplicate_node,
-        inputs=[st],
-        outputs=[st, history_dd, gallery, mini_gallery, selected_layer_idx, chips_md],
-    )
-
-    # Rename node
-    rename_btn.click(
-        fn=on_rename_node,
-        inputs=[rename_inp, st],
-        outputs=[st, history_dd],
-    ).then(
-        fn=lambda: "",
-        inputs=[],
-        outputs=[rename_inp],
+    # Refine
+    refine_button.click(
+        fn=on_refine_click,
+        inputs=[
+            refine_sub_layers,
+            seed,
+            randomize_seed,
+            prompt,
+            neg_prompt,
+            true_guidance_scale,
+            num_inference_steps,
+            cfg_norm,
+            use_en_prompt,
+            resolution,
+            gpu_duration,
+            state,
+        ],
+        outputs=[
+            state,
+            refined_accordion,
+            refined_gallery,
+            history_dropdown,
+            history_chips,
+            export_file,
+            export_zip_file,
+        ],
     )
 
-    # Refine selected layer
-    refine_btn.click(
-        fn=on_refine_click,
+    # Redo refine
+    redo_button.click(
+        fn=on_redo_refine,
         inputs=[
-            sub_layers_count,
             seed,
             randomize_seed,
             prompt,
@@ -1003,51 +1299,130 @@ It is not designed to control the semantic content of individual layers explicit
             use_en_prompt,
             resolution,
             gpu_duration,
-            st,
+            state,
+        ],
+        outputs=[
+            state,
+            refined_accordion,
+            refined_gallery,
+            history_dropdown,
+            history_chips,
+            export_file,
+            export_zip_file,
         ],
+    )
+
+    # History select
+    history_dropdown.change(
+        fn=on_history_select,
+        inputs=[history_dropdown, state],
         outputs=[
-            st,
-            history_dd,
+            state,
             gallery,
-            mini_gallery,
-            selected_layer_idx,
-            chips_md,
-            refine_info,
-            refined_acc,
+            layer_strip,
+            layer_dropdown,
+            history_dropdown,
+            history_chips,
+            export_file,
+            export_zip_file,
+            refined_accordion,
             refined_gallery,
+            selected_layer_label,
+        ],
+    )
+
+    # Back to parent
+    back_parent_btn.click(
+        fn=on_back_to_parent,
+        inputs=[state],
+        outputs=[
+            state,
+            gallery,
+            layer_strip,
+            layer_dropdown,
+            history_dropdown,
+            history_chips,
+            export_file,
+            export_zip_file,
+            refined_accordion,
+            refined_gallery,
+            selected_layer_label,
+        ],
+    )
+
+    # Duplicate node (branch)
+    duplicate_btn.click(
+        fn=on_duplicate_node,
+        inputs=[state],
+        outputs=[
+            state,
+            gallery,
+            layer_strip,
+            layer_dropdown,
+            history_dropdown,
+            history_chips,
             export_file,
             export_zip_file,
+            refined_accordion,
+            refined_gallery,
+            selected_layer_label,
         ],
     )
 
-    # Export current node
-    export_btn.click(
-        fn=on_export_selected,
-        inputs=[st],
-        outputs=[export_file, export_zip_file],
+    # Rename node
+    rename_btn.click(
+        fn=on_rename_node,
+        inputs=[rename_text, state],
+        outputs=[
+            state,
+            gallery,
+            layer_strip,
+            layer_dropdown,
+            history_dropdown,
+            history_chips,
+            export_file,
+            export_zip_file,
+            refined_accordion,
+            refined_gallery,
+            selected_layer_label,
+        ],
     )
 
-    # Save session
+    # Save node
     save_btn.click(
-        fn=on_save_session,
-        inputs=[st],
-        outputs=[save_status, session_id_box],
+        fn=on_save_current,
+        inputs=[state],
+        outputs=[save_status],
+    )
+
+    # Refresh sessions list
+    refresh_sessions_btn.click(
+        fn=on_refresh_sessions,
+        inputs=[],
+        outputs=[sessions_dropdown, sessions_status],
     )
 
     # Load session
-    load_btn.click(
+    load_session_btn.click(
         fn=on_load_session,
-        inputs=[load_session_id, st],
-        outputs=[st, save_status, history_dd, gallery, mini_gallery, selected_layer_idx, chips_md],
-    ).then(
-        fn=lambda st: st.get("session_id", ""),
-        inputs=[st],
-        outputs=[session_id_box],
-    ).then(
-        fn=lambda: "",
-        inputs=[],
-        outputs=[load_session_id],
+        inputs=[sessions_dropdown, state],
+        outputs=[
+            state,
+            gallery,
+            layer_strip,
+            layer_dropdown,
+            history_dropdown,
+            history_chips,
+            export_file,
+            export_zip_file,
+            refined_accordion,
+            refined_gallery,
+            selected_layer_label,
+        ],
     )
 
+# Serialize GPU tasks; helps stability on some ZeroGPU envs
+demo.queue(concurrency_count=1, max_size=20)
+
 if __name__ == "__main__":
     demo.launch()