Update app.py

app.py

@@ -5,28 +5,20 @@
 # ------------------------------
 # 0) Imports & environment
 # ------------------------------
-import os 
+import os
 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
 os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
 os.makedirs(os.environ["HF_HOME"], exist_ok=True)
-import io, csv, time, json, base64, re, zipfile
-from typing import Generator, List, Tuple, Dict, Any
-from threading import Thread
-# Persist model caches between restarts
 
+import io, csv, time, json, base64, re, zipfile
+from typing import List, Tuple, Dict, Any
 
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import LlavaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
-
-# Optional deps for import/export (we handle gracefully if missing)
-try:
-    import pandas as pd  # not required at runtime; kept for future use
-except Exception:
-    pd = None
+from transformers import LlavaForConditionalGeneration, AutoProcessor
 
-# Liger is optional; skip if missing
+# Optional: Liger kernel (ignored if missing)
 try:
     from liger_kernel.transformers import apply_liger_kernel_to_llama
 except Exception:
@@ -46,11 +38,10 @@ except Exception:
 # ------------------------------
 APP_DIR = os.getcwd()
 SESSION_FILE = "/tmp/forge_session.json"
-# --- Branding 
 
+# Branding: fixed logo height
 LOGO_HEIGHT_PX = int(os.getenv("FORGE_LOGO_PX", 60))
 
-
 # Settings live in a user cache dir (persists better than /tmp)
 CONFIG_DIR = os.path.expanduser("~/.cache/forgecaptions")
 os.makedirs(CONFIG_DIR, exist_ok=True)
@@ -58,7 +49,7 @@ SETTINGS_FILE = os.path.join(CONFIG_DIR, "settings.json")
 
 JOURNAL_FILE = "/tmp/forge_journal.json"
 
-# IMPORTANT: keep generated assets in /tmp so Gradio can serve them safely
+# Generated assets in /tmp so Gradio can serve them safely
 THUMB_CACHE = "/tmp/forgecaptions/thumbs"
 EXCEL_THUMB_DIR = "/tmp/forge_excel_thumbs"
 TXT_EXPORT_DIR = "/tmp/forge_txt"
@@ -80,8 +71,6 @@ print(f"[ForgeCaptions] Gradio version: {gr.__version__}")
 
 # ------------------------------
 # 2) Model loader (GPU-safe lazy init)
-#    - Processor on CPU
-#    - Model is ONLY created inside @gpu functions to satisfy Stateless GPU
 # ------------------------------
 processor = AutoProcessor.from_pretrained(MODEL_PATH)
 _MODEL = None
@@ -104,8 +93,8 @@ def get_model():
                 low_cpu_mem_usage=True,
                 device_map=0,
             )
+            # Try to enable Liger on the LLM submodule (best-effort, silent if missing)
             try:
-                from liger_kernel.transformers import apply_liger_kernel_to_llama
                 lm = getattr(_MODEL, "language_model", None) or getattr(_MODEL, "model", None)
                 if lm is not None:
                     ok = apply_liger_kernel_to_llama(lm)
@@ -137,19 +126,19 @@ STYLE_OPTIONS = [
     "Flux.1-Dev",
     "Stable Diffusion",
     "MidJourney",
-    "E-commerce product", 
+    "E-commerce product",
     "Portrait (photography)",
     "Landscape (photography)",
     "Art analysis (no artist names)",
     "Social caption",
     "Aesthetic tags (comma-sep)"
-    ]
+]
 
 CAPTION_TYPE_MAP: Dict[str, str] = {
     "Descriptive": "Write a detailed description for this image.",
     "Character training": (
         "Write a thorough, training-ready caption for a character dataset. "
-        "Describe subject appearance (physique, face/hair), clothing and accessories, actions/pose/gesture, camera angle/focal cues."
+        "Describe subject appearance (physique, face/hair), clothing and accessories, actions/pose/gesture, camera angle/focal cues. "
         "If multiple subjects are present, describe each briefly (most prominent first) and distinguish them by visible traits."
     ),
     "Flux.1-Dev": "Write a Flux.1-Dev style prompt that would reproduce this image faithfully.",
@@ -163,7 +152,6 @@ CAPTION_TYPE_MAP: Dict[str, str] = {
     "Social caption": "Write an engaging caption describing the visible content. No hashtags.",
 }
 
-
 LENGTH_CHOICES = ["any", "very short", "short", "medium-length", "long", "very long"] + [str(i) for i in range(20, 261, 10)]
 
 _LENGTH_HINTS = {
@@ -173,7 +161,6 @@ _LENGTH_HINTS = {
     "long":         "Write a detailed caption (≈80–120 words).",
     "very long":    "Write a very detailed caption (≈150–250 words).",
 }
-
 def _length_hint(choice: str) -> str:
     if not choice or choice == "any":
         return ""
@@ -181,10 +168,9 @@ def _length_hint(choice: str) -> str:
         return f"Limit the caption to at most {choice} words."
     return _LENGTH_HINTS.get(choice, "")
 
-
 EXTRA_CHOICES = [
     "Only include a character's modifiable, style-level attributes (hair style/color, makeup, clothing/accessories, pose, expression). Do NOT mention identity traits (skin tone, age, body type).",
-    "Use profanity when describing sexual acts or genitalia (pussy, fucking, cum, cock, etc.).",   # You do NOT have permission to remove
+    "Use profanity when describing sexual acts or genitalia (pussy, fucking, cum, cock, etc.).",
     "Be sexually graphic and describe sexual position when visible.",
     "Exclude brand names, watermarks, artist names, file/meta info, quality terms, hashtags, and resolution notes.",
     "Arrange information in the order: Subject; Clothing/Accessories; Action/Pose; Setting/Environment; Lighting/Camera.",
@@ -204,7 +190,6 @@ EXTRA_CHOICES = [
 ]
 NAME_OPTION = "If there is a person/character in the image you must refer to them as {name}."
 
-# Place near EXTRA_CHOICES
 DEFAULT_EXTRA_CHOICES = [
     "Use profanity when describing sexual acts or genitalia (pussy, fucking, cum, cock, etc.).",
     "Be sexually graphic and describe sexual position when visible.",
@@ -259,15 +244,16 @@ def load_settings() -> dict:
         "shape_aliases_enabled": True,
         "shape_aliases": [],
         "excel_thumb_px": 128,
-        "logo_px": 60,
+        "logo_px": LOGO_HEIGHT_PX,
         "shape_aliases_persist": True,
         "extras": DEFAULT_EXTRA_CHOICES,
+        "caption_length": "long",
     }
 
     for k, v in defaults.items():
         cfg.setdefault(k, v)
 
-    # Normalize styles to a valid list
+    # Normalize
     styles = cfg.get("styles") or []
     if not isinstance(styles, list):
         styles = [styles]
@@ -348,13 +334,6 @@ def logo_b64_img() -> str:
 # 6) Shape Aliases (plural-aware + '-shaped' variants)
 # ------------------------------
 def _plural_token_regex(tok: str) -> str:
-    """
-    Build a regex for a token that also matches simple English plurals.
-    Rules:
-      - endswith s/x/z/ch/sh → add '(?:es)?'
-      - consonant + y → '(?:y|ies)'
-      - default → 's?'
-    """
     t = (tok or "").strip()
     if not t: return ""
     t_low = t.lower()
@@ -365,11 +344,6 @@ def _plural_token_regex(tok: str) -> str:
     return re.escape(t) + r"s?"
 
 def _compile_shape_aliases_from_file():
-    """
-    Build regex list from settings["shape_aliases"].
-    Left cell accepts comma OR pipe separated synonyms (multi-word OK).
-    Matches are case-insensitive, catches simple plurals, and allows '-shaped' or ' shaped'.
-    """
     s = load_settings()
     if not s.get("shape_aliases_enabled", True):
         return []
@@ -409,11 +383,6 @@ def get_shape_alias_rows_ui_defaults():
     return rows, enabled
 
 def save_shape_alias_rows(enabled, df_rows, persist):
-    """
-    Save or just apply alias rows.
-    - If persist=True → write to SETTINGS_FILE and recompile from file
-    - If persist=False → do NOT touch disk; just compile & apply in-memory
-    """
     cleaned = []
     for r in (df_rows or []):
         if not r:
@@ -431,7 +400,6 @@ def save_shape_alias_rows(enabled, df_rows, persist):
         save_settings(cfg)
         status = "💾 Saved to disk (will persist across restarts)."
 
-    # Recompile in-memory, regardless of persist
     global _SHAPE_ALIASES
     if bool(enabled):
         compiled = []
@@ -453,7 +421,7 @@ def save_shape_alias_rows(enabled, df_rows, persist):
 
 
 # ------------------------------
-# 7) Prompt builder (instruction text shown/used for model)
+# 7) Prompt builder
 # ------------------------------
 def final_instruction(style_list: List[str], extra_opts: List[str], name_value: str, length_choice: str = "long") -> str:
     styles = style_list or ["Character training"]
@@ -463,10 +431,10 @@ def final_instruction(style_list: List[str], extra_opts: List[str], name_value:
         core += " " + " ".join(extra_opts)
     if NAME_OPTION in (extra_opts or []):
         core = core.replace("{name}", (name_value or "{NAME}").strip())
-    if "Aesthetic tags (comma-sep)" not in styles: # If they're asking for comma-separated tags, ignore word-length guidance.
+    if "Aesthetic tags (comma-sep)" not in styles:
         lh = _length_hint(length_choice or "any")
         if lh:
-            core += " " + lh    
+            core += " " + lh
     return core
 
 
@@ -486,9 +454,6 @@ def _build_inputs(im: Image.Image, instr: str, dtype) -> Dict[str, Any]:
 
 @torch.no_grad()
 def caption_once(im: Image.Image, instr: str, temp: float, top_p: float, max_tokens: int) -> str:
-    """
-    NOTE: Not GPU-decorated on purpose; call this only from within a @gpu function.
-    """
     model, device, dtype = get_model()
     inputs = _build_inputs(im, instr, dtype)
     inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
@@ -525,8 +490,8 @@ def run_batch(
     top_p: float,
     max_tokens: int,
     max_side: int,
-    time_budget_s: float | None = None,                    # respects Zero-GPU window (None = unlimited)
-    progress: gr.Progress = gr.Progress(track_tqdm=True),  # drives the progress bar
+    time_budget_s: float | None = None,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),
 ) -> Tuple[List[dict], list, list, str, List[str], int, int]:
     """
     Returns:
@@ -580,17 +545,14 @@ def run_batch(
         total,
     )
 
-@gpu
-@torch.no_grad()
 
 # ------------------------------
-# 9) Export/Import helpers (CSV/XLSX/TXT ZIP)
+# 9) Export helpers (CSV/XLSX/TXT ZIP)
 # ------------------------------
 def _rows_to_table(rows: List[dict]) -> list:
     return [[r.get("filename",""), r.get("caption","")] for r in (rows or [])]
 
 def _table_to_rows(table_value: Any, rows: List[dict]) -> List[dict]:
-    # Expect list-of-lists (Dataframe type="array")
     tbl = table_value or []
     new = []
     for i, r in enumerate(rows or []):
@@ -706,79 +668,6 @@ def export_txt_zip(table_value: Any, dataset_name: str) -> str:
                 z.write(os.path.join(TXT_EXPORT_DIR, fn), arcname=fn)
     return zpath
 
-def import_captions_file(file_path: str, session_rows: List[dict]) -> Tuple[List[dict], list, list, str]:
-    """
-    Import captions from CSV or XLSX and merge by filename into current session.
-    - If filename exists → update its caption
-    - Otherwise append a new row (without image path/thumbnail)
-    """
-    if not file_path or not os.path.exists(file_path):
-        table_rows = _rows_to_table(session_rows)
-        gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
-                         for r in session_rows if (r.get("thumb_path") or r.get("path"))]
-        return session_rows, gallery_pairs, table_rows, "No file selected."
-
-    ext = os.path.splitext(file_path)[1].lower()
-    imported: List[Tuple[str, str]] = []
-
-    try:
-        if ext == ".csv":
-            with open(file_path, "r", encoding="utf-8") as f:
-                reader = csv.reader(f)
-                rows = list(reader)
-            if rows and len(rows[0]) >= 2 and str(rows[0][0]).lower().strip() == "filename":
-                rows = rows[1:]
-            for r in rows:
-                if not r or len(r) < 2:
-                    continue
-                fn = str(r[0]).strip()
-                cap = str(r[1]).strip()
-                if fn:
-                    imported.append((fn, cap))
-
-        elif ext in (".xlsx", ".xls"):
-            try:
-                from openpyxl import load_workbook
-            except Exception as e:
-                raise RuntimeError("XLSX import requires 'openpyxl' in requirements.txt.") from e
-            wb = load_workbook(file_path, read_only=True, data_only=True)
-            ws = wb.active
-            rows = list(ws.iter_rows(values_only=True))
-            if rows and len(rows[0]) >= 2 and str(rows[0][0]).lower().strip() == "filename":
-                rows = rows[1:]
-            for r in rows:
-                if not r or len(r) < 2:
-                    continue
-                fn = str(r[0]).strip() if r[0] is not None else ""
-                cap = str(r[1]).strip() if r[1] is not None else ""
-                if fn:
-                    imported.append((fn, cap))
-        else:
-            return session_rows, _rows_to_table(session_rows), _rows_to_table(session_rows), f"Unsupported file type: {ext}"
-    except Exception as e:
-        table_rows = _rows_to_table(session_rows)
-        gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
-                         for r in session_rows if (r.get("thumb_path") or r.get("path"))]
-        return session_rows, gallery_pairs, table_rows, f"Import failed: {e}"
-
-    # Merge
-    idx_by_name = {r.get("filename",""): i for i, r in enumerate(session_rows)}
-    updates, inserts = 0, 0
-    for fn, cap in imported:
-        if fn in idx_by_name:
-            session_rows[idx_by_name[fn]]["caption"] = cap
-            updates += 1
-        else:
-            session_rows.append({"filename": fn, "caption": cap, "path": "", "thumb_path": ""})
-            inserts += 1
-
-    save_session(session_rows)
-    gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
-                     for r in session_rows if (r.get("thumb_path") or r.get("path"))]
-    table_rows = _rows_to_table(session_rows)
-    stamp = f"Imported {len(imported)} rows • updated {updates} • added {inserts} • {time.strftime('%H:%M:%S')}"
-    return session_rows, gallery_pairs, table_rows, stamp
-
 
 # ------------------------------
 # 10) UI header helper (fixed logo size)
@@ -801,13 +690,15 @@ def _render_header_html(px: int) -> str:
     width: auto;
     object-fit: contain;
     display: block;
+    max-width: 320px; /* cap very wide logos */
   }}
   @media (max-width: 640px) {{
-    .cf-logo {{ height: {max(60, int(px) - 12)}px; }} /* optional small-screen tweak */
+    .cf-logo {{ height: {max(48, int(px) - 8)}px; }}
   }}
 </style>
 """
 
+
 # ------------------------------
 # 11) UI (Blocks)
 # ------------------------------
@@ -819,8 +710,8 @@ BASE_CSS = """
 .cf-hero{display:flex; align-items:center; justify-content:center; gap:16px;
   margin:4px 0 12px; text-align:center;}
 .cf-hero .cf-text{text-align:center;}
-.cf-title{margin:0;font-size:3.25rem;line-height:1;letter-spacing:.2px}
-.cf-sub{margin:6px 0 0;font-size:1.1rem;color:#cfd3da}
+.cf-title{margin:0;font-size:3.0rem;line-height:1;letter-spacing:.2px}
+.cf-sub{margin:6px 0 0;font-size:1.05rem;color:#cfd3da}
 
 /* Results area + robust scrollbars */
 .cf-scroll{border:1px solid #e6e6e6; border-radius:10px; padding:8px}
@@ -831,14 +722,9 @@ BASE_CSS = """
 """
 
 with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
-    # Ensure Spaces sees a GPU function (without touching CUDA in main)
-    demo.load(inputs=None, outputs=None)
-
     # ---- Header
     settings = load_settings()
-    header_html = gr.HTML(_render_header_html(LOGO_HEIGHT_PX))
-
-
+    header_html = gr.HTML(_render_header_html(settings.get("logo_px", LOGO_HEIGHT_PX)))
 
     # ---- Controls group
     with gr.Group():
@@ -867,19 +753,18 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                     trig       = gr.Textbox(label="Trigger word", value=settings.get("trigger",""))
                     add_start  = gr.Textbox(label="Add text to start", value=settings.get("begin",""))
                     add_end    = gr.Textbox(label="Add text to end", value=settings.get("end",""))
-                    
 
-            # RIGHT: instructions + dataset + general sliders + logo controls
+            # RIGHT: instructions + dataset + general sliders
             with gr.Column(scale=1):
                 with gr.Accordion("Model Instructions", open=False):
                     instruction_preview = gr.Textbox(
-                        label=None, 
+                        label=None,
                         lines=12,
                         value=final_instruction(
                             settings.get("styles", ["Character training"]),
                             settings.get("extras", []),
                             settings.get("name",""),
-                            settings.get("caption_length", "long"),     
+                            settings.get("caption_length", "long"),
                         ),
                     )
                 dataset_name = gr.Textbox(label="Dataset name (export title prefix)",
@@ -896,14 +781,13 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 gpu_budget = gr.Slider(20, 110, value=55, step=5, label="Max seconds per GPU call")
                 no_time_limit = gr.Checkbox(value=False, label="No time limit (ignore above)")
 
-
     # Persist instruction + general settings
     def _refresh_instruction(styles, extra, name_value, trigv, begv, endv, excel_px, ms, cap_len):
         instr = final_instruction(styles or ["Character training"], extra or [], name_value, cap_len)
         cfg = load_settings()
         cfg.update({
             "styles": styles or ["Character training"],
-            "extras": extra or [],
+            "extras": _valid_extras(extra),
             "name": name_value,
             "trigger": trigv, "begin": begv, "end": endv,
             "excel_thumb_px": int(excel_px),
@@ -914,9 +798,11 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
         return instr
 
     for comp in [style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side, caption_length]:
-        comp.change(_refresh_instruction,
-                    inputs=[style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side, caption_length],
-                    outputs=[instruction_preview])
+        comp.change(
+            _refresh_instruction,
+            inputs=[style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side, caption_length],
+            outputs=[instruction_preview]
+        )
 
     def _save_dataset_name(name):
         cfg = load_settings()
@@ -925,8 +811,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
         return gr.update()
     dataset_name.change(_save_dataset_name, inputs=[dataset_name], outputs=[])
 
-
-    # ---- Shape Aliases (with plural matching + persist) ----
+    # ---- Shape Aliases (with plural matching + persist)
     with gr.Accordion("Shape Aliases", open=False):
         gr.Markdown(
             "### 🔷 Shape Aliases\n"
@@ -934,8 +819,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             "**How to use:**\n"
             "- Left column = a single token **or** comma/pipe-separated synonyms, e.g. `diamond, rhombus | lozenge`\n"
             "- Right column = replacement name, e.g. `family-emblem`\n"
-            "Matches are case-insensitive, catches simple plurals (`box`→`boxes`, `lady`→`ladies`), "
-            "and also matches `*-shaped` or `* shaped` variants."
+            "Matches are case-insensitive, catches simple plurals, and also matches `*-shaped` / `* shaped` variants."
         )
 
         init_rows, init_enabled = get_shape_alias_rows_ui_defaults()
@@ -959,10 +843,8 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             clear_btn   = gr.Button("Clear", variant="secondary")
             save_btn    = gr.Button("💾 Save", variant="primary")
 
-        # status line for saves
         save_status = gr.Markdown("")
 
-        # --- local handlers (must stay inside Blocks context) ---
         def _add_row(cur):
             cur = (cur or []) + [["", ""]]
             return gr.update(value=cur, row_count=(max(1, len(cur)), "dynamic"))
@@ -980,7 +862,6 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             return gr.update()
         persist_aliases.change(_save_alias_persist_flag, inputs=[persist_aliases], outputs=[])
 
-        # Persist rows if persist_aliases checked; otherwise apply in-memory only
         save_btn.click(
             save_shape_alias_rows,
             inputs=[enable_aliases, alias_table, persist_aliases],
@@ -999,28 +880,16 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 outputs=[single_caption_out]
             )
 
-#        with gr.Tab("Batch"):
-#            with gr.Accordion("Uploaded images", open=True):
-#                input_files = gr.File(label="Drop images (or click to select)", file_types=["image"], file_count="multiple", type="filepath")
-#            run_button = gr.Button("Caption batch", variant="primary")
-
-#            with gr.Accordion("Import captions from CSV/XLSX (merge by filename)", open=False):
-#                import_file = gr.File(label="Choose .csv or .xlsx", file_types=[".csv", ".xlsx"], type="filepath")
-#                import_btn = gr.Button("Import into current session")
-
         with gr.Tab("Batch"):
             with gr.Accordion("Uploaded images", open=True):
-                input_files = gr.File(label="Drop images (or click to select)", file_types=["image"], file_count="multiple",)
-            preview_gallery = gr.Gallery(
-                        label="Preview (un-captioned)",
-                        show_label=True,
-                        columns=5,
-                        height=220,
-                    )
-            input_files.change(on_files_changed, inputs=[input_files], outputs=[preview_gallery])
+                input_files = gr.File(
+                    label="Drop images (or click to select)",
+                    file_types=["image"],
+                    file_count="multiple",
+                    type="filepath"
+                )
             run_button = gr.Button("Caption batch", variant="primary")
 
-        
     # ---- Results area (gallery left / table right)
     rows_state  = gr.State(load_session())
     autosave_md = gr.Markdown("Ready.")
@@ -1043,7 +912,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 headers=["filename", "caption"],
                 interactive=True,
                 wrap=True,
-                type="array",          # prevents pandas truth ambiguity
+                type="array",
                 elem_id="cfTable"
             )
 
@@ -1066,7 +935,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             export_txt_btn  = gr.Button("Export captions as .txt (zip)")
             txt_zip         = gr.File(label="TXT zip", visible=False)
 
-    # ---- Robust scroll sync (works with Gradio v5 Gallery)
+    # ---- Robust scroll sync
     gr.HTML("""
 <script>
 (function () {
@@ -1075,9 +944,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
     if (!host) return null;
     return host.querySelector('[data-testid="gallery"]') || host;
   }
-  function findTbl() {
-    return document.querySelector("#cfTableWrap");
-  }
+  function findTbl() { return document.querySelector("#cfTableWrap"); }
   function syncScroll(a, b) {
     if (!a || !b) return;
     let lock = false;
@@ -1121,7 +988,6 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
         files = files or []
         budget = None if no_limit else float(budget_s)
 
-        # Manual step → process first chunk only
         if mode == "Manual (step)" and files:
             chunks = _split_chunks(files, int(csize))
             batch = chunks[0]
@@ -1135,7 +1001,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             prog = f"Batch progress: {done}/{total} processed in this step • Remaining overall: {len(remaining)}"
             return new_rows, gal, tbl, stamp, remaining, panel_vis, gr.update(value=msg), gr.update(value=prog)
 
-        # Auto 
+        # Auto
         new_rows, gal, tbl, stamp, leftover, done, total = run_batch(
             files, rows or [], instr, t, p, m, int(ms), budget
         )
@@ -1147,21 +1013,9 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
     run_button.click(
         _run_click,
         inputs=[input_files, rows_state, instruction_preview, max_side, chunk_mode, chunk_size, gpu_budget, no_time_limit],
-        outputs=[rows_state, gallery, table, autosave_md, remaining_state, step_panel, step_msg, progress_md],
-    ).then(
-        lambda rows: [(Image.open(r["path"]).convert("RGB"), r["caption"]) for r in rows],
-        inputs=[rows_state],
-        outputs=[gallery],
-    )
-    table.change(
-        sync_table_to_session,
-        inputs=[table, rows_state],
-        outputs=[rows_state, captions_text],
-    ).then(
-        lambda rows: [(Image.open(r["path"]).convert("RGB"), r["caption"]) for r in rows],
-        inputs=[rows_state],
-        outputs=[gallery],
+        outputs=[rows_state, gallery, table, autosave_md, remaining_state, step_panel, step_msg, progress_md]
     )
+
     def _step_next(remain, rows, instr, ms, csize, budget_s, no_limit):
         t, p, m = _tpms()
         remain = remain or []
@@ -1171,7 +1025,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             return (
                 rows,
                 gr.update(value="No files remaining."),
-                gr.update(visible=True),
+                gr.update(visible=False),
                 [],
                 [],
                 [],
@@ -1207,16 +1061,8 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
         gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
                          for r in session_rows if (r.get("thumb_path") or r.get("path"))]
         return session_rows, gallery_pairs, f"Saved • {time.strftime('%H:%M:%S')}"
-    table.change(sync_table_to_session, inputs=[table, rows_state], outputs=[rows_state, gallery, autosave_md])
-
-    def new_session() -> Tuple[List[dict], list, list, str]:
-        return [], [], _rows_to_table([]), ""
 
-    # ---- Import hook
-    def _do_import(fpath, rows):
-        new_rows, gal, tbl, stamp = import_captions_file(fpath, rows or [])
-        return new_rows, gal, tbl, stamp
-    import_btn.click(_do_import, inputs=[import_file, rows_state], outputs=[rows_state, gallery, table, autosave_md])
+    table.change(sync_table_to_session, inputs=[table, rows_state], outputs=[rows_state, gallery, autosave_md])
 
     # ---- Exports
     export_csv_btn.click(
@@ -1234,7 +1080,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
 
 
 # ------------------------------
-# 12) Launch (SSR disabled for stability on Spaces)
+# 12) Launch
 # ------------------------------
 if __name__ == "__main__":
     demo.queue(max_size=64).launch(
@@ -1243,6 +1089,5 @@ if __name__ == "__main__":
         ssr_mode=False,
         debug=True,
         show_error=True,
-        # Allow Gradio to serve generated files from /tmp caches
         allowed_paths=[THUMB_CACHE, EXCEL_THUMB_DIR, TXT_EXPORT_DIR],
     )