Update app.py

app.py

@@ -1,7 +1,10 @@
 import os, io, csv, time, json, hashlib, base64, zipfile, re
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Any
+
+# Caches
+os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
 
-import spaces
 import gradio as gr
 from PIL import Image
 import torch
@@ -10,17 +13,14 @@ from transformers import LlavaForConditionalGeneration, AutoProcessor
 # ────────────────────────────────────────────────────────
 # Paths & caches
 # ────────────────────────────────────────────────────────
-os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
-os.makedirs(os.environ["HF_HOME"], exist_ok=True)
-
 APP_DIR = os.getcwd()
 SESSION_FILE = "/tmp/session.json"
 SETTINGS_FILE = "/tmp/cf_settings.json"
 JOURNAL_FILE = "/tmp/cf_journal.json"
-TXT_EXPORT_DIR = "/tmp/cf_txt"
-THUMB_CACHE = os.path.expanduser("~/.cache/captionforge/thumbs")
+THUMB_CACHE = os.path.expanduser("~/.cache/forgecaptions/thumbs")
+EXCEL_THUMB_DIR = "/tmp/forge_excel_thumbs"
 os.makedirs(THUMB_CACHE, exist_ok=True)
-os.makedirs(TXT_EXPORT_DIR, exist_ok=True)
+os.makedirs(EXCEL_THUMB_DIR, exist_ok=True)
 
 # ────────────────────────────────────────────────────────
 # Model setup
@@ -35,8 +35,8 @@ def _detect_gpu():
 
 BACKEND, VRAM_GB, GPU_NAME = _detect_gpu()
 DTYPE   = torch.bfloat16 if BACKEND == "cuda" else torch.float32
-MAX_SIDE_CAP = 1024 if BACKEND == "cuda" else 640
 DEVICE  = "cuda" if BACKEND == "cuda" else "cpu"
