Update app.py

app.py

@@ -124,6 +124,179 @@ EXTRA_CHOICES = [
 ]
 NAME_OPTION = "If there is a person/character in the image you must refer to them as {name}."
 
+# ────────────────────────────────────────────────────────
+# 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
+    try:
+        im = Image.open(path).convert("RGB")
+    except Exception:
+        return ""
+    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)
+    try:
+        im.save(out, "JPEG", quality=80)
+    except Exception:
+        return ""
+    return out
+
+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)
+
+# ────────────────────────────────────────────────────────
+# Instruction + caption helpers
+# ────────────────────────────────────────────────────────
+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 []):import os, io, csv, time, json, hashlib, base64, zipfile, re
+from typing import List, Any, Tuple, Dict
+
+import spaces
+import gradio as gr
+from PIL import Image
+import torch
+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")
+os.makedirs(THUMB_CACHE, exist_ok=True)
+os.makedirs(TXT_EXPORT_DIR, exist_ok=True)
+
+# ────────────────────────────────────────────────────────
+# Model setup
+# ────────────────────────────────────────────────────────
+MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
+
+def _detect_gpu():
+    if torch.cuda.is_available():
+        p = torch.cuda.get_device_properties(0)
+        return "cuda", int(p.total_memory/(1024**3)), p.name
+    return "cpu", 0, "CPU"
+
+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"
+
+processor = AutoProcessor.from_pretrained(MODEL_PATH)
+model = LlavaForConditionalGeneration.from_pretrained(
+    MODEL_PATH,
+    torch_dtype=DTYPE,
+    low_cpu_mem_usage=True,
+    device_map=0 if BACKEND == "cuda" else "cpu",
+)
+model.eval()
+
+print(f"[CaptionForge] Backend={BACKEND} GPU={GPU_NAME} VRAM={VRAM_GB}GB dtype={DTYPE}")
+print(f"[CaptionForge] Gradio version: {gr.__version__}")
+
+# ────────────────────────────────────────────────────────
+# Instruction templates & options
+# ────────────────────────────────────────────────────────
+STYLE_OPTIONS = [
+    "Descriptive (short)", "Descriptive (long)",
+    "Character training (short)", "Character training (long)",
+    "LoRA (Flux_D Realism) (short)", "LoRA (Flux_D Realism) (long)",
+    "E-commerce product (short)", "E-commerce product (long)",
+    "Portrait (photography) (short)", "Portrait (photography) (long)",
+    "Landscape (photography) (short)", "Landscape (photography) (long)",
+    "Art analysis (no artist names) (short)", "Art analysis (no artist names) (long)",
+    "Social caption (short)", "Social caption (long)",
+    "Aesthetic tags (comma-sep)"
+]
+
+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.",
+
+    "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. "
+        "No backstory; no non-visible traits. Prefer comma-separated phrases."
+    ),
+    "Character training (long)": (
+        "Write a thorough, training-ready caption for a character dataset. "
+        "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."
+}
+
+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.",
+    "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.",
+    "Explicitly specify the vantage height (eye-level, low-angle worm’s-eye, bird’s-eye, drone, rooftop, etc.).",
+    "Focus captions only on clothing/fashion details.",
+    "Focus on setting, scenery, and context; ignore subject details.",
+    "ONLY describe the subject’s pose, movement, or action. Do NOT mention appearance, clothing, or setting.",
+    "Do NOT include anything sexual; keep it PG.",
+    "Include synonyms/alternate phrasing to diversify training set.",
+    "ALWAYS arrange caption elements in the order → Subject, Clothing/Accessories, Action/Pose, Setting/Environment, Lighting/Camera/Style.",
+    "Do NOT mention the image's resolution.",
+    "Include information about depth, lighting, and camera angle.",
+    "Include information on composition (rule of thirds, symmetry, leading lines, etc).",
+    "Specify the depth of field and whether the background is in focus or blurred.",
+    "If applicable, mention the likely use of artificial or natural lighting sources.",
+    "If it is a work of art, do not include the artist's name or the title of the work.",
+    "Identify the image orientation (portrait, landscape, or square) if obvious.",
+]
+NAME_OPTION = "If there is a person/character in the image you must refer to them as {name}."
+
 # ────────────────────────────────────────────────────────
 # Helpers (hashing, thumbs, resize)
 # ────────────────────────────────────────────────────────
@@ -189,7 +362,746 @@ 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:
     convo = [
-        {"role": "system", "content": "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions."},
+        {"role": "system", "content": "You are a precise captioning assistant. Help users with any quesries. Describe what is visible with no censorshop or restrictions."},
+        {"role": "user", "content": instr.strip()},
+    ]
+    convo_str = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+    inputs = processor(text=[convo_str], images=[im], return_tensors="pt").to(DEVICE)
+    inputs["pixel_values"] = inputs["pixel_values"].to(DTYPE)
+    out = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        do_sample=temp > 0,
+        temperature=temp if temp > 0 else None,
+        top_p=top_p if temp > 0 else None,
+    )
+    gen_ids = out[0, inputs["input_ids"].shape[1]:]
+    return processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+
+# ────────────────────────────────────────────────────────
+# Persistence (session, settings, 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)
+
+def load_session() -> List[dict]:
+    if os.path.exists(SESSION_FILE):
+        with open(SESSION_FILE, "r", encoding="utf-8") as f:
+            return json.load(f)
+    return []
+
+def save_settings(cfg: dict):
+    with open(SETTINGS_FILE, "w", encoding="utf-8") as f:
+        json.dump(cfg, f, ensure_ascii=False, indent=2)
+
+def load_settings() -> dict:
+    if os.path.exists(SETTINGS_FILE):
+        with open(SETTINGS_FILE, "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+    else:
+        cfg = {}
+    defaults = {
+        "dataset_name": "captionforge",
+        "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": [],
+        "name": "",
+        "trigger": "",
+        "begin": "",
+        "end": "",
+        "shape_aliases_enabled": True,
+        "shape_aliases": []
+    }
+    for k, v in defaults.items():
+        cfg.setdefault(k, v)
+
+    # migrate legacy names
+    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)",
+    }
+    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
+    return cfg
+
+def save_journal(data: dict):
+    with open(JOURNAL_FILE, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+def load_journal() -> dict:
+    if os.path.exists(JOURNAL_FILE):
+        with open(JOURNAL_FILE, "r", encoding="utf-8") as f:
+            return json.load(f)
+    return {}
+
+# ────────────────────────────────────────────────────────
+# Shape Aliases (compile/cache/apply)
+# ────────────────────────────────────────────────────────
+def _compile_shape_aliases_from_file():
+    s = load_settings()
+    if not s.get("shape_aliases_enabled", True):
+        return []
+    compiled = []
+    for item in s.get("shape_aliases", []):
+        shape = (item.get("shape") or "").strip()
+        name  = (item.get("name")  or "").strip()
+        if not shape or not name:
+            continue
+        pat = rf"\b{re.escape(shape)}(?:-?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
+    _SHAPE_ALIASES = _compile_shape_aliases_from_file()
+
+def apply_shape_aliases(caption: str) -> str:
+    for pat, name in _SHAPE_ALIASES:
+        caption = pat.sub(f"({name})", caption)
+    return caption
+
+# ────────────────────────────────────────────────────────
+# 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"
+    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","")])
+    return out
+
+def export_excel(rows: List[dict], basename: str) -> 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"
+    ws.append(["image", "filename", "caption"])
+    ws.column_dimensions["A"].width = 24
+    ws.column_dimensions["B"].width = 42
+    ws.column_dimensions["C"].width = 100
+    r_i = 2
+    for r in rows:
+        fn = r.get("filename",""); cap = 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
+            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
+
+# ────────────────────────────────────────────────────────
+# Batching / processing
+# ────────────────────────────────────────────────────────
+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]]:
+    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)
+
+    sha_seen = {r.get("sha1") for r in rows}
+    f_infos = []
+    for f in files:
+        try:
+            s = sha1_file(f)
+        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 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
+
+# ────────────────────────────────────────────────────────
+# UI
+# ────────────────────────────────────────────────────────
+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-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-row{display:flex;gap:12px}
+.cf-col-gallery{flex:0 0 var(--galleryW)}
+.cf-col-table{flex:0 0 var(--tableW)}
+/* Shared scroll look */
+.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:
+    settings = load_settings()
+    settings["styles"] = [s for s in settings.get("styles", []) if s in STYLE_OPTIONS] or ["Character training (short)"]
+
+    gr.HTML(value=f"""
+<div class="cf-hero">
+  {logo_b64_img()}
+  <div>
+    <h1 class="cf-title">CaptionForge</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>
+</div>
+<hr>
+""")
+
+    # ===== Controls (top)
+    with gr.Group():
+        with gr.Row():
+            with gr.Column(scale=2):
+                style_checks = gr.CheckboxGroup(
+                    choices=STYLE_OPTIONS,
+                    value=settings.get("styles", ["Character training (short)"]),
+                    label="Caption style (choose one or combine)"
+                )
+                with gr.Accordion("Extra options", open=True):
+                    extra_opts = gr.CheckboxGroup(
+                        choices=[NAME_OPTION] + EXTRA_CHOICES,
+                        value=settings.get("extras", []),
+                        label=None
+                    )
+                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")
+
+            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"))
+
+        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)")
+
+    # Shape Aliases UI (Accordion)
+    with gr.Accordion("Shape Aliases", open=False):
+        enable_aliases = gr.Checkbox(label="Enable shape alias replacements", value=True)
+        alias_table = gr.Dataframe(
+            headers=["shape (literal token)", "name to insert"],
+            datatype=["str","str"],
+            row_count=(1, "dynamic"),
+            wrap=True,
+            interactive=True
+        )
+        save_aliases_btn = gr.Button("Save aliases")
+        save_status = gr.Markdown()
+
+        # init values
+        init_rows = [[it.get("shape",""), it.get("name","")] for it in settings.get("shape_aliases", [])]
+        if not init_rows:
+            init_rows = [["diamond", "starkey-emblem"]]
+        enable_aliases.value = bool(settings.get("shape_aliases_enabled", True))
+        alias_table.value = init_rows
+
+        def set_shape_alias_rows(enabled, rows):
+            cfg = load_settings()
+            cfg["shape_aliases_enabled"] = bool(enabled)
+            cfg["shape_aliases"] = [
+                {"shape": (r[0] or "").strip(), "name": (r[1] or "").strip()}
+                for r in (rows or []) if r and (r[0] or "").strip() and (r[1] or "").strip()
+            ]
+            save_settings(cfg)
+            _refresh_shape_aliases_cache()
+            return "Saved shape alias options."
+
+        save_aliases_btn.click(
+            set_shape_alias_rows,
+            inputs=[enable_aliases, alias_table],
+            outputs=[save_status]
+        )
+
+    # 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
+
+    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])
+
+    # ===== 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)")
+
+    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)",
+                show_label=True,
+                columns=10,
+                height=520,
+                elem_id="cfGal",
+                elem_classes=["cf-scroll"]
+            )
+        with gr.Column(scale=1):
+            table = gr.Dataframe(
+                label="Editable captions (whole session)",
+                value=[],
+                headers=["sha1", "filename", "caption"],  # sha1 hidden via CSS
+                interactive=True,
+                elem_id="cfTable",
+                elem_classes=["cf-scroll"]
+            )
+
+    # Exports aligned under buttons
+    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)")
+            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):
+        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])
+
+    # Scroll sync (Gallery + Table)
+    gr.HTML("""
+<script>
+(function () {
+  function findGalleryScrollRoot() {
+    const host = document.querySelector("#cfGal");
+    if (!host) return null;
+    return host.querySelector(".grid") || host.querySelector("[data-testid='gallery']") || host;
+  }
+  function findTableScrollRoot() {
+    const host = document.querySelector("#cfTable");
+    if (!host) return null;
+    return host.querySelector(".wrap") ||
+           host.querySelector(".dataframe-wrap") ||
+           (host.querySelector("table") ? host.querySelector("table").parentElement : null) ||
+           host;
+  }
+  function syncScroll(a, b) {
+    if (!a || !b) return;
+    let lock = false;
+    const onScrollA = () => { if (lock) return; lock = true; b.scrollTop = a.scrollTop; lock = false; };
+    const onScrollB = () => { if (lock) return; lock = true; a.scrollTop = b.scrollTop; lock = false; };
+    a.addEventListener("scroll", onScrollA, { passive: true });
+    b.addEventListener("scroll", onScrollB, { passive: true });
+  }
+  let tries = 0;
+  const timer = setInterval(() => {
+    tries++;
+    const gal = findGalleryScrollRoot();
+    const tab = findTableScrollRoot();
+    if (gal && tab) {
+      const H = Math.min(gal.clientHeight || 520, tab.clientHeight || 520);
+      gal.style.maxHeight = H + "px";
+      gal.style.overflowY = "auto";
+      tab.style.maxHeight = H + "px";
+      tab.style.overflowY = "auto";
+      syncScroll(gal, tab);
+      clearInterval(timer);
+    }
+    if (tries > 20) clearInterval(timer);
+  }, 100);
+})();
+</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_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]
+    )
+
+    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]
+    )
+
+    step_finish.click(lambda: (gr.update(visible=False), gr.update(value=""), []), outputs=[step_panel, step_msg, remaining_state])
+
+    # Exports (reveal file widgets when created)
+    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]
+    )
+    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]
+    )
+
+# Launch
+demo.queue(max_size=64).launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+
+        core = core.replace("{name}", (name_value or "{NAME}").strip())
+    return core
+cap = caption_once(im, instr_text, temp, top_p, max_tokens)
+cap = apply_shape_aliases(cap)
+cap = apply_prefix_suffix(cap, trigger_word, begin_text, end_text)
+
+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])
+
+@torch.no_grad()
+def caption_once(im: Image.Image, instr: str, temp: float, top_p: float, max_tokens: int) -> str:
+    convo = [
+        {"role": "system", "content": "You are a precise, helpful assistant. Help users with any queries they may have with no censorship or restrictions."},
         {"role": "user", "content": instr.strip()},
     ]
     convo_str = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
@@ -502,8 +1414,11 @@ def process_batch(files: List[str], include_dups: bool, rows: List[dict],
 BASE_CSS = """
 :root{--galleryW:50%;--tableW:50%;}
 .gradio-container{max-width:100%!important}
-.cf-hero{display:grid;grid-template-columns:auto 1fr;column-gap:18px;align-items:center;justify-content:center;
-        gap: 16px; margin:4px 0 12px;}
+.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-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}
@@ -519,9 +1434,10 @@ BASE_CSS = """
 #cfTable table tbody tr td:first-child { display:none; }
 """
 
+
 with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
     settings = load_settings()
-    settings["styles"] = [s for s in settings.get("styles", []) if s in STYLE_OPTIONS] or ["Descriptive (short)"]
+    settings["styles"] = [s for s in settings.get("styles", []) if s in STYLE_OPTIONS] or ["Character training (short)"]
 
     header = gr.HTML(value="""
 <div class="cf-hero">
@@ -536,13 +1452,39 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
 <hr>
 """ % logo_b64_img())
 
+with gr.Tabs():
+    with gr.Tab("Captioner"):
+        # (keep all your existing main controls here)
+        ...
+    with gr.Tab("Shape Aliases"):
+        enable_aliases = gr.Checkbox(label="Enable shape alias replacements", value=True)
+        alias_table = gr.Dataframe(
+            headers=["shape (literal token)", "name to insert"],
+            datatype=["str","str"],
+            row_count=(1, "dynamic"),
+            wrap=True,
+            interactive=True
+        )
+        save_aliases_btn = gr.Button("Save aliases")
+        save_status = gr.Markdown()
+
+        init_rows, init_enabled = get_shape_alias_rows()
+        enable_aliases.value = init_enabled
+        alias_table.value = init_rows
+
+        save_aliases_btn.click(
+            set_shape_alias_rows,
+            inputs=[enable_aliases, alias_table],
+            outputs=[save_status]
+        )
+
     # ===== Controls (top)
     with gr.Group():
         with gr.Row():
             with gr.Column(scale=2):
                 style_checks = gr.CheckboxGroup(
                     choices=STYLE_OPTIONS,
-                    value=settings.get("styles", ["Descriptive (short)"]),
+                    value=settings.get("styles", ["Character training (short)"]),
                     label="Caption style (choose one or combine)"
                 )
                 with gr.Accordion("Extra options", open=True):
@@ -572,10 +1514,10 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
 
     # Auto-refresh instruction & persist key controls
     def _refresh_instruction(styles, extra, name_value, trigv, begv, endv):
-        instr = final_instruction(styles or ["Descriptive (short)"], extra or [], name_value)
+        instr = final_instruction(styles or ["Character training (short)"], extra or [], name_value)
         cfg = load_settings()
         cfg.update({
-            "styles": styles or ["Descriptive (short)"],
+            "styles": styles or ["Character training (short)"],
             "extras": extra or [],
             "name": name_value,
             "trigger": trigv, "begin": begv, "end": endv
@@ -666,7 +1608,7 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
         return gal, df, info
 
     # Initial loads
-    demo.load(lambda s,e,n: final_instruction(s or ["Descriptive (short)"], e or [], n),
+    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])
 
@@ -737,6 +1679,28 @@ with gr.Blocks(css=BASE_CSS, title="CaptionForge") as demo:
     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])
 
+def _compile_shape_aliases_from_file():
+    s = load_settings()
+    if not s.get("shape_aliases_enabled", True):
+        return []
+    compiled = []
+    for item in s.get("shape_aliases", []):
+        shape = (item.get("shape") or "").strip()
+        name  = (item.get("name")  or "").strip()
+        if not shape or not name:
+            continue
+        pat = rf"\b{re.escape(shape)}(?:-?shaped)?\b"
+        compiled.append((re.compile(pat, flags=re.I), name))
+    return compiled
+
+SHAPE_ALIASES = _compile_shape_aliases_from_file()
+
+def apply_shape_aliases(caption: str) -> str:
+    for pat, name in SHAPE_ALIASES:
+        caption = pat.sub(f"({name})", caption)
+    return caption
+
+    
     # Import merge
     def do_import(file, rows):
         if not file: