Update tab_batch.py

app.py

@@ -12,6 +12,7 @@ from tab_batch import render_batch_processor
 from tab_timeline import render_timeline_tab
 from tab_timeline_wip import render_timeline_wip
 from tab_comfy import render_comfy_monitor
+from tab_raw import render_raw_editor
 
 # ==========================================
 # 1. PAGE CONFIGURATION
@@ -139,6 +140,10 @@ with st.sidebar:
         st.session_state.file_selector = json_files[0].name
     
     selected_file_name = st.radio("Select File", [f.name for f in json_files], key="file_selector")
+    
+    # --- GLOBAL MONITOR TOGGLE (NEW) ---
+    st.markdown("---")
+    show_monitor = st.checkbox("Show Comfy Monitor", value=True)
 
 # ==========================================
 # 4. MAIN APP LOGIC
@@ -160,8 +165,6 @@ if selected_file_name:
         st.session_state.edit_history_idx = None
         
         # --- AUTO-SWITCH TAB LOGIC ---
-        # If the file has 'batch_data' or is a list, force Batch tab.
-        # Otherwise, force Single tab.
         is_batch = "batch_data" in data or isinstance(data, list)
         if is_batch:
             st.session_state.active_tab_name = "🚀 Batch Processor"
@@ -173,31 +176,30 @@ if selected_file_name:
 
     st.title(f"Editing: {selected_file_name}")
 
-    # --- CONTROLLED NAVIGATION (REPLACES ST.TABS) ---
-    # Using radio buttons allows us to change 'active_tab_name' programmatically above.
+    # --- CONTROLLED NAVIGATION ---
+    # Removed "🔌 Comfy Monitor" from this list
     tabs_list = [
         "📝 Single Editor", 
         "🚀 Batch Processor", 
         "🕒 Timeline", 
         "🧪 Interactive Timeline",
-        "🔌 Comfy Monitor"
+        "💻 Raw Editor"
     ]
     
-    # Ensure active tab is valid (safety check)
     if st.session_state.active_tab_name not in tabs_list:
         st.session_state.active_tab_name = tabs_list[0]
 
     current_tab = st.radio(
         "Navigation", 
         tabs_list,
-        key="active_tab_name", # Binds to session state
+        key="active_tab_name", 
         horizontal=True,
         label_visibility="collapsed"
     )
     
     st.markdown("---")
 
-    # --- RENDER SELECTED TAB ---
+    # --- RENDER EDITOR TABS ---
     if current_tab == "📝 Single Editor":
         render_single_editor(data, file_path)
         
@@ -209,6 +211,12 @@ if selected_file_name:
         
     elif current_tab == "🧪 Interactive Timeline":
         render_timeline_wip(data, file_path)
+        
+    elif current_tab == "💻 Raw Editor":
+        render_raw_editor(data, file_path)
 
-    elif current_tab == "🔌 Comfy Monitor":
-        render_comfy_monitor()
+    # --- GLOBAL PERSISTENT MONITOR ---
+    if show_monitor:
+        st.markdown("---")
+        with st.expander("🔌 ComfyUI Monitor", expanded=True):
+            render_comfy_monitor()

tab_batch.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import random
+import copy
 from utils import DEFAULTS, save_json, load_json
 from history_tree import HistoryTree 
 
@@ -96,6 +97,7 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
     st.markdown("---")
     st.info(f"Batch contains {len(batch_list)} sequences.")
 
+    # Updated LoRA keys to match new logic
     lora_keys = ["lora 1 high", "lora 1 low", "lora 2 high", "lora 2 low", "lora 3 high", "lora 3 low"]
     standard_keys = {
         "general_prompt", "general_negative", "current_prompt", "negative", "prompt", "seed",
@@ -112,7 +114,7 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
         prefix = f"{selected_file_name}_seq{i}_v{st.session_state.ui_reset_token}" 
 
         with st.expander(f"🎬 Sequence #{seq_num}", expanded=False):
-            # --- NEW: ACTION ROW WITH CLONING ---
+            # --- ACTION ROW ---
             act_c1, act_c2, act_c3, act_c4 = st.columns([1.2, 1.8, 1.2, 0.5])
             
             # 1. Copy Source
@@ -131,18 +133,14 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
                     st.toast("Copied!", icon="📥")
                     st.rerun()
 
-            # 2. Cloning Tools (Next / End)
+            # 2. Cloning Tools
             with act_c2:
                 cl_1, cl_2 = st.columns(2)
-                
-                # Clone Next
                 if cl_1.button("👯 Next", key=f"{prefix}_c_next", help="Clone and insert below", use_container_width=True):
                     new_seq = seq.copy()
-                    # Calculate new max sequence number
                     max_sn = 0
                     for s in batch_list: max_sn = max(max_sn, int(s.get("sequence_number", 0)))
                     new_seq["sequence_number"] = max_sn + 1
-                    
                     batch_list.insert(i + 1, new_seq)
                     data["batch_data"] = batch_list
                     save_json(file_path, data)
@@ -150,13 +148,11 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
                     st.toast("Cloned to Next!", icon="👯")
                     st.rerun()
 
-                # Clone End
                 if cl_2.button("⏬ End", key=f"{prefix}_c_end", help="Clone and add to bottom", use_container_width=True):
                     new_seq = seq.copy()
                     max_sn = 0
                     for s in batch_list: max_sn = max(max_sn, int(s.get("sequence_number", 0)))
                     new_seq["sequence_number"] = max_sn + 1
-                    
                     batch_list.append(new_seq)
                     data["batch_data"] = batch_list
                     save_json(file_path, data)
@@ -192,7 +188,7 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
                 seq["negative"] = st.text_area("Specific Negative", value=seq.get("negative", ""), height=60, key=f"{prefix}_sn")
             
             with c2:
-                seq["sequence_number"] = st.number_input("Seq Num", value=int(seq_num), key=f"{prefix}_sn_val")
+                seq["sequence_number"] = st.number_input("Sequence Number", value=int(seq_num), key=f"{prefix}_sn_val")
                 
                 s_row1, s_row2 = st.columns([3, 1])
                 seed_key = f"{prefix}_seed"
@@ -211,32 +207,66 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
                 seq["flf"] = st.text_input("FLF", value=str(seq.get("flf", DEFAULTS["flf"])), key=f"{prefix}_flf")
                 
                 if "video file path" in seq or "vace" in selected_file_name:
-                    seq["video file path"] = st.text_input("Video Path", value=seq.get("video file path", ""), key=f"{prefix}_vid")
+                    seq["video file path"] = st.text_input("Video File Path", value=seq.get("video file path", ""), key=f"{prefix}_vid")
                     with st.expander("VACE Settings"):
-                        seq["frame_to_skip"] = st.number_input("Skip", value=int(seq.get("frame_to_skip", 81)), key=f"{prefix}_fts")
-                        seq["input_a_frames"] = st.number_input("In A", value=int(seq.get("input_a_frames", 0)), key=f"{prefix}_ia")
-                        seq["input_b_frames"] = st.number_input("In B", value=int(seq.get("input_b_frames", 0)), key=f"{prefix}_ib")
-                        seq["reference switch"] = st.number_input("Switch", value=int(seq.get("reference switch", 1)), key=f"{prefix}_rsw")
-                        seq["vace schedule"] = st.number_input("Sched", value=int(seq.get("vace schedule", 1)), key=f"{prefix}_vsc")
-                        seq["reference path"] = st.text_input("Ref Path", value=seq.get("reference path", ""), key=f"{prefix}_rp")
-                        seq["reference image path"] = st.text_input("Ref Img", value=seq.get("reference image path", ""), key=f"{prefix}_rip")
+                        seq["frame_to_skip"] = st.number_input("Frame to Skip", value=int(seq.get("frame_to_skip", 81)), key=f"{prefix}_fts")
+                        seq["input_a_frames"] = st.number_input("Input A Frames", value=int(seq.get("input_a_frames", 0)), key=f"{prefix}_ia")
+                        seq["input_b_frames"] = st.number_input("Input B Frames", value=int(seq.get("input_b_frames", 0)), key=f"{prefix}_ib")
+                        seq["reference switch"] = st.number_input("Reference Switch", value=int(seq.get("reference switch", 1)), key=f"{prefix}_rsw")
+                        seq["vace schedule"] = st.number_input("VACE Schedule", value=int(seq.get("vace schedule", 1)), key=f"{prefix}_vsc")
+                        seq["reference path"] = st.text_input("Reference Path", value=seq.get("reference path", ""), key=f"{prefix}_rp")
+                        seq["reference image path"] = st.text_input("Reference Image Path", value=seq.get("reference image path", ""), key=f"{prefix}_rip")
                 
                 if "i2v" in selected_file_name and "vace" not in selected_file_name:
-                    seq["reference image path"] = st.text_input("Ref Img", value=seq.get("reference image path", ""), key=f"{prefix}_ri2")
-                    seq["flf image path"] = st.text_input("FLF Img", value=seq.get("flf image path", ""), key=f"{prefix}_flfi")
+                    seq["reference image path"] = st.text_input("Reference Image Path", value=seq.get("reference image path", ""), key=f"{prefix}_ri2")
+                    seq["flf image path"] = st.text_input("FLF Image Path", value=seq.get("flf image path", ""), key=f"{prefix}_flfi")
 
-            # --- LoRA Settings (Reverted to plain text) ---
+            # --- UPDATED: LoRA Settings with Tag Wrapping ---
             with st.expander("💊 LoRA Settings"):
                 lc1, lc2, lc3 = st.columns(3)
-                with lc1:
-                    seq["lora 1 high"] = st.text_input("LoRA 1 Name", value=seq.get("lora 1 high", ""), key=f"{prefix}_l1h")
-                    seq["lora 1 low"] = st.text_input("LoRA 1 Strength", value=str(seq.get("lora 1 low", "")), key=f"{prefix}_l1l")
-                with lc2:
-                    seq["lora 2 high"] = st.text_input("LoRA 2 Name", value=seq.get("lora 2 high", ""), key=f"{prefix}_l2h")
-                    seq["lora 2 low"] = st.text_input("LoRA 2 Strength", value=str(seq.get("lora 2 low", "")), key=f"{prefix}_l2l")
-                with lc3:
-                    seq["lora 3 high"] = st.text_input("LoRA 3 Name", value=seq.get("lora 3 high", ""), key=f"{prefix}_l3h")
-                    seq["lora 3 low"] = st.text_input("LoRA 3 Strength", value=str(seq.get("lora 3 low", "")), key=f"{prefix}_l3l")
+                
+                # Helper to render the tag wrapper UI
+                def render_lora_col(col_obj, lora_idx):
+                    with col_obj:
+                        st.caption(f"**LoRA {lora_idx}**")
+                        
+                        # --- HIGH ---
+                        k_high = f"lora {lora_idx} high"
+                        raw_h = str(seq.get(k_high, ""))
+                        # Strip tags for display
+                        disp_h = raw_h.replace("<lora:", "").replace(">", "")
+                        
+                        st.write("High:")
+                        rh1, rh2, rh3 = st.columns([0.25, 1, 0.1])
+                        rh1.markdown("<div style='text-align: right; padding-top: 8px;'><code>&lt;lora:</code></div>", unsafe_allow_html=True)
+                        val_h = rh2.text_input(f"L{lora_idx}H", value=disp_h, key=f"{prefix}_l{lora_idx}h", label_visibility="collapsed")
+                        rh3.markdown("<div style='padding-top: 8px;'><code>&gt;</code></div>", unsafe_allow_html=True)
+                        
+                        if val_h:
+                            seq[k_high] = f"<lora:{val_h}>"
+                        else:
+                            seq[k_high] = ""
+
+                        # --- LOW ---
+                        k_low = f"lora {lora_idx} low"
+                        raw_l = str(seq.get(k_low, ""))
+                        # Strip tags for display
+                        disp_l = raw_l.replace("<lora:", "").replace(">", "")
+                        
+                        st.write("Low:")
+                        rl1, rl2, rl3 = st.columns([0.25, 1, 0.1])
+                        rl1.markdown("<div style='text-align: right; padding-top: 8px;'><code>&lt;lora:</code></div>", unsafe_allow_html=True)
+                        val_l = rl2.text_input(f"L{lora_idx}L", value=disp_l, key=f"{prefix}_l{lora_idx}l", label_visibility="collapsed")
+                        rl3.markdown("<div style='padding-top: 8px;'><code>&gt;</code></div>", unsafe_allow_html=True)
+                        
+                        if val_l:
+                            seq[k_low] = f"<lora:{val_l}>"
+                        else:
+                            seq[k_low] = ""
+
+                render_lora_col(lc1, 1)
+                render_lora_col(lc2, 2)
+                render_lora_col(lc3, 3)
 
             # --- CUSTOM PARAMETERS ---
             st.markdown("---")
@@ -289,7 +319,7 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
             tree_data = data.get("history_tree", {})
             htree = HistoryTree(tree_data)
             
-            snapshot_payload = data.copy()
+            snapshot_payload = copy.deepcopy(data)
             if "history_tree" in snapshot_payload: del snapshot_payload["history_tree"]
             
             htree.commit(snapshot_payload, note=commit_msg if commit_msg else "Batch Update")
@@ -301,4 +331,4 @@ def render_batch_processor(data, file_path, json_files, current_dir, selected_fi
                 del st.session_state.restored_indicator
             
             st.toast("Batch Saved & Snapshot Created!", icon="🚀")
-            st.rerun()
+            st.rerun()

tab_comfy.py

@@ -2,14 +2,31 @@ import streamlit as st
 import requests
 from PIL import Image
 from io import BytesIO
+import urllib.parse
+import time  # <--- NEW IMPORT
 from utils import save_config
 
-def render_single_instance(instance_config, index, all_instances):
+def render_single_instance(instance_config, index, all_instances, timeout_minutes):
     url = instance_config.get("url", "http://127.0.0.1:8188")
     name = instance_config.get("name", f"Server {index+1}")
     
     COMFY_URL = url.rstrip("/")
     
+    # --- TIMEOUT LOGIC ---
+    # Generate unique keys for session state
+    toggle_key = f"live_toggle_{index}"
+    start_time_key = f"live_start_{index}"
+    
+    # Check if we need to auto-close
+    if st.session_state.get(toggle_key, False) and timeout_minutes > 0:
+        start_time = st.session_state.get(start_time_key, 0)
+        elapsed = time.time() - start_time
+        if elapsed > (timeout_minutes * 60):
+            st.session_state[toggle_key] = False
+            # We don't need st.rerun() here because the fragment loop will pick up the state change on the next pass
+            # but an explicit rerun makes it snappy.
+            st.rerun()
+
     c_head, c_set = st.columns([3, 1])
     c_head.markdown(f"### 🔌 {name}")
     
@@ -29,7 +46,7 @@ def render_single_instance(instance_config, index, all_instances):
             save_config(
                 st.session_state.current_dir, 
                 st.session_state.config['favorites'], 
-                {"comfy_instances": all_instances}
+                st.session_state.config 
             )
             st.toast("Server config saved!", icon="💾")
             st.rerun()
@@ -41,7 +58,7 @@ def render_single_instance(instance_config, index, all_instances):
             save_config(
                 st.session_state.current_dir, 
                 st.session_state.config['favorites'], 
-                {"comfy_instances": all_instances}
+                st.session_state.config
             )
             st.rerun()
 
@@ -64,18 +81,32 @@ def render_single_instance(instance_config, index, all_instances):
             col1.metric("Status", "🔴 Offline")
             col2.metric("Pending", "-")
             col3.metric("Running", "-")
-            st.error(f"Could not connect to {COMFY_URL}")
-            return 
-
-    # --- 2. LIVE VIEW (WITH TOGGLE) ---
+            st.error(f"Could not connect to API at {COMFY_URL}")
+    
+    # --- 2. LIVE VIEW (VIA REMOTE BROWSER) ---
     st.write("") 
     c_label, c_ctrl = st.columns([1, 2])
     c_label.subheader("📺 Live View")
     
-    # LIVE PREVIEW TOGGLE
-    enable_preview = c_ctrl.checkbox("Enable Live Preview", value=True, key=f"live_toggle_{index}")
+    # Capture the toggle interaction to set start time
+    def on_toggle_change():
+        if st.session_state[toggle_key]:
+            st.session_state[start_time_key] = time.time()
+
+    enable_preview = c_ctrl.checkbox(
+        "Enable Live Preview", 
+        value=False, 
+        key=toggle_key,
+        on_change=on_toggle_change
+    )
     
     if enable_preview:
+        # Display Countdown if timeout is active
+        if timeout_minutes > 0:
+            elapsed = time.time() - st.session_state.get(start_time_key, time.time())
+            remaining = (timeout_minutes * 60) - elapsed
+            st.caption(f"⏱️ Auto-off in: **{int(remaining)}s**")
+
         # Height Slider
         iframe_h = st.slider(
             "Height (px)", 
@@ -83,16 +114,21 @@ def render_single_instance(instance_config, index, all_instances):
             key=f"h_slider_{index}"
         )
 
+        # Get Configured Viewer URL
+        viewer_base = st.session_state.config.get("viewer_url", "http://192.168.1.51:5800")
+        final_src = viewer_base
+
+        st.info(f"Viewing via Remote Browser: `{final_src}`")
         st.markdown(
             f"""
-            <iframe src="{COMFY_URL}" width="100%" height="{iframe_h}px" 
-                style="border: 1px solid #444; border-radius: 8px; box-shadow: 0 4px 6px rgba(0,0,0,0.3);">
+            <iframe src="{final_src}" width="100%" height="{iframe_h}px" 
+                style="border: 2px solid #666; border-radius: 8px; box-shadow: 0 4px 6px rgba(0,0,0,0.3);">
             </iframe>
             """,
             unsafe_allow_html=True
         )
     else:
-        st.info("Live Preview is disabled. Enable it above to see the interface.")
+        st.info("Live Preview is disabled.")
 
     st.markdown("---")
 
@@ -130,7 +166,42 @@ def render_single_instance(instance_config, index, all_instances):
         except Exception as e:
             st.error(f"Error fetching image: {e}")
 
-def render_comfy_monitor():
+# Check for fragment support (Streamlit 1.37+)
+if hasattr(st, "fragment"):
+    # This decorator ensures this function re-runs every 10 seconds automatically
+    # allowing it to catch the timeout even if you are away from the keyboard.
+    @st.fragment(run_every=300)
+    def _monitor_fragment():
+        _render_content()
+else:
+    # Fallback for older Streamlit versions (Won't auto-refresh while idle)
+    def _monitor_fragment():
+        _render_content()
+
+def _render_content():
+    # --- GLOBAL SETTINGS FOR MONITOR ---
+    with st.expander("🔧 Monitor Settings", expanded=False):
+        c_set1, c_set2 = st.columns(2)
+        
+        current_viewer = st.session_state.config.get("viewer_url", "http://192.168.1.51:5800")
+        new_viewer = c_set1.text_input("Remote Browser URL", value=current_viewer, help="e.g., http://192.168.1.51:5800")
+        
+        # New Timeout Slider
+        current_timeout = st.session_state.config.get("monitor_timeout", 0)
+        new_timeout = c_set2.slider("Live Preview Timeout (Minutes)", 0, 60, value=current_timeout, help="0 = Always On. Sets how long the preview stays open before auto-closing.")
+
+        if st.button("💾 Save Monitor Settings"):
+            st.session_state.config["viewer_url"] = new_viewer
+            st.session_state.config["monitor_timeout"] = new_timeout
+            save_config(
+                st.session_state.current_dir, 
+                st.session_state.config['favorites'], 
+                st.session_state.config
+            )
+            st.success("Settings saved!")
+            st.rerun()
+
+    # --- INSTANCE MANAGEMENT ---
     if "comfy_instances" not in st.session_state.config:
         st.session_state.config["comfy_instances"] = [
             {"name": "Main Server", "url": "http://192.168.1.100:8188"}
@@ -140,9 +211,11 @@ def render_comfy_monitor():
     tab_names = [i["name"] for i in instances] + ["➕ Add Server"]
     tabs = st.tabs(tab_names)
     
+    timeout_val = st.session_state.config.get("monitor_timeout", 0)
+
     for i, tab in enumerate(tabs[:-1]):
         with tab:
-            render_single_instance(instances[i], i, instances)
+            render_single_instance(instances[i], i, instances, timeout_val)
             
     with tabs[-1]:
         st.header("Add New ComfyUI Instance")
@@ -157,9 +230,13 @@ def render_comfy_monitor():
                     save_config(
                         st.session_state.current_dir, 
                         st.session_state.config['favorites'], 
-                        {"comfy_instances": instances}
+                        st.session_state.config
                     )
                     st.success("Server Added!")
                     st.rerun()
                 else:
                     st.error("Please fill in both Name and URL.")
+
+def render_comfy_monitor():
+    # We call the wrapper which decides if it's a fragment or not
+    _monitor_fragment()

tab_raw.py

@@ -0,0 +1,78 @@
+import streamlit as st
+import json
+import copy
+from utils import save_json, get_file_mtime
+
+def render_raw_editor(data, file_path):
+    st.subheader(f"💻 Raw Editor: {file_path.name}")
+    
+    # Toggle to hide massive history objects
+    # This is crucial because history trees can get huge and make the text area laggy.
+    col_ctrl, col_info = st.columns([1, 2])
+    with col_ctrl:
+        hide_history = st.checkbox(
+            "Hide History (Safe Mode)", 
+            value=True, 
+            help="Hides 'history_tree' and 'prompt_history' to keep the editor fast and prevent accidental deletion of version control."
+        )
+
+    # Prepare display data
+    if hide_history:
+        display_data = copy.deepcopy(data)
+        # Safely remove heavy keys for the view only
+        if "history_tree" in display_data: del display_data["history_tree"]
+        if "prompt_history" in display_data: del display_data["prompt_history"]
+    else:
+        display_data = data
+
+    # Convert to string
+    # ensure_ascii=False ensures emojis and special chars render correctly
+    try:
+        json_str = json.dumps(display_data, indent=4, ensure_ascii=False)
+    except Exception as e:
+        st.error(f"Error serializing JSON: {e}")
+        json_str = "{}"
+
+    # The Text Editor
+    # We use ui_reset_token in the key to force the text area to reload content on save
+    new_json_str = st.text_area(
+        "JSON Content", 
+        value=json_str, 
+        height=650, 
+        key=f"raw_edit_{file_path.name}_{st.session_state.ui_reset_token}"
+    )
+    
+    st.markdown("---")
+    
+    if st.button("💾 Save Raw Changes", type="primary", use_container_width=True):
+        try:
+            # 1. Parse the text back to JSON
+            input_data = json.loads(new_json_str)
+            
+            # 2. If we were in Safe Mode, we must merge the hidden history back in
+            if hide_history:
+                if "history_tree" in data:
+                    input_data["history_tree"] = data["history_tree"]
+                if "prompt_history" in data:
+                    input_data["prompt_history"] = data["prompt_history"]
+            
+            # 3. Save to Disk
+            save_json(file_path, input_data)
+            
+            # 4. Update Session State
+            # We clear and update the existing dictionary object so other tabs see the changes
+            data.clear()
+            data.update(input_data)
+            
+            # 5. Update Metadata to prevent conflict warnings
+            st.session_state.last_mtime = get_file_mtime(file_path)
+            st.session_state.ui_reset_token += 1
+            
+            st.toast("Raw JSON Saved Successfully!", icon="✅")
+            st.rerun()
+            
+        except json.JSONDecodeError as e:
+            st.error(f"❌ Invalid JSON Syntax: {e}")
+            st.error("Please fix the formatting errors above before saving.")
+        except Exception as e:
+            st.error(f"❌ Unexpected Error: {e}")

tab_timeline.py

@@ -59,6 +59,11 @@ def render_timeline_tab(data, file_path):
                 with c3:
                     if not is_head:
                         if st.button("⏪", key=f"log_rst_{n['id']}", help="Restore this version"):
+                            # --- FIX: Cleanup 'batch_data' if restoring a Single File ---
+                            if "batch_data" not in n["data"] and "batch_data" in data:
+                                del data["batch_data"]
+                            # -------------------------------------------------------------
+                            
                             data.update(n["data"])
                             htree.head_id = n['id']
                             data["history_tree"] = htree.to_dict()
@@ -102,6 +107,11 @@ def render_timeline_tab(data, file_path):
         with col_act:
             st.write(""); st.write("")
             if st.button("⏪ Restore Version", type="primary", use_container_width=True):
+                # --- FIX: Cleanup 'batch_data' if restoring a Single File ---
+                if "batch_data" not in node_data and "batch_data" in data:
+                    del data["batch_data"]
+                # -------------------------------------------------------------
+
                 data.update(node_data)
                 htree.head_id = selected_node['id']
                 data["history_tree"] = htree.to_dict()

tab_timeline_wip.py

@@ -54,18 +54,21 @@ def render_timeline_wip(data, file_path):
                 type="STRAIGHT"
             ))
 
+    # --- UPDATED CONFIGURATION ---
     config = Config(
         width="100%",
-        height="400px", 
+        # Increased height from 400px to 600px for better visibility
+        height="600px", 
         directed=True, 
         physics=False, 
         hierarchical=True, 
         layout={
             "hierarchical": {
                 "enabled": True,
-                "levelSeparation": 150,
-                "nodeSpacing": 100,
-                "treeSpacing": 100,
+                # Increased separation to widen the tree structure
+                "levelSeparation": 200,  # Was 150
+                "nodeSpacing": 150,      # Was 100
+                "treeSpacing": 150,      # Was 100
                 "direction": "LR", 
                 "sortMethod": "directed"
             }
@@ -96,6 +99,11 @@ def render_timeline_wip(data, file_path):
         with c_h2:
             st.write(""); st.write("")
             if st.button("⏪ Restore This Version", type="primary", use_container_width=True, key=f"rst_{target_node_id}"):
+                # --- FIX: Cleanup 'batch_data' if restoring a Single File ---
+                if "batch_data" not in node_data and "batch_data" in data:
+                    del data["batch_data"]
+                # -------------------------------------------------------------
+
                 data.update(node_data)
                 htree.head_id = target_node_id
                 
@@ -172,4 +180,4 @@ def render_timeline_wip(data, file_path):
         else:
             # Single File Preview
             prefix = f"p_{target_node_id}_single"
-            render_preview_fields(node_data, prefix)
+            render_preview_fields(node_data, prefix)