Update app.py

app.py

@@ -1,9 +1,24 @@
+# =====================================================================
+# ForgeCaptions - Gradio app for single & batch image captioning
+# =====================================================================
+# CHANGELOG (this version)
+# - GPU-safe: all CUDA only inside @spaces.GPU functions.
+# - Restored: Single tab + Batch chunking (Auto / All-at-once / Manual step).
+# - Shape Aliases: supports comma/pipe-separated synonyms per row.
+# - Default caption style: "Character training (long)".
+# - Model Instructions + Caption Style in minimizable accordions.
+# - Excel export: thumbnail size slider controls image scaling & row height.
+# - Header logo scaled to the full text stack (centered).
+# - Kept gallery & table positions unchanged; scroll sync retained.
+# =====================================================================
+
+# ------------------------------
+# 0) Imports & environment
+# ------------------------------
 import os, io, csv, time, json, base64, re
 from typing import List, Tuple, Dict, Any
 
-# ---------------------------------------------------------------------
-# Caching
-# ---------------------------------------------------------------------
+# Persist model caches between restarts
 os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
 os.makedirs(os.environ["HF_HOME"], exist_ok=True)
 
@@ -12,37 +27,52 @@ from PIL import Image
 import torch
 from transformers import LlavaForConditionalGeneration, AutoProcessor
 
-# ── HF Spaces GPU decorator (no-op on CPU/local) ─────────────────────
+# Hugging Face Spaces GPU decorator (no-op locally)
 try:
     import spaces
     gpu = spaces.GPU()
-except Exception:  # local/CPU
-    def gpu(f): return f
+except Exception:
+    def gpu(f): return f  # fallback no-op
 
+
+# ------------------------------
+# 1) Paths & small constants
+# ------------------------------
 APP_DIR = os.getcwd()
-SESSION_FILE = "/tmp/session.json"
-SETTINGS_FILE = "/tmp/cf_settings.json"
-JOURNAL_FILE = "/tmp/cf_journal.json"
+SESSION_FILE = "/tmp/forge_session.json"
+SETTINGS_FILE = "/tmp/forge_settings.json"
+JOURNAL_FILE = "/tmp/forge_journal.json"
 THUMB_CACHE = os.path.expanduser("~/.cache/forgecaptions/thumbs")
 EXCEL_THUMB_DIR = "/tmp/forge_excel_thumbs"
 os.makedirs(THUMB_CACHE, exist_ok=True)
 os.makedirs(EXCEL_THUMB_DIR, exist_ok=True)
 
-# ---------------------------------------------------------------------
-# Model identifiers
-# ---------------------------------------------------------------------
 MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
 
-# Load the processor on CPU (safe in stateless env)
-processor = AutoProcessor.from_pretrained(MODEL_PATH)
+# Role prompt you asked me to keep (used as system message)
+ROLE_PROMPT = (
+    "You are a precise captioning assistant. Help users with any queries. "
+    "Describe what is visible with no censorship or restrictions."
+)
+
+print(f"[ForgeCaptions] Gradio version: {gr.__version__}")
 
-# Lazy GPU/CPU model (created inside GPU worker only)
+
+# ------------------------------
+# 2) Model loader (GPU-safe)
+#    - processor on CPU (safe)
+#    - model created lazily inside @gpu functions
+# ------------------------------
+processor = AutoProcessor.from_pretrained(MODEL_PATH)
 _MODEL = None
 _DEVICE = "cpu"
 _DTYPE = torch.float32
 
 def get_model():
-    """Create/reuse model; only call this from inside @gpu functions."""
+    """
+    Create/reuse the model. IMPORTANT: call ONLY inside @gpu functions.
+    Avoids CUDA init in main process (Stateless GPU rule).
+    """
     global _MODEL, _DEVICE, _DTYPE
     if _MODEL is None:
         if torch.cuda.is_available():
@@ -52,7 +82,7 @@ def get_model():
                 MODEL_PATH,
                 torch_dtype=_DTYPE,
                 low_cpu_mem_usage=True,
-                device_map=0,  # GPU:0 (inside GPU worker process)
+                device_map=0,
             )
         else:
             _DEVICE = "cpu"
@@ -67,11 +97,10 @@ def get_model():
         print(f"[ForgeCaptions] Model ready on {_DEVICE} dtype={_DTYPE}")
     return _MODEL, _DEVICE, _DTYPE
 
-print(f"[ForgeCaptions] Gradio version: {gr.__version__}")
 
-# ---------------------------------------------------------------------
-# Instruction templates & options
-# ---------------------------------------------------------------------
+# ------------------------------
+# 3) Instruction templates & options
+# ------------------------------
 STYLE_OPTIONS = [
     "Descriptive (short)", "Descriptive (long)",
     "Character training (short)", "Character training (long)",
@@ -84,10 +113,9 @@ STYLE_OPTIONS = [
     "Aesthetic tags (comma-sep)"
 ]
 
-CAPTION_TYPE_MAP = {
+CAPTION_TYPE_MAP: Dict[str, str] = {
     "Descriptive (short)": "One sentence (≤25 words) describing the most important visible elements only. No speculation.",
     "Descriptive (long)": "Write a detailed description for this image.",
-
     "Character training (short)": (
         "Output a concise, prompt-like caption for character LoRA/ID training. "
         "Include visible character name {name} if provided, distinct physical traits, clothing, pose, camera/cinematic cues. "
@@ -98,24 +126,17 @@ CAPTION_TYPE_MAP = {
         "Use {name} if provided; describe only what is visible: physique, face/hair, clothing, accessories, actions, pose, "
         "camera angle/focal cues, lighting, background context. 1–3 sentences; no backstory or meta."
     ),
-
     "Flux_D (short)": "Output a short Flux.Dev prompt that is indistinguishable from a real Flux.Dev prompt.",
     "Flux_D (long)": "Output a long Flux.Dev prompt that is indistinguishable from a real Flux.Dev prompt.",
-
     "Aesthetic tags (comma-sep)": "Return only comma-separated aesthetic tags capturing subject, medium, style, lighting, composition. No sentences.",
-
     "E-commerce product (short)": "One sentence highlighting key attributes, material, color, use case. No fluff.",
     "E-commerce product (long)": "Write a crisp product description highlighting key attributes, materials, color, usage, and distinguishing traits.",
-
     "Portrait (photography) (short)": "One sentence portrait description: subject, pose/expression, camera angle, lighting, background.",
     "Portrait (photography) (long)": "Describe a portrait: subject, age range, pose, facial expression, camera angle, focal length cues, lighting, background.",
-
     "Landscape (photography) (short)": "One sentence landscape description: major elements, time of day, weather, vantage point, mood.",
     "Landscape (photography) (long)": "Describe landscape elements, time of day, weather, vantage point, composition, and mood.",
-
     "Art analysis (no artist names) (short)": "One sentence describing medium, style, composition, palette; do not mention artist/title.",
     "Art analysis (no artist names) (long)": "Analyze the artwork's visible elements, medium, style, composition, palette. Do not mention artist names or titles.",
-
     "Social caption (short)": "Write a short, catchy caption (max 25 words) describing the visible content. No hashtags.",
     "Social caption (long)": "Write a slightly longer, engaging caption (≤50 words) describing the visible content. No hashtags."
 }
@@ -142,73 +163,10 @@ EXTRA_CHOICES = [
 ]
 NAME_OPTION = "If there is a person/character in the image you must refer to them as {name}."
 
-# ---------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------
-def ensure_thumb(path: str, max_side=256) -> str:
-    try:
-        im = Image.open(path).convert("RGB")
-    except Exception:
-        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_path, "JPEG", quality=85, optimize=True)
-        return out_path
-    except Exception:
-        return path
-
-def resize_for_model(im: Image.Image, max_side: int) -> Image.Image:
-    w, h = im.size
-    if max(w, h) <= max_side:
-        return im
-    s = max_side / max(w, h)
-    return im.resize((int(w*s), int(h*s)), Image.LANCZOS)
-
-def apply_prefix_suffix(caption: str, trigger_word: str, begin_text: str, end_text: str) -> str:
-    parts = []
-    if trigger_word.strip():
-        parts.append(trigger_word.strip())
-    if begin_text.strip():
-        parts.append(begin_text.strip())
-    parts.append(caption.strip())
-    if end_text.strip():
-        parts.append(end_text.strip())
-    return " ".join([p for p in parts if p])
-
-# Instruction + caption
-def final_instruction(style_list: List[str], extra_opts: List[str], name_value: str) -> str:
-    styles = style_list or ["Descriptive (short)"]
-    parts = [CAPTION_TYPE_MAP.get(s, "") for s in styles]
-    core = " ".join(p for p in parts if p).strip()
-    if extra_opts:
-        core += " " + " ".join(extra_opts)
-    if NAME_OPTION in (extra_opts or []):
-        core = core.replace("{name}", (name_value or "{NAME}").strip())
-    return core
-
-def logo_b64_img() -> str:
-    candidates = [
-        os.path.join(APP_DIR, "forgecaptions-logo.png"),
-        os.path.join(APP_DIR, "captionforge-logo.png"),
-        "/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 ""
 
-# ---------------------------------------------------------------------
-# Persistence
-# ---------------------------------------------------------------------
+# ------------------------------
+# 4) Persistence helpers (settings/session/journal)
+# ------------------------------
 def save_session(rows: List[dict]):
     with open(SESSION_FILE, "w", encoding="utf-8") as f:
         json.dump(rows, f, ensure_ascii=False, indent=2)
@@ -229,13 +187,14 @@ def load_settings() -> dict:
             cfg = json.load(f)
     else:
         cfg = {}
+    # sensible defaults for this app/version
     defaults = {
         "dataset_name": "forgecaptions",
         "temperature": 0.6,
         "top_p": 0.9,
         "max_tokens": 256,
         "max_side": 896,
-        "styles": ["Character training (long)"],  # ← default changed
+        "styles": ["Character training (long)"],  # default you requested
         "extras": [],
         "name": "",
         "trigger": "",
@@ -247,23 +206,9 @@ def load_settings() -> dict:
     }
     for k, v in defaults.items():
         cfg.setdefault(k, v)
-
-    legacy_map = {
-        "Descriptive": "Descriptive (short)",
-        "LoRA (Flux_D Realism)": "LoRA (Flux_D Realism) (short)",
-        "Portrait (photography)": "Portrait (photography) (short)",
-        "Landscape (photography)": "Landscape (photography) (short)",
-        "Art analysis (no artist names)": "Art analysis (no artist names) (short)",
-        "E-commerce product": "E-commerce product (short)",
-    }
+    # validate styles against allowed set
     styles = cfg.get("styles") or []
-    migrated = []
-    for s in styles if isinstance(styles, list) else [styles]:
-        migrated.append(legacy_map.get(s, s))
-    migrated = [s for s in migrated if s in STYLE_OPTIONS]
-    if not migrated:
-        migrated = ["Descriptive (short)"]
-    cfg["styles"] = migrated
+    cfg["styles"] = [s for s in (styles if isinstance(styles, list) else [styles]) if s in STYLE_OPTIONS] or ["Character training (long)"]
     return cfg
 
 def save_journal(data: dict):
@@ -276,10 +221,72 @@ def load_journal() -> dict:
             return json.load(f)
     return {}
 
-# ---------------------------------------------------------------------
-# Shape Aliases
-# ---------------------------------------------------------------------
+
+# ------------------------------
+# 5) Small utilities (thumbs, resize, prefix/suffix)
+# ------------------------------
+def ensure_thumb(path: str, max_side=256) -> str:
+    try:
+        im = Image.open(path).convert("RGB")
+    except Exception:
+        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_path, "JPEG", quality=85, optimize=True)
+        return out_path
+    except Exception:
+        return path
+
+def resize_for_model(im: Image.Image, max_side: int) -> Image.Image:
+    w, h = im.size
+    if max(w, h) <= max_side:
+        return im
+    s = max_side / max(w, h)
+    return im.resize((int(w*s), int(h*s)), Image.LANCZOS)
+
+def apply_prefix_suffix(caption: str, trigger_word: str, begin_text: str, end_text: str) -> str:
+    parts = []
+    if trigger_word.strip():
+        parts.append(trigger_word.strip())
+    if begin_text.strip():
+        parts.append(begin_text.strip())
+    parts.append(caption.strip())
+    if end_text.strip():
+        parts.append(end_text.strip())
+    return " ".join([p for p in parts if p])
+
+def logo_b64_img() -> str:
+    """
+    Load a PNG logo if present (falls back gracefully).
+    """
+    candidates = [
+        os.path.join(APP_DIR, "forgecaptions-logo.png"),
+        os.path.join(APP_DIR, "captionforge-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 ""
+
+
+# ------------------------------
+# 6) Shape Aliases (comma/pipe synonyms per row)
+# ------------------------------
 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, whole-word, and allow '-shaped' or ' shaped'.
+    """
     s = load_settings()
     if not s.get("shape_aliases_enabled", True):
         return []
@@ -289,18 +296,14 @@ def _compile_shape_aliases_from_file():
         name = (item.get("name")  or "").strip()
         if not raw or not name:
             continue
-        # allow comma or pipe separated synonyms in one cell
         tokens = [t.strip() for t in re.split(r"[|,]", raw) if t.strip()]
         if not tokens:
             continue
-        # de-dup and prefer longer phrases first (prevents "diamond" eating "diamond emblem")
-        tokens = sorted(set(tokens), key=lambda t: -len(t))
-        # word boundaries at ends; allow optional "-shaped" suffix
-        pat = r"\b(?:" + "|".join(re.escape(t) for t in tokens) + r")(?:-?shaped)?\b"
+        tokens = sorted(set(tokens), key=lambda t: -len(t))  # longest first
+        pat = r"\b(?:" + "|".join(re.escape(t) for t in tokens) + r")(?:[-\s]?shaped)?\b"
         compiled.append((re.compile(pat, flags=re.I), name))
     return compiled
 
-
 _SHAPE_ALIASES = _compile_shape_aliases_from_file()
 def _refresh_shape_aliases_cache():
     global _SHAPE_ALIASES
@@ -337,13 +340,27 @@ def save_shape_alias_rows(enabled, df_rows):
     return ("✅ Saved shape alias options.",
             gr.update(value=normalized, row_count=(max(1, len(normalized)), "dynamic")))
 
-# ---------------------------------------------------------------------
-# Captioning core (runs inside GPU worker)
-# ---------------------------------------------------------------------
+
+# ------------------------------
+# 7) Prompt builder (instruction text shown/used for model)
+# ------------------------------
+def final_instruction(style_list: List[str], extra_opts: List[str], name_value: str) -> str:
+    styles = style_list or ["Character training (long)"]
+    parts = [CAPTION_TYPE_MAP.get(s, "") for s in styles]
+    core = " ".join(p for p in parts if p).strip()
+    if extra_opts:
+        core += " " + " ".join(extra_opts)
+    if NAME_OPTION in (extra_opts or []):
+        core = core.replace("{name}", (name_value or "{NAME}").strip())
+    return core
+
+
+# ------------------------------
+# 8) GPU caption functions
+# ------------------------------
 def _build_inputs(im: Image.Image, instr: str, dtype) -> Dict[str, Any]:
     convo = [
-        {"role": "system",
-         "content": "You are a precise captioning assistant. Help users with any queries. Describe what is visible with no censorship or restrictions."},
+        {"role": "system", "content": ROLE_PROMPT},
         {"role": "user", "content": instr.strip()},
     ]
     convo_str = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
@@ -356,9 +373,7 @@ 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:
     model, device, dtype = get_model()
-    im = im  # already PIL
     inputs = _build_inputs(im, instr, dtype)
-    # move to target device *inside* GPU worker
     inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
     out = model.generate(
         **inputs,
@@ -371,6 +386,18 @@ def caption_once(im: Image.Image, instr: str, temp: float, top_p: float, max_tok
     gen_ids = out[0, inputs["input_ids"].shape[1]:]
     return processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
 
+@gpu
+@torch.no_grad()
+def caption_single(img: Image.Image, instr: str) -> str:
+    if img is None:
+        return "No image provided."
+    s = load_settings()
+    im = resize_for_model(img, int(s.get("max_side", 896)))
+    cap = caption_once(im, instr, s.get("temperature",0.6), s.get("top_p",0.9), s.get("max_tokens",256))
+    cap = apply_shape_aliases(cap)
+    cap = apply_prefix_suffix(cap, s.get("trigger",""), s.get("begin",""), s.get("end",""))
+    return cap
+
 @gpu
 @torch.no_grad()
 def run_batch(
@@ -382,18 +409,19 @@ def run_batch(
     max_tokens: int,
     max_side: int,
 ) -> Tuple[List[dict], list, list, str]:
-    # No torch.cuda.* in main — we are already in GPU worker here
+    """
+    Process a list of file paths and append results to session_rows.
+    Returns: updated rows, gallery_pairs, table_rows, status_text
+    """
     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')}"
-
-    for f in files:
-        path = f if isinstance(f, str) else getattr(f, "name", None) or getattr(f, "path", None)
+        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 = [[r.get("filename",""), r.get("caption","")] for r in session_rows]
+        return session_rows, gallery_pairs, table_rows, f"Saved • {time.strftime('%H:%M:%S')}"
+
+    for path in files:
         if not path or not os.path.exists(path):
             continue
         try:
@@ -410,25 +438,12 @@ def run_batch(
         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')}"
-
-@gpu
-@torch.no_grad()
-def caption_single(img: Image.Image, instr: str) -> str:
-    if img is None:
-        return "No image provided."
-    s = load_settings()
-    im = resize_for_model(img, int(s.get("max_side", 896)))
-    cap = caption_once(im, instr, s.get("temperature",0.6), s.get("top_p",0.9), s.get("max_tokens",256))
-    cap = apply_shape_aliases(cap)
-    cap = apply_prefix_suffix(cap, s.get("trigger",""), s.get("begin",""), s.get("end",""))
-    return cap
+    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 = [[r.get("filename",""), r.get("caption","")] for r in session_rows]
+    return session_rows, gallery_pairs, table_rows, f"Saved • {time.strftime('%H:%M:%S')}"
 
-# tiny warmup so Spaces sees a GPU function at startup
+# Ensure Spaces detects at least one GPU function at startup
 @gpu
 @torch.no_grad()
 def _gpu_startup_warm():
@@ -439,9 +454,10 @@ def _gpu_startup_warm():
     except Exception as e:
         print("[ForgeCaptions] GPU warmup skipped:", e)
 
-# ---------------------------------------------------------------------
-# Export helpers
-# ---------------------------------------------------------------------
+
+# ------------------------------
+# 9) Export helpers (CSV/XLSX)
+# ------------------------------
 def _rows_to_table(rows: List[dict]) -> list:
     return [[r.get("filename",""), r.get("caption","")] for r in (rows or [])]
 
@@ -460,9 +476,7 @@ 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(["filename", "caption"])
-        w.writerows(data)
+        w = csv.writer(f); w.writerow(["filename", "caption"]); w.writerows(data)
     return out
 
 def _resize_for_excel(path: str, px: int) -> str:
@@ -504,11 +518,11 @@ def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_p
     ws.column_dimensions["B"].width = 42
     ws.column_dimensions["C"].width = 100
 
-    row_h = int(int(thumb_px) * 0.75)  # px→pt-ish
+    # Convert pixel target to approx. row points (Excel ≈ 0.75 * px)
+    row_h = int(int(thumb_px) * 0.75)
     r_i = 2
     for r in (session_rows or []):
-        fn = r.get("filename","")
-        cap = caption_by_file.get(fn, r.get("caption",""))
+        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")
@@ -526,53 +540,63 @@ def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_p
     wb.save(out)
     return out
 
-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
-# ---------------------------------------------------------------------
+# ------------------------------
+# 10) UI (Blocks)
+# ------------------------------
 BASE_CSS = """
 :root{--galleryW:50%;--tableW:50%;}
 .gradio-container{max-width:100%!important}
-.cf-hero{display:flex; align-items:center; justify-content:center; gap:16px;
-  margin:4px 0 12px; text-align:center;}
-.cf-hero > div { text-align:center; }
-.cf-logo{height:calc(3.25rem + 3 * 1.1rem + 18px);width:auto;object-fit:contain}
+.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-logo{
+  /* Make logo fill roughly the full text stack; clamped for sanity */
+  height: clamp(120px, calc(3.25rem + 3 * 1.1rem + 24px), 180px);
+  width:auto; object-fit:contain; display:block; flex:0 0 auto;
+}
 .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}
+
+/* Results area */
 .cf-scroll{max-height:70vh; overflow-y:auto; border:1px solid #e6e6e6; border-radius:10px; padding:8px}
 #cfGal .grid > div { height: 96px; }
 """
 
 with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
-    # ensure Spaces sees a GPU function at start (without touching CUDA in main)
+    # Ensure Spaces sees a GPU function (without touching CUDA in main)
     demo.load(_gpu_startup_warm, inputs=None, outputs=None)
 
-    settings = load_settings()
-    settings["styles"] = [s for s in settings.get("styles", []) if s in STYLE_OPTIONS] or ["Character training (long)"]
-
+    # ---- Header (logo + title center). Script sets logo height to match text exactly.
     gr.HTML(value=f"""
 <div class="cf-hero">
   {logo_b64_img()}
-  <div>
+  <div class="cf-text">
     <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 export</div>
   </div>
 </div>
-<hr>""")
+<hr>
+<script>
+setTimeout(() => {{
+  const logo = document.querySelector(".cf-logo");
+  const text = document.querySelector(".cf-text");
+  if (logo && text) logo.style.height = text.getBoundingClientRect().height + "px";
+}}, 0);
+</script>
+""")
+
+    # ---- Settings state (loaded once)
+    settings = load_settings()
 
-    # ── Controls
+    # ---- Controls group (left/right columns)
     with gr.Group():
         with gr.Row():
+            # LEFT: Style + Extra + Name/Prefix/Suffix (accordions minimizable)
             with gr.Column(scale=2):
                 with gr.Accordion("Caption style (choose one or combine)", open=True):
                     style_checks = gr.CheckboxGroup(
@@ -592,6 +616,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                     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: Instruction preview + dataset + sliders
             with gr.Column(scale=1):
                 with gr.Accordion("Model Instructions", open=False):
                     instruction_preview = gr.Textbox(label=None, lines=12)
@@ -600,8 +625,15 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 max_side   = gr.Slider(256, 1024, settings.get("max_side", 896), 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 settings: temperature 0.6 • top-p 0.9 • max tokens 256")
-
+                # Chunking controls (restored)
+                chunk_mode = gr.Radio(
+                    choices=["Auto", "Manual (all at once)", "Manual (step)"],
+                    value="Manual (step)",
+                    label="Batch mode"
+                )
+                chunk_size = gr.Slider(1, 50, value=10, step=1, label="Chunk size")
+
+    # -- Keep instruction text in sync with controls and persist to settings
     def _refresh_instruction(styles, extra, name_value, trigv, begv, endv, excel_px, ms):
         instr = final_instruction(styles or ["Character training (long)"], extra or [], name_value)
         cfg = load_settings()
@@ -624,43 +656,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
     demo.load(lambda s,e,n: final_instruction(s or ["Character training (long)"], e or [], n),
               inputs=[style_checks, extra_opts, name_input], outputs=[instruction_preview])
 
-    # ── Shape Aliases (improved)
-    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"
-            "- Left column = a single token **or** comma/pipe-separated synonyms, e.g. `penis, cock | phallic`\n"
-            "- Right column = replacement name, e.g. `family-emblem`\n\n"
-            "Matches are case-insensitive, use whole words, and also catch `*-shaped` (e.g., `diamond-shaped`).\n"
-            "Multi-word phrases are supported."
-        )
-        init_rows, init_enabled = get_shape_alias_rows_ui_defaults()
-        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,
-            col_count=(2, "fixed"),
-            row_count=(max(1, len(init_rows)), "dynamic"),
-            datatype=["str","str"],
-            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("")
-        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])
-
-    # ── Tabs: Single & Batch
+    # ---- Tabs: Single & Batch
     with gr.Tabs():
         with gr.Tab("Single"):
             input_image_single = gr.Image(type="pil", label="Input Image", height=512, width=512)
@@ -677,9 +673,10 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 input_files = gr.File(label="Drop images", file_types=["image"], file_count="multiple", type="filepath")
             run_button = gr.Button("Caption batch", variant="primary")
 
-    # ── Results + Table (same position)
+    # ---- Results (UNCHANGED POSITION): Gallery left, Table right
     rows_state  = gr.State(load_session())
     autosave_md = gr.Markdown("Ready.")
+    remaining_state = gr.State([])  # for manual step mode
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -702,7 +699,14 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 elem_classes=["cf-scroll"]
             )
 
-    # Exports
+    # ---- Step panel (restored)
+    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")
+
+    # ---- Exports
     with gr.Row():
         with gr.Column():
             export_csv_btn  = gr.Button("Export CSV")
@@ -711,13 +715,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
             export_xlsx_btn = gr.Button("Export Excel (.xlsx) with thumbnails")
             xlsx_file       = gr.File(label="Excel file", visible=False)
 
-    def _initial_gallery(rows):
-        rows = rows or []
-        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
+    # ---- Scroll sync (gallery ↔ table)
     gr.HTML("""
 <script>
 (function () {
@@ -762,29 +760,116 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
 </script>
 """)
 
-    # Batch run → rows + gallery + table
-    def _run_click(files, rows, instr, ms):
+    # ---- Batch chunking logic (restored)
+    def _split_chunks(files, csize: int):
+        files = files or []
+        c = max(1, int(csize))
+        return [files[i:i+c] for i in range(0, len(files), c)]
+
+    def _tpms():
         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
+        return s.get("temperature", 0.6), s.get("top_p", 0.9), s.get("max_tokens", 256)
+
+    def _run_click(files, rows, instr, ms, mode, csize):
+        t, p, m = _tpms()
+        files = files or []
+        # Manual step → process first chunk only
+        if mode == "Manual (step)" and files:
+            chunks = _split_chunks(files, int(csize))
+            batch = chunks[0]
+            remaining = sum(chunks[1:], [])
+            new_rows, gal, tbl, stamp = run_batch(batch, rows or [], instr, t, p, m, int(ms))
+            panel_vis = gr.update(visible=bool(remaining))
+            msg = f"{len(remaining)} files remain. Process next chunk?"
+            return new_rows, gal, tbl, stamp, remaining, panel_vis, gr.update(value=msg)
+        # Auto / all-at-once → process everything in one go
+        else:
+            new_rows, gal, tbl, stamp = run_batch(files, rows or [], instr, t, p, m, int(ms))
+            return new_rows, gal, tbl, stamp, [], gr.update(visible=False), gr.update(value="")
 
     run_button.click(
         _run_click,
-        inputs=[input_files, rows_state, instruction_preview, max_side],
-        outputs=[rows_state, gallery, table, autosave_md]
+        inputs=[input_files, rows_state, instruction_preview, max_side, chunk_mode, chunk_size],
+        outputs=[rows_state, gallery, table, autosave_md, remaining_state, step_panel, step_msg]
+    )
+
+    def _step_next(remain, rows, instr, ms, csize):
+        t, p, m = _tpms()
+        remain = remain or []
+        if not remain:
+            return rows, gr.update(value="No files remaining."), gr.update(visible=False), [], [], [], "Saved."
+        batch = remain[:int(csize)]
+        leftover = remain[int(csize):]
+        new_rows, gal, tbl, stamp = run_batch(batch, rows or [], instr, t, p, m, int(ms))
+        panel_vis = gr.update(visible=bool(leftover))
+        msg = f"{len(leftover)} files remain. Process next chunk?" if leftover else "All done."
+        return new_rows, msg, panel_vis, leftover, gal, tbl, stamp
+
+    step_next.click(
+        _step_next,
+        inputs=[remaining_state, rows_state, instruction_preview, max_side, chunk_size],
+        outputs=[rows_state, step_msg, step_panel, remaining_state, gallery, table, autosave_md]
     )
 
-    # Table edits sync
+    def _step_finish():
+        return gr.update(visible=False), gr.update(value=""), []
+
+    step_finish.click(
+        _step_finish,
+        inputs=None,
+        outputs=[step_panel, step_msg, remaining_state]
+    )
+
+    # ---- Table edits → persist + refresh gallery
+    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')}"
+
     table.change(
         sync_table_to_session,
         inputs=[table, rows_state],
         outputs=[rows_state, gallery, autosave_md]
     )
 
-    # Exports
+    # ---- Shape Aliases accordion (with examples & buttons)
+    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"
+            "- Left column = a single token **or** comma/pipe-separated synonyms, e.g. `diamond, rhombus | lozenge`\n"
+            "- Right column = replacement name, e.g. `starkey-emblem`\n"
+            "Matches are case-insensitive, whole-word, and also catch `*-shaped` or `* shaped`."
+        )
+        init_rows, init_enabled = get_shape_alias_rows_ui_defaults()
+        enable_aliases = gr.Checkbox(label="Enable shape alias replacements", value=init_enabled)
+        alias_table = gr.Dataframe(
+            headers=["shape (token or synonyms)", "name to insert"],
+            value=init_rows,
+            col_count=(2, "fixed"),
+            row_count=(max(1, len(init_rows)), "dynamic"),
+            datatype=["str","str"],
+            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("")
+        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])
+
+    # ---- Exports
     export_csv_btn.click(
         lambda tbl: (export_csv_from_table(tbl), gr.update(visible=True)),
         inputs=[table], outputs=[csv_file, csv_file]
@@ -794,7 +879,10 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
         inputs=[table, rows_state, excel_thumb_px], outputs=[xlsx_file, xlsx_file]
     )
 
-# Launch (SSR off for stability on Spaces)
+
+# ------------------------------
+# 11) Launch (SSR disabled for stability on Spaces)
+# ------------------------------
 if __name__ == "__main__":
     demo.queue(max_size=64).launch(
         server_name="0.0.0.0",