Update app.py

app.py

@@ -226,12 +226,8 @@ def ds_list_sessions(max_sessions: int = 50) -> Tuple[List[str], str]:
     repo_id = ds_repo_id()
     try:
         files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
-        # sessions/<session_id>/session.json
         sess = set()
         for p in files:
-            if p.startswith("sessions/") and p.count("/") >= 2 and p.endswith("/session.json"):
-                # rarely list_repo_files returns directories without trailing slash, but for files it's path
-                pass
             if p.startswith("sessions/") and p.endswith("session.json"):
                 parts = p.split("/")
                 if len(parts) >= 3:
@@ -241,6 +237,44 @@ def ds_list_sessions(max_sessions: int = 50) -> Tuple[List[str], str]:
     except Exception as e:
         return [], f"List sessions failed: {repr(e)}"
 
+def ds_get_root_index() -> Tuple[Optional[Dict[str, Any]], str]:
+    """
+    Read dataset root index.json (used for "last session").
+    Expected keys: id or last_session_id or session_id.
+    """
+    b, msg = ds_download_bytes("index.json")
+    if b is None:
+        return None, msg
+    try:
+        obj = json.loads(b.decode("utf-8"))
+        if not isinstance(obj, dict):
+            return None, "index.json is not an object"
+        return obj, "OK"
+    except Exception as e:
+        return None, f"Failed to parse index.json: {repr(e)}"
+
+def ds_get_last_session_id() -> Tuple[Optional[str], str]:
+    idx, msg = ds_get_root_index()
+    if not idx:
+        return None, msg
+    for k in ("id", "last_session_id", "session_id"):
+        v = idx.get(k)
+        if isinstance(v, str) and v.strip():
+            return v.strip(), "OK"
+    return None, "index.json missing id/last_session_id/session_id"
+
+def ds_set_root_index(session_id: str) -> Tuple[bool, str]:
+    """
+    Update dataset root index.json so we can auto-load the last session after refresh/restart.
+    """
+    payload = {
+        "id": session_id,
+        "last_session_id": session_id,
+        "updated_at": time.time(),
+    }
+    b = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
+    return ds_upload_bytes("index.json", b, f"update index.json {session_id}")
+
 
 # -------------------------
 # Node / History model
@@ -504,7 +538,16 @@ def _persist_session_manifest(state: Dict[str, Any]) -> Tuple[bool, str]:
     }
     b = json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8")
     path = f"{_session_base(session_id)}/session.json"
-    return ds_upload_bytes(path, b, f"save session manifest {session_id}")
+    ok_m, msg_m = ds_upload_bytes(path, b, f"save session manifest {session_id}")
+    if not ok_m:
+        return False, msg_m
+
+    # Update root index.json (best-effort, do not fail save if it can't be updated)
+    ok_i, msg_i = ds_set_root_index(session_id)
+    if not ok_i:
+        return True, f"{msg_m} (warning: index.json update failed: {msg_i})"
+
+    return True, f"{msg_m} + updated index.json"
 
 def _load_session_manifest(session_id: str) -> Tuple[Optional[Dict[str, Any]], str]:
     b, msg = ds_download_bytes(f"{_session_base(session_id)}/session.json")
@@ -551,7 +594,6 @@ def _persistence_status_text() -> str:
 
 def on_refresh_sessions():
     sessions, msg = ds_list_sessions()
-    # For dropdown choices, use sessions as both label+value
     return gr.update(choices=sessions, value=(sessions[0] if sessions else None)), msg
 
 def on_init_dataset():
@@ -683,7 +725,6 @@ def on_decompose_click(
     )
 
 def on_layer_pick_from_dropdown(state: Dict[str, Any], layer_name: str):
-    # layer_name = "Layer k"
     node_id = state.get("selected_node_id")
     imgs = _get_node_images(state, node_id) if node_id else []
     n = len(imgs)
@@ -703,7 +744,6 @@ def on_layer_pick_from_gallery(state: Dict[str, Any], evt: gr.SelectData):
     n = len(imgs)
     idx = int(evt.index) if evt and evt.index is not None else 0
     idx = max(0, min(n - 1, idx)) if n > 0 else 0
-    # also update dropdown value
     dd_value = f"Layer {idx+1}" if n > 0 else None
     return gr.update(value=idx), gr.update(value=dd_value), _layer_label(idx, n)
 
@@ -739,7 +779,7 @@ def _refine_from_source(
         neg_prompt=neg_prompt,
         true_guidance_scale=true_guidance_scale,
         num_inference_steps=num_inference_steps,
-        layer=sub_layers,  # << refine layers count
+        layer=sub_layers,
         cfg_norm=cfg_norm,
         use_en_prompt=use_en_prompt,
         resolution=resolution,
@@ -790,6 +830,7 @@ def on_refine_click(
             gr.update(),
             gr.update(),
             gr.update(),
+            gr.update(),
             gr.update(visible=False),
             [],
             None,
@@ -824,7 +865,6 @@ def on_refine_click(
 
     sub_layers = _clamp_int(sub_layers, default=3, lo=2, hi=10)
 
-    # Run refine
     layers_out, used_seed, settings_snapshot = _refine_from_source(
         state,
         source_node_id=source_node_id,
@@ -842,7 +882,6 @@ def on_refine_click(
         randomize_seed=randomize_seed,
     )
 
-    # Create child node and select it
     child_name = f"refine ({state['nodes'][source_node_id]['meta']['name']}) L{idx+1}"
     child_id = _add_node(
         state,
@@ -858,37 +897,32 @@ def on_refine_click(
     state["selected_node_id"] = child_id
     state["last_refined_node_id"] = child_id
 
-    # UI update
     n_layers = len(layers_out)
     layer_choices, layer_value = _build_layer_dropdown(n_layers)
     hist_choices = _history_choices(state)
     chips = _make_chips(state)
     selected_label = _layer_label(0, n_layers)
 
-    # exports for current node (child)
     pptx_path, zip_path, exp_msg = _current_node_export(state, child_id)
 
     status = f"Refined into {n_layers} sub-layer(s). Seed={used_seed}. {exp_msg}"
-
-    # refined block visible + show refined gallery
     refined_visible = gr.update(visible=True)
 
-    # Also update dropdown that selects nodes (history)
     return (
         state,
-        layers_out,                 # main gallery shows child node now
-        layers_out,                 # mini picker gallery
+        layers_out,
+        layers_out,
         gr.update(choices=layer_choices, value=layer_value),
         gr.update(value=0),
         selected_label,
         gr.update(choices=[c[1] for c in hist_choices], value=child_id),
         chips,
         refined_visible,
-        layers_out,                 # refined gallery
+        layers_out,
         pptx_path,
         zip_path,
         status,
-        gr.update(),                # no change to seed_used textbox here
+        gr.update(),
     )
 
 def on_history_select(state: Dict[str, Any], node_id: str):
@@ -919,7 +953,6 @@ def on_history_select(state: Dict[str, Any], node_id: str):
 
     pptx_path, zip_path, exp_msg = _current_node_export(state, node_id)
 
-    # hide refined block when user manually switches nodes (keeps UX clean)
     return (
         state,
         imgs,
@@ -952,7 +985,6 @@ def on_duplicate_node(state: Dict[str, Any]):
     new_id = _duplicate_node(state, node_id)
     if not new_id:
         return state, gr.update(), "Duplicate failed."
-    # select duplicated
     return on_history_select(state, new_id)
 
 def on_rename_node(state: Dict[str, Any], new_name: str):
@@ -960,10 +992,7 @@ def on_rename_node(state: Dict[str, Any], new_name: str):
     if not node_id:
         return state, gr.update(), "No selected node."
     _rename_node(state, node_id, new_name)
-    # update history dropdown labels
     hist_choices = _history_choices(state)
-    # gr.Dropdown choices are values; we keep node_id list, but users see ids.
-    # We'll show names in a markdown below.
     chips = _make_chips(state)
     return state, gr.update(choices=[c[1] for c in hist_choices], value=node_id), chips, "Renamed."
 
@@ -1054,7 +1083,6 @@ def on_redo_refine(
         randomize_seed=randomize_seed,
     )
 
-    # create new child under same source (new branch)
     child_name = f"redo refine ({state['nodes'][source_node_id]['meta']['name']}) L{int(source_layer_idx)+1}"
     child_id = _add_node(
         state,
@@ -1070,7 +1098,6 @@ def on_redo_refine(
     state["selected_node_id"] = child_id
     state["last_refined_node_id"] = child_id
 
-    # update UI (same as refine)
     n_layers = len(layers_out)
     layer_choices, layer_value = _build_layer_dropdown(n_layers)
     hist_choices = _history_choices(state)
@@ -1111,7 +1138,7 @@ def on_save_current(state: Dict[str, Any]):
     ok2, msg2 = _persist_session_manifest(state)
     if not ok2:
         return msg2
-    return f"✅ Saved node + session manifest. {msg1}"
+    return f"✅ Saved node + session manifest. {msg1} | {msg2}"
 
 def on_load_session(state: Dict[str, Any], session_id: str):
     if not session_id:
@@ -1121,7 +1148,6 @@ def on_load_session(state: Dict[str, Any], session_id: str):
     if manifest is None:
         return state, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(visible=False), [], None, None, msg
 
-    # rebuild state from manifest (lazy-load images only for selected)
     new_state = _init_state()
     new_state["session_id"] = manifest.get("session_id") or session_id
     new_state["created_at"] = manifest.get("created_at")
@@ -1146,7 +1172,6 @@ def on_load_session(state: Dict[str, Any], session_id: str):
         else:
             return new_state, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(visible=False), [], None, None, f"Loaded manifest but failed to load node images: {msg2}"
 
-        # Also pre-load root images if root != selected (nice UX)
         root = new_state.get("root_node_id")
         if root and root != sel and root in nodes_meta and not new_state["nodes"][root]["images"]:
             rl = int(nodes_meta[root].get("num_layers", 0))
@@ -1180,20 +1205,14 @@ def on_load_session(state: Dict[str, Any], session_id: str):
     )
 
 def on_history_need_images(state: Dict[str, Any], node_id: str):
-    """
-    When switching nodes after loading manifest, images may not be loaded yet.
-    We'll lazy-load from dataset if missing.
-    """
     if not node_id or node_id not in state.get("nodes", {}):
         return state, "Unknown node."
     imgs = state["nodes"][node_id].get("images", [])
     if imgs:
         return state, "OK"
-    # try load from dataset
     session_id = state.get("session_id")
     if not session_id:
         return state, "No session_id."
-    # get num_layers from meta? we stored only in manifest; not guaranteed here. fallback: try read session.json again
     manifest, msg = _load_session_manifest(session_id)
     if not manifest:
         return state, f"Cannot load manifest: {msg}"
@@ -1207,6 +1226,45 @@ def on_history_need_images(state: Dict[str, Any], node_id: str):
     state["nodes"][node_id]["images"] = imgs2
     return state, "Loaded images."
 
+def on_autoload_last_session(state: Dict[str, Any]):
+    if not isinstance(state, dict):
+        state = _init_state()
+    if not ds_enabled():
+        return (
+            state,
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(visible=False),
+            [],
+            None,
+            None,
+            "Dataset persistence disabled (no autoload).",
+        )
+    sid, msg = ds_get_last_session_id()
+    if not sid:
+        return (
+            state,
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(visible=False),
+            [],
+            None,
+            None,
+            f"No last session to autoload: {msg}",
+        )
+    return on_load_session(state, sid)
+
+
 # -------------------------
 # Build UI
 # -------------------------
@@ -1336,7 +1394,6 @@ It is not designed to control the semantic content of individual layers explicit
 
             with gr.Group():
                 gr.Markdown("### Refine (Recursive Decomposition)")
-                # refine params are NOT in Advanced Settings (per your UX request)
                 sub_layers = gr.Slider(
                     label="Sub-layers (Refine)",
                     minimum=2,
@@ -1384,7 +1441,6 @@ It is not designed to control the semantic content of individual layers explicit
                 export_zip = gr.File(label="Download ZIP")
 
             status = gr.Markdown("")
-
             seed_used = gr.Textbox(label="Seed used", value="", interactive=False)
 
     # Examples
@@ -1422,6 +1478,27 @@ It is not designed to control the semantic content of individual layers explicit
         ],
     )
 
+    # Auto-load last session (reads root index.json; supports id/last_session_id/session_id)
+    demo.load(
+        fn=on_autoload_last_session,
+        inputs=[state],
+        outputs=[
+            state,
+            gallery,
+            layer_picker,
+            layer_dropdown,
+            selected_layer_idx,
+            selected_layer_label,
+            history_dd,
+            chips_md,
+            refined_block,
+            refined_gallery,
+            export_pptx,
+            export_zip,
+            status,
+        ],
+    )
+
     # Decompose
     btn_decompose.click(
         fn=on_decompose_click,
@@ -1510,8 +1587,7 @@ It is not designed to control the semantic content of individual layers explicit
 
     # History select (lazy load images first)
     def _history_select_with_lazy(state, node_id):
-        state, msg = on_history_need_images(state, node_id)
-        # ignore msg for now; on_history_select will error if still missing
+        state, _ = on_history_need_images(state, node_id)
         return on_history_select(state, node_id)
 
     history_dd.change(