+MAX_SIDE_CAP = 1024 if BACKEND == "cuda" else 640
 
 processor = AutoProcessor.from_pretrained(MODEL_PATH)
 model = LlavaForConditionalGeneration.from_pretrained(
@@ -47,8 +47,8 @@ model = LlavaForConditionalGeneration.from_pretrained(
 )
 model.eval()
 
-print(f"[CaptionForge] Backend={BACKEND} GPU={GPU_NAME} VRAM={VRAM_GB}GB dtype={DTYPE}")
-print(f"[CaptionForge] Gradio version: {gr.__version__}")
+print(f"[ForgeCaptions] Backend={BACKEND} GPU={GPU_NAME} VRAM={VRAM_GB}GB dtype={DTYPE}")
+print(f"[ForgeCaptions] Gradio version: {gr.__version__}")
 
 # ────────────────────────────────────────────────────────
 # Instruction templates & options
@@ -66,8 +66,8 @@ STYLE_OPTIONS = [
 ]
 
 CAPTION_TYPE_MAP = {
-    "Descriptive (short)": "One sentence (≤25 words) describing the most important visible elements only. No speculation.",
-    "Descriptive (long)": "Write a detailed description for this image.",
+    "Descriptive (short)": "Write a short describition of the most important visible elements only. No speculation.",
+    "Descriptive (long)": "Write a long, highly detailed description for this image.",
 
     "Character training (short)": (
         "Output a concise, prompt-like caption for character LoRA/ID training. "
@@ -103,7 +103,7 @@ CAPTION_TYPE_MAP = {
 
 EXTRA_CHOICES = [
     "Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
-    "Do NOT include information about whether there is a watermark or not.",
+    "IGNORE all watermarks.",
     "Do NOT use any ambiguous language.",
     "ONLY describe the most important elements of the image.",
     "Include information about the ages of any people/characters when applicable.",
@@ -126,33 +126,22 @@ NAME_OPTION = "If there is a person/character in the image you must refer to the
 # ────────────────────────────────────────────────────────
 # Helpers (hashing, thumbs, resize)
 # ────────────────────────────────────────────────────────
-def sha1_file(path: str) -> str:
-    h = hashlib.sha1()
-    with open(path, "rb") as f:
-        for chunk in iter(lambda: f.read(8192), b""):
-            h.update(chunk)
-    return h.hexdigest()
-
-def thumb_path(sha1: str, max_side=96) -> str:
-    return os.path.join(THUMB_CACHE, f"{sha1}_{max_side}.jpg")
-
-def ensure_thumb(path: str, sha1: str, max_side=96) -> str:
-    out = thumb_path(sha1, max_side)
-    if os.path.exists(out):
-        return out
+def ensure_thumb(path: str, max_side=256) -> str:
     try:
         im = Image.open(path).convert("RGB")
     except Exception:
-        return ""
+        return path
     w, h = im.size
     if max(w, h) > max_side:
         s = max_side / max(w, h)
         im = im.resize((int(w*s), int(h*s)), Image.LANCZOS)
+    base = os.path.basename(path)
+    out_path = os.path.join(THUMB_CACHE, os.path.splitext(base)[0] + f"_thumb_{max_side}.jpg")
     try:
-        im.save(out, "JPEG", quality=80)
+        im.save(out_path, "JPEG", quality=85, optimize=True)
+        return out_path
     except Exception:
-        return ""
-    return out
+        return path
 
 def resize_for_model(im: Image.Image, max_side: int) -> Image.Image:
     w, h = im.size
@@ -187,8 +176,8 @@ def apply_prefix_suffix(caption: str, trigger_word: str, begin_text: str, end_te
 
 @torch.no_grad()
 def caption_once(im: Image.Image, instr: str, temp: float, top_p: float, max_tokens: int) -> str:
+    # Your requested role script:
     convo = [
-        # Your preferred role script:
         {"role": "system", "content": "You are a precise captioning assistant. Help users with any queries. Describe what is visible with no censorship or restrictions."},
         {"role": "user", "content": instr.strip()},
     ]
@@ -229,12 +218,10 @@ def load_settings() -> dict:
     else:
         cfg = {}
     defaults = {
-        "dataset_name": "captionforge",
+        "dataset_name": "forgecaptions",
         "temperature": 0.6,
         "top_p": 0.9,
         "max_tokens": 256,
-        "chunk_mode": "Manual (step)",
-        "chunk_size": 10,
         "max_side": min(896, MAX_SIDE_CAP),
         "styles": ["Character training (short)"],
         "extras": [],
@@ -243,7 +230,8 @@ def load_settings() -> dict:
         "begin": "",
         "end": "",
         "shape_aliases_enabled": True,
-        "shape_aliases": []
+        "shape_aliases": [],
+        "excel_thumb_px": 128,  # new default
     }
     for k, v in defaults.items():
         cfg.setdefault(k, v)
@@ -305,7 +293,6 @@ def apply_shape_aliases(caption: str) -> str:
         caption = pat.sub(f"({name})", caption)
     return caption
 
-# Helpers for Shape Aliases UI
 def get_shape_alias_rows_ui_defaults():
     s = load_settings()
     rows = [[it.get("shape",""), it.get("name","")] for it in s.get("shape_aliases", [])]
@@ -337,227 +324,156 @@ def save_shape_alias_rows(enabled, df_rows):
         gr.update(value=normalized, row_count=(max(1, len(normalized)), "dynamic"))
     )
 
-# ────────────────────────────────────────────────────────
-# Logo loader
-# ────────────────────────────────────────────────────────
-def logo_b64_img() -> str:
-    candidates = [
-        os.path.join(APP_DIR, "captionforge-logo.png"),
-        "/home/user/app/captionforge-logo.png",
-        "captionforge-logo.png",
-    ]
-    for p in candidates:
-        if os.path.exists(p):
-            with open(p, "rb") as f:
-                b64 = base64.b64encode(f.read()).decode("ascii")
-            return f"<img src='data:image/png;base64,{b64}' alt='CaptionForge' class='cf-logo'>"
-    return ""
-
 # ────────────────────────────────────────────────────────
 # Import / Export helpers
 # ────────────────────────────────────────────────────────
-def sanitize_basename(s: str) -> str:
-    s = (s or "").strip() or "captionforge"
-    return re.sub(r"[^A-Za-z0-9._-]+", "_", s)[:120]
-
-def ts_stamp() -> str:
-    return time.strftime("%Y%m%d_%H%M%S")
-
-def export_csv(rows: List[dict], basename: str) -> str:
-    base = sanitize_basename(basename)
-    out = f"/tmp/{base}_{ts_stamp()}.csv"
+def export_csv_from_table(table_value: Any) -> str:
+    data = table_value or []
+    out = f"/tmp/forgecaptions_{int(time.time())}.csv"
     with open(out, "w", newline="", encoding="utf-8") as f:
         w = csv.writer(f)
-        w.writerow(["sha1", "filename", "caption"])
-        for r in rows:
-            w.writerow([r.get("sha1",""), r.get("filename",""), r.get("caption","")])
+        w.writerow(["filename", "caption"])
+        w.writerows(data)
     return out
 
-def export_excel(rows: List[dict], basename: str) -> str:
+def _resize_for_excel(path: str, px: int) -> str:
+    """Create a temp resized copy for Excel embedding."""
+    try:
+        im = Image.open(path).convert("RGB")
+    except Exception:
+        return path
+    w, h = im.size
+    if max(w, h) > px:
+        s = px / max(w, h)
+        im = im.resize((int(w*s), int(h*s)), Image.LANCZOS)
+    base = os.path.basename(path)
+    out_path = os.path.join(EXCEL_THUMB_DIR, f"{os.path.splitext(base)[0]}_{px}px.jpg")
+    try:
+        im.save(out_path, "JPEG", quality=85, optimize=True)
+        return out_path
+    except Exception:
+        return path
+
+def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_px: int) -> str:
     try:
         from openpyxl import Workbook
         from openpyxl.drawing.image import Image as XLImage
     except Exception as e:
         raise RuntimeError("Excel export requires 'openpyxl' in requirements.txt.") from e
-    base = sanitize_basename(basename)
-    out = f"/tmp/{base}_{ts_stamp()}.xlsx"
-    wb = Workbook(); ws = wb.active; ws.title = "CaptionForge"
+
+    # Respect user edits (table wins)
+    caption_by_file = {}
+    for row in (table_value or []):
+        if not row:
+            continue
+        fn = str(row[0]) if len(row) > 0 else ""
+        cap = str(row[1]) if len(row) > 1 and row[1] is not None else ""
+        if fn:
+            caption_by_file[fn] = cap
+
+    wb = Workbook(); ws = wb.active; ws.title = "ForgeCaptions"
     ws.append(["image", "filename", "caption"])
     ws.column_dimensions["A"].width = 24
     ws.column_dimensions["B"].width = 42
     ws.column_dimensions["C"].width = 100
+
+    # Approx px→points (Excel row height is points; ~0.75 pt per px @ 96dpi)
+    row_h = int(thumb_px * 0.75)
+
     r_i = 2
-    for r in rows:
-        fn = r.get("filename",""); cap = r.get("caption","")
+    for r in (session_rows or []):
+        fn = r.get("filename","")
+        cap = caption_by_file.get(fn, r.get("caption",""))
         ws.cell(row=r_i, column=2, value=fn)
         ws.cell(row=r_i, column=3, value=cap)
         img_path = r.get("thumb_path") or r.get("path")
         if img_path and os.path.exists(img_path):
             try:
-                ximg = XLImage(img_path); ws.add_image(ximg, f"A{r_i}"); ws.row_dimensions[r_i].height = 110
+                resized = _resize_for_excel(img_path, int(thumb_px))
+                xlimg = XLImage(resized)
+                ws.add_image(xlimg, f"A{r_i}")
+                ws.row_dimensions[r_i].height = row_h
             except Exception:
                 pass
         r_i += 1
-    wb.save(out); return out
-
-def export_txt_zip(rows: List[dict], basename: str) -> str:
-    base = sanitize_basename(basename)
-    for name in os.listdir(TXT_EXPORT_DIR):
-        try: os.remove(os.path.join(TXT_EXPORT_DIR, name))
-        except Exception: pass
-    used: Dict[str,int] = {}
-    for r in rows:
-        orig = r.get("filename") or r.get("sha1") or "item"
-        stem = re.sub(r"\.[A-Za-z0-9]+$", "", orig)
-        stem = sanitize_basename(stem)
-        if stem in used:
-            used[stem] += 1
-            stem = f"{stem}_{used[stem]}"
-        else:
-            used[stem] = 0
-        txt_path = os.path.join(TXT_EXPORT_DIR, f"{stem}.txt")
-        with open(txt_path, "w", encoding="utf-8") as f:
-            f.write(r.get("caption", ""))
-    zpath = f"/tmp/{base}_{ts_stamp()}_txt.zip"
-    with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as z:
-        for name in os.listdir(TXT_EXPORT_DIR):
-            if name.endswith(".txt"):
-                z.write(os.path.join(TXT_EXPORT_DIR, name), arcname=name)
-    return zpath
-
-def import_csv_jsonl(path: str, rows: List[dict]) -> List[dict]:
-    """Merge captions from CSV/JSONL into current rows by sha1→filename, updating only 'caption'."""
-    if not path or not os.path.exists(path):
-        return rows or []
-    rows = rows or []
-    by_sha = {r.get("sha1"): r for r in rows if r.get("sha1")}
-    by_fn  = {r.get("filename"): r for r in rows if r.get("filename")}
-
-    def _apply(item: dict):
-        if not isinstance(item, dict):
-            return
-        cap = (item.get("caption") or "").strip()
-        if not cap:
-            return
-        sha = (item.get("sha1") or "").strip()
-        fn  = (item.get("filename") or "").strip()
-        tgt = by_sha.get(sha) if sha else None
-        if not tgt and fn:
-            tgt = by_fn.get(fn)
-        if tgt is not None:
-            prev = tgt.get("caption", "")
-            if cap != prev:
-                tgt.setdefault("history", []).append(prev)
-                tgt["caption"] = cap
-
-    ext = os.path.splitext(path)[1].lower()
-    try:
-        if ext == ".csv":
-            with open(path, "r", encoding="utf-8") as f:
-                for row in csv.DictReader(f):
-                    _apply(row)
-        elif ext in (".jsonl", ".ndjson"):
-            with open(path, "r", encoding="utf-8") as f:
-                for line in f:
-                    line = line.strip()
-                    if not line:
-                        continue
-                    try:
-                        obj = json.loads(line)
-                    except Exception:
-                        continue
-                    _apply(obj)
-        else:
-            with open(path, "r", encoding="utf-8") as f:
-                data = json.load(f)
-            if isinstance(data, list):
-                for obj in data:
-                    _apply(obj)
-    except Exception as e:
-        print("[CaptionForge] Import merge failed:", e)
-    save_session(rows)
-    return rows
+
+    out = f"/tmp/forgecaptions_{int(time.time())}.xlsx"
+    wb.save(out)
+    return out
+
+# Rows<->Table helpers
+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]:
+    tbl = table_value or []
+    new = []
+    for i, r in enumerate(rows or []):
+        r = dict(r)
+        if i < len(tbl) and len(tbl[i]) >= 2:
+            r["filename"] = str(tbl[i][0]) if tbl[i][0] is not None else r.get("filename","")
+            r["caption"]  = str(tbl[i][1]) if tbl[i][1] is not None else r.get("caption","")
+        new.append(r)
+    return new
 
 # ────────────────────────────────────────────────────────
-# Batching / processing
+# Batch captioning (returns rows, gallery, table)
 # ────────────────────────────────────────────────────────
-def adaptive_defaults(vram_gb: int) -> Tuple[int, int]:
-    if vram_gb >= 40: return 24, min(1024, MAX_SIDE_CAP)
-    if vram_gb >= 24: return 16, min(1024, MAX_SIDE_CAP)
-    if vram_gb >= 12: return 10, 896
-    if vram_gb >= 8:  return 6, 768
-    return 4, 640
-
-def warmup_if_needed(temp: float, top_p: float, max_tokens: int, max_side: int):
-    try:
-        im = Image.new("RGB", (min(128, max_side), min(128, max_side)), (127,127,127))
-        _ = caption_once(im, "Warm up.", temp=0.0, top_p=top_p, max_tokens=min(16, max_tokens))
-    except Exception as e:
-        print("[CaptionForge] Warm-up skipped:", e)
-
-@spaces.GPU()
 @torch.no_grad()
-def process_batch(files: List[str], include_dups: bool, rows: List[dict],
-                  instr_text: str, temp: float, top_p: float, max_tokens: int,
-                  trigger_word: str, begin_text: str, end_text: str,
-                  mode: str, chunk_size: int, max_side: int, step_once: bool,
-                  progress=gr.Progress(track_tqdm=True)) -> Tuple[List[dict], List[str], List[str]]:
+def run_batch(
+    files: List[Any],
+    session_rows: List[dict],
+    instr_text: str,
+    temp: float,
+    top_p: float,
+    max_tokens: int,
+    max_side: int,
+) -> Tuple[List[dict], list, list, str]:
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    files = files or []
-    if not files: return rows, [], []
 
-    auto_mode = mode == "Auto"
-    if auto_mode:
-        chunk_size, max_side = adaptive_defaults(VRAM_GB)
+    session_rows = session_rows or []
+    files = files or []
+    if not files:
+        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, _rows_to_table(session_rows), f"Saved • {time.strftime('%H:%M:%S')}"
 
-    sha_seen = {r.get("sha1") for r in rows}
-    f_infos = []
     for f in files:
+        path = f if isinstance(f, str) else getattr(f, "name", None) or getattr(f, "path", None)
+        if not path or not os.path.exists(path):
+            continue
         try:
-            s = sha1_file(f)
+            im = Image.open(path).convert("RGB")
         except Exception:
             continue
-        if (not include_dups) and (s in sha_seen):
-            continue
-        f_infos.append((f, os.path.basename(f), s))
-    if not f_infos: return rows, [], []
-
-    warmup_if_needed(temp, top_p, max_tokens, max_side)
-
-    chunks = [f_infos[i:i+chunk_size] for i in range(0, len(f_infos), chunk_size)]
-    if step_once:
-        chunks = chunks[:1]
-
-    error_log: List[str] = []
-    remaining: List[str] = []
-
-    for ci, chunk in enumerate(chunks):
-        for (f, fname, sha) in progress.tqdm(chunk, desc=f"Chunk {ci+1}/{len(chunks)}"):
-            thumb = ensure_thumb(f, sha, 96)
-            attempt = 0
-            cap = None
-            while attempt < 2:
-                attempt += 1
-                try:
-                    im = Image.open(f).convert("RGB")
-                    im = resize_for_model(im, max_side)
-                    cap = caption_once(im, instr_text, temp, top_p, max_tokens)
-                    cap = apply_shape_aliases(cap)  # Apply custom shape→name mappings
-                    cap = apply_prefix_suffix(cap, trigger_word, begin_text, end_text)
-                    break
-                except Exception as e:
-                    error_log.append(f"{fname} ({sha[:8]}): {type(e).__name__}: {e} [attempt {attempt}]")
-            rows.append({
-                "sha1": sha, "filename": fname, "path": f, "thumb_path": thumb,
-                "caption": cap or "", "history": []
-            })
-        save_session(rows)
-
-    if step_once and len(f_infos) > len(chunks[0]):
-        remaining = [fi[0] for fi in f_infos[len(chunks[0]):]]
-    save_journal({"remaining_files": remaining})
-    return rows, remaining, error_log
+        im = resize_for_model(im, max_side)
+        cap = caption_once(im, instr_text, temp, top_p, max_tokens)
+        cap = apply_shape_aliases(cap)
+        s = load_settings()
+        cap = apply_prefix_suffix(cap, s.get("trigger",""), s.get("begin",""), s.get("end",""))
+
+        filename = os.path.basename(path)
+        thumb = ensure_thumb(path, 256)
+        session_rows.append({"filename": filename, "caption": cap, "path": path, "thumb_path": thumb})
+
+    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"))
+    ]
+    return session_rows, gallery_pairs, _rows_to_table(session_rows), f"Saved • {time.strftime('%H:%M:%S')}"
+
+def sync_table_to_session(table_value: Any, session_rows: List[dict]) -> Tuple[List[dict], list, str]:
+    session_rows = _table_to_rows(table_value, session_rows or [])
+    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"))
+    ]
+    return session_rows, gallery_pairs, f"Saved • {time.strftime('%H:%M:%S')}"
 
 # ────────────────────────────────────────────────────────
 # UI
@@ -580,12 +496,24 @@ BASE_CSS = """
 .cf-scroll{max-height:70vh; overflow-y:auto; border:1px solid #e6e6e6; border-radius:10px; padding:8px}
 /* Uniform sizes */
 #cfGal .grid > div { height: 96px; }
-/* Hide SHA1 in table (first column) */
-#cfTable table thead tr th:first-child { display:none; }
-#cfTable table tbody tr td:first-child { display:none; }
 """
 
-with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
+def logo_b64_img() -> str:
+    candidates = [
+        os.path.join(APP_DIR, "forgecaptions-logo.png"),
+        os.path.join(APP_DIR, "captionforge-logo.png"),  # fallback if you kept the old name
+        "/home/user/app/forgecaptions-logo.png",
+        "forgecaptions-logo.png",
+        "captionforge-logo.png",
+    ]
+    for p in candidates:
+        if os.path.exists(p):
+            with open(p, "rb") as f:
+                b64 = base64.b64encode(f.read()).decode("ascii")
+            return f"<img src='data:image/png;base64,{b64}' alt='ForgeCaptions' class='cf-logo'>"
+    return ""
+
+with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
     settings = load_settings()
     settings["styles"] = [s for s in settings.get("styles", []) if s in STYLE_OPTIONS] or ["Character training (short)"]
 
@@ -593,10 +521,10 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
 <div class="cf-hero">
   {logo_b64_img()}
   <div>
-    <h1 class="cf-title">CaptionForge</h1>
+    <h1 class="cf-title">ForgeCaptions</h1>
     <div class="cf-sub">Batch captioning</div>
     <div class="cf-sub">Scrollable editor & autosave</div>
-    <div class="cf-sub">CSV / Excel / TXT export</div>
+    <div class="cf-sub">CSV / Excel export</div>
   </div>
 </div>
 <hr>
@@ -619,123 +547,116 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
                     )
                 with gr.Accordion("Name & Prefix/Suffix", open=True):
                     name_input = gr.Textbox(label="Person / Character Name", value=settings.get("name", ""))
-                    trig       = gr.Textbox(label="Trigger word")
-                    add_start  = gr.Textbox(label="Add text to start")
-                    add_end    = gr.Textbox(label="Add text to end")
+                    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",""))
 
             with gr.Column(scale=1):
-                instruction_preview = gr.Textbox(label="Model Instructions", lines=14)
-                dataset_name = gr.Textbox(label="Dataset name (used for export file titles)", value=settings.get("dataset_name", "captionforge"))
+                instruction_preview = gr.Textbox(label="Model Instructions", lines=12)
+                dataset_name = gr.Textbox(label="Dataset name (used for export file titles)", value=settings.get("dataset_name", "forgecaptions"))
+                max_side   = gr.Slider(256, MAX_SIDE_CAP, settings.get("max_side", min(896, MAX_SIDE_CAP)), step=32, label="Max side (resize)")
+                excel_thumb_px = gr.Slider(64, 256, value=settings.get("excel_thumb_px", 128), step=8, label="Excel thumbnail size (px)")
+                gr.Markdown("Generation (saved in settings): temperature 0.6 • top-p 0.9 • max_tokens 256")
 
-        with gr.Row():
-            chunk_mode = gr.Radio(
-                choices=["Auto", "Manual (all at once)", "Manual (step)"],
-                value=settings.get("chunk_mode", "Manual (step)"),
-                label="Batch mode"
-            )
-            chunk_size = gr.Slider(1, 50, settings.get("chunk_size", 10), step=1, label="Chunk size")
-            max_side   = gr.Slider(256, MAX_SIDE_CAP, settings.get("max_side", min(896, MAX_SIDE_CAP)), step=32, label="Max side (resize)")
+    # Auto-refresh instruction & persist key controls
+    def _refresh_instruction(styles, extra, name_value, trigv, begv, endv, excel_px, ms):
+        instr = final_instruction(styles or ["Character training (short)"], extra or [], name_value)
+        cfg = load_settings()
+        cfg.update({
+            "styles": styles or ["Character training (short)"],
+            "extras": extra or [],
+            "name": name_value,
+            "trigger": trigv, "begin": begv, "end": endv,
+            "excel_thumb_px": int(excel_px),
+            "max_side": int(ms),
+        })
+        save_settings(cfg)
+        return instr
+
+    for comp in [style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side]:
+        comp.change(
+            _refresh_instruction,
+            inputs=[style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side],
+            outputs=[instruction_preview]
+        )
+
+    # Set initial instruction text on load
+    demo.load(lambda s,e,n: final_instruction(s or ["Character training (short)"], e or [], n),
+              inputs=[style_checks, extra_opts, name_input], outputs=[instruction_preview])
 
-    # === Shape Aliases UI (Accordion) ===
+    # ===== Shape Aliases
     with gr.Accordion("Shape Aliases", open=False):
         gr.Markdown(
             "### 🔷 Shape Aliases\n"
             "Replace literal **shape tokens** in captions with a preferred **name**.\n\n"
-            "**How to use:**\n"
-            "1) Left column: exact shape token (e.g., `diamond`, `heart`).\n"
-            "2) Right column: replacement name (e.g., `starkey-emblem`).\n"
-            "3) Click **Save**.\n\n"
-            "_Rows with blanks are ignored; add as many rows as you like._"
+            "**How to use:** left column = shape token (e.g., `diamond`), right column = replacement name (e.g., `starkey-emblem`)."
         )
-
         init_rows, init_enabled = get_shape_alias_rows_ui_defaults()
-
-        enable_aliases = gr.Checkbox(
-            label="Enable shape alias replacements",
-            value=init_enabled,
-        )
-
+        enable_aliases = gr.Checkbox(label="Enable shape alias replacements", value=init_enabled)
         alias_table = gr.Dataframe(
             headers=["shape (literal token)", "name to insert"],
-            value=init_rows,                               # set value at creation time
-            col_count=(2, "fixed"),                        # exactly two columns
-            row_count=(max(1, len(init_rows)), "dynamic"), # user can add/remove rows
+            value=init_rows,
+            col_count=(2, "fixed"),
+            row_count=(max(1, len(init_rows)), "dynamic"),
             datatype=["str", "str"],
-            type="array",                                  # treat rows as simple lists
+            type="array",
             interactive=True
         )
-
         with gr.Row():
             add_row_btn = gr.Button("＋ Add row", variant="secondary")
             clear_btn   = gr.Button("Clear", variant="secondary")
             save_btn    = gr.Button("💾 Save", variant="primary")
-
         save_status = gr.Markdown("")
-
-        # Button handlers
         def _add_row(cur):
             cur = (cur or []) + [["", ""]]
             return gr.update(value=cur, row_count=(max(1, len(cur)), "dynamic"))
-
         def _clear_rows():
             return gr.update(value=[["", ""]], row_count=(1, "dynamic"))
-
         add_row_btn.click(_add_row, inputs=[alias_table], outputs=[alias_table])
         clear_btn.click(_clear_rows, outputs=[alias_table])
-        save_btn.click(save_shape_alias_rows, inputs=[enable_aliases, alias_table],
-                       outputs=[save_status, alias_table])
-
-    # Auto-refresh instruction & persist key controls
-    def _refresh_instruction(styles, extra, name_value, trigv, begv, endv):
-        instr = final_instruction(styles or ["Character training (short)"], extra or [], name_value)
-        cfg = load_settings()
-        cfg.update({
-            "styles": styles or ["Character training (short)"],
-            "extras": extra or [],
-            "name": name_value,
-            "trigger": trigv, "begin": begv, "end": endv
-        })
-        save_settings(cfg)
-        return instr
+        save_btn.click(save_shape_alias_rows, inputs=[enable_aliases, alias_table], outputs=[save_status, alias_table])
+
+    # ===== Tabs: Single & Batch (keeps gallery/table position below)
+    with gr.Tabs():
+        with gr.Tab("Single"):
+            input_image_single = gr.Image(type="pil", label="Input Image", height=512, width=512)
+            single_caption_btn = gr.Button("Caption")
+            single_caption_out = gr.Textbox(label="Caption (single)")
+
+            def _caption_single(img, instr):
+                if img is None:
+                    return "No image provided."
+                s = load_settings()
+                im = resize_for_model(img, int(s.get("max_side", MAX_SIDE_CAP)))
+                t = s.get("temperature", 0.6)
+                p = s.get("top_p", 0.9)
+                m = s.get("max_tokens", 256)
+                cap = caption_once(im, instr, t, p, m)
+                cap = apply_shape_aliases(cap)
+                cap = apply_prefix_suffix(cap, s.get("trigger",""), s.get("begin",""), s.get("end",""))
+                return cap
+
+            single_caption_btn.click(
+                _caption_single,
+                inputs=[input_image_single, instruction_preview],
+                outputs=[single_caption_out]
+            )
 
-    for comp in [style_checks, extra_opts, name_input, trig, add_start, add_end]:
-        comp.change(_refresh_instruction, inputs=[style_checks, extra_opts, name_input, trig, add_start, add_end], outputs=[instruction_preview])
+        with gr.Tab("Batch"):
+            with gr.Accordion("Uploaded images", open=True):
+                input_files = gr.File(label="Drop images", file_types=["image"], file_count="multiple", type="filepath")
+            run_button = gr.Button("Caption batch", variant="primary")
 
-    # ===== File inputs & actions
-    with gr.Accordion("Uploaded images", open=True) as uploads_acc:
-        input_files = gr.File(label="Drop images", file_types=["image"], file_count="multiple", type="filepath")
-        import_file = gr.File(label="Import CSV or JSONL (merge captions)")
+    # ===== Results/Table (position unchanged)
+    rows_state  = gr.State(load_session())
+    autosave_md = gr.Markdown("Ready.")
 
-    with gr.Row():
-        import_btn = gr.Button("Import → Merge")
-        resume_btn = gr.Button("Resume last run", variant="secondary")
-        run_button = gr.Button("Caption batch", variant="primary")
-
-    # Step-chunk controls
-    step_panel = gr.Group(visible=False)
-    with step_panel:
-        step_msg = gr.Markdown("")
-        step_next = gr.Button("Process next chunk")
-        step_finish = gr.Button("Finish")
-
-    # States
-    rows_state          = gr.State(load_session())
-    visible_count_state = gr.State(100)
-    pending_files       = gr.State([])
-    remaining_state     = gr.State([])
-    autosave_md         = gr.Markdown("Ready.")
-
-    # Error log
-    with gr.Accordion("Error log", open=False):
-        log_md = gr.Markdown("")
-
-    # Side-by-side gallery + table with shared scroll
     with gr.Row():
         with gr.Column(scale=1):
             gallery = gr.Gallery(
-                label="Results (images)",
+                label="Results (image + caption)",
                 show_label=True,
-                columns=10,
+                columns=3,
                 height=520,
                 elem_id="cfGal",
                 elem_classes=["cf-scroll"]
@@ -743,49 +664,30 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
         with gr.Column(scale=1):
             table = gr.Dataframe(
                 label="Editable captions (whole session)",
-                value=[],
-                headers=["sha1", "filename", "caption"],  # sha1 hidden via CSS
+                value=_rows_to_table(load_session()),
+                headers=["filename", "caption"],
                 interactive=True,
+                wrap=True,
                 elem_id="cfTable",
                 elem_classes=["cf-scroll"]
             )
 
-    # Exports aligned under buttons
+    # Exports
     with gr.Row():
         with gr.Column():
             export_csv_btn  = gr.Button("Export CSV")
             csv_file        = gr.File(label="CSV file", visible=False)
         with gr.Column():
-            export_xlsx_btn = gr.Button("Export Excel (.xlsx)")
+            export_xlsx_btn = gr.Button("Export Excel (.xlsx) with thumbnails")
             xlsx_file       = gr.File(label="Excel file", visible=False)
-        with gr.Column():
-            export_txt_btn  = gr.Button("Export captions as .txt (zip)")
-            txt_zip         = gr.File(label="TXT zip", visible=False)
 
-    # Render helpers
-    def _render(rows, vis_n):
+    # Initial gallery render
+    def _initial_gallery(rows):
         rows = rows or []
-        n = max(0, int(vis_n) if vis_n else 0)
-        win = rows[:n]
-
-        # Gallery expects list of paths
-        gal = []
-        for r in win:
-            p = r.get("thumb_path") or r.get("path")
-            if isinstance(p, str) and os.path.exists(p):
-                gal.append(p)
-
-        # Dataframe rows (sha1 kept for mapping; hidden via CSS)
-        df = [[str(r.get("sha1","")), str(r.get("filename","")), str(r.get("caption",""))] for r in win]
-        info = f"Showing {len(win)} of {len(rows)}"
-        return gal, df, info
-
-    # Initial loads
-    demo.load(lambda s,e,n: final_instruction(s or ["Character training (short)"], e or [], n),
-              inputs=[style_checks, extra_opts, name_input], outputs=[instruction_preview])
-    demo.load(_render, inputs=[rows_state, visible_count_state], outputs=[gallery, table, autosave_md])
+        return [((r.get("thumb_path") or r.get("path")), r.get("caption","")) for r in rows if (r.get("thumb_path") or r.get("path"))]
+    demo.load(_initial_gallery, inputs=[rows_state], outputs=[gallery])
 
-    # Scroll sync (Gallery + Table)
+    # Scroll sync
     gr.HTML("""
 <script>
 (function () {
@@ -830,123 +732,36 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
 </script>
 """)
 
-    # Table edits -> save + rerender
-    def on_table_edit(df_rows, rows):
-        by_sha = {r.get("sha1"): r for r in (rows or [])}
-        for row in (df_rows or []):
-            if not row:
-                continue
-            sha = str(row[0]) if len(row) > 0 else ""
-            tgt = by_sha.get(sha)
-            if not tgt:
-                continue
-            prev = tgt.get("caption", "")
-            newc = row[2] if len(row) > 2 else prev
-            if newc != prev:
-                hist = tgt.setdefault("history", [])
-                hist.append(prev)
-                tgt["caption"] = newc
-        save_session(rows or [])
-        return rows, f"Saved • {time.strftime('%H:%M:%S')}"
-    table.change(on_table_edit, inputs=[table, rows_state], outputs=[rows_state, autosave_md])\
-         .then(_render, inputs=[rows_state, visible_count_state], outputs=[gallery, table, autosave_md])
-
-    # Import merge
-    def do_import(file, rows):
-        if not file:
-            return rows, gr.update(value="No file"), rows
-        rows = import_csv_jsonl(file.name, rows or [])
-        return rows, gr.update(value=f"Merged from {os.path.basename(file.name)} at {time.strftime('%H:%M:%S')}"), rows
-    import_btn.click(do_import, inputs=[import_file, rows_state], outputs=[rows_state, autosave_md, rows_state])\
-              .then(_render, inputs=[rows_state, visible_count_state], outputs=[gallery, table, autosave_md])
-
-    # Prepare run (hash dupes)
-    def prepare_run(files, rows, pending):
-        files = files or []
-        rows = rows or []
-        pending = pending or []
-        if not files:
-            return [], gr.update(open=True), []
-        existing = {r.get("sha1") for r in rows if r.get("sha1")}
-        pending_sha = set()
-        for f in pending:
-            try:
-                pending_sha.add(sha1_file(f))
-            except Exception:
-                pass
-        keep = []
-        for f in files:
-            try:
-                s = sha1_file(f)
-            except Exception:
-                continue
-            if s not in existing and s not in pending_sha:
-                keep.append(f)
-        return keep, gr.update(open=False), keep
-    input_files.change(prepare_run, inputs=[input_files, rows_state, pending_files], outputs=[pending_files, uploads_acc, pending_files])
-
-    # Step visibility helper (dynamic message)
-    def _step_visibility(remain, mode):
-        if (mode == "Manual (step)") and remain:
-            return gr.update(visible=True), gr.update(value=f"{len(remain)} files remain. Process next chunk?")
-        return gr.update(visible=False), gr.update(value="")
-
-    # Run helpers (use settings for gen params)
-    def _run_once(pending, rows, instr, trigv, begv, endv, mode, csize, mside):
-        cfg = load_settings()
-        t = cfg.get("temperature", 0.6)
-        p = cfg.get("top_p", 0.9)
-        m = cfg.get("max_tokens", 256)
-        step_once = (mode == "Manual (step)")
-        new_rows, remaining, errors = process_batch(
-            pending, False, rows or [], instr, t, p, m, trigv, begv, endv,
-            mode, int(csize), int(mside), step_once
-        )
-        log = "\n".join(errors) if errors else "(no errors)"
-        return new_rows, remaining, log, f"Saved • {time.strftime('%H:%M:%S')}"
+    # Run batch
+    def _run_click(files, rows, instr, ms):
+        s = load_settings()
+        t = s.get("temperature", 0.6)
+        p = s.get("top_p", 0.9)
+        m = s.get("max_tokens", 256)
+        new_rows, gal, tbl, stamp = run_batch(files, rows or [], instr, t, p, m, int(ms))
+        return new_rows, gal, tbl, stamp
 
     run_button.click(
-        _run_once,
-        inputs=[pending_files, rows_state, instruction_preview, trig, add_start, add_end, chunk_mode, chunk_size, max_side],
-        outputs=[rows_state, remaining_state, log_md, autosave_md]
-    ).then(
-        _render, inputs=[rows_state, visible_count_state], outputs=[gallery, table, autosave_md]
-    ).then(
-        _step_visibility, inputs=[remaining_state, chunk_mode], outputs=[step_panel, step_msg]
+        _run_click,
+        inputs=[input_files, rows_state, instruction_preview, max_side],
+        outputs=[rows_state, gallery, table, autosave_md]
     )
 
-    def _resume(rows):
-        j = load_journal()
-        rem = j.get("remaining_files", [])
-        return rem, ("Nothing to resume" if not rem else f"Loaded {len(rem)} remaining")
-    resume_btn.click(_resume, inputs=[rows_state], outputs=[pending_files, autosave_md])
-
-    def _step_next(remain, rows, instr, trigv, begv, endv, mode, csize, mside):
-        return _run_once(remain, rows, instr, trigv, begv, endv, mode, csize, mside)
-    step_next.click(
-        _step_next,
-        inputs=[remaining_state, rows_state, instruction_preview, trig, add_start, add_end, chunk_mode, chunk_size, max_side],
-        outputs=[rows_state, remaining_state, log_md, autosave_md]
-    ).then(
-        _render, inputs=[rows_state, visible_count_state], outputs=[gallery, table, autosave_md]
-    ).then(
-        _step_visibility, inputs=[remaining_state, chunk_mode], outputs=[step_panel, step_msg]
+    # Table edits sync → rows + gallery
+    table.change(
+        sync_table_to_session,
+        inputs=[table, rows_state],
+        outputs=[rows_state, gallery, autosave_md]
     )
 
-    step_finish.click(lambda: (gr.update(visible=False), gr.update(value=""), []), outputs=[step_panel, step_msg, remaining_state])
-
-    # Exports (reveal file widgets when created)
+    # Exports (use slider value)
     export_csv_btn.click(
-        lambda rows, base: (export_csv(rows or [], base), gr.update(visible=True)),
-        inputs=[rows_state, dataset_name], outputs=[csv_file, csv_file]
+        lambda tbl: (export_csv_from_table(tbl), gr.update(visible=True)),
+        inputs=[table], outputs=[csv_file, csv_file]
     )
     export_xlsx_btn.click(
-        lambda rows, base: (export_excel(rows or [], base), gr.update(visible=True)),
-        inputs=[rows_state, dataset_name], outputs=[xlsx_file, xlsx_file]
-    )
-    export_txt_btn.click(
-        lambda rows, base: (export_txt_zip(rows or [], base), gr.update(visible=True)),
-        inputs=[rows_state, dataset_name], outputs=[txt_zip, txt_zip]
+        lambda tbl, rows, px: (export_excel_with_thumbs(tbl, rows or [], int(px)), gr.update(visible=True)),
+        inputs=[table, rows_state, excel_thumb_px], outputs=[xlsx_file, xlsx_file]
     )
 
 # Launch