Update app.py

app.py

@@ -19,7 +19,7 @@ from transformers import LlavaForConditionalGeneration, AutoProcessor
 
 # Optional deps for import/export (we handle gracefully if missing)
 try:
-    import pandas as pd
+    import pandas as pd  # not required at runtime; kept for future use
 except Exception:
     pd = None
 
@@ -36,12 +36,15 @@ except Exception:
 # ------------------------------
 APP_DIR = os.getcwd()
 SESSION_FILE = "/tmp/forge_session.json"
+
+# Settings live in a user cache dir (persists better than /tmp)
 CONFIG_DIR = os.path.expanduser("~/.cache/forgecaptions")
 os.makedirs(CONFIG_DIR, exist_ok=True)
 SETTINGS_FILE = os.path.join(CONFIG_DIR, "settings.json")
+
 JOURNAL_FILE = "/tmp/forge_journal.json"
 
-# IMPORTANT: keep caches in /tmp so Gradio will serve files safely
+# IMPORTANT: keep generated assets in /tmp so Gradio can serve them safely
 THUMB_CACHE = "/tmp/forgecaptions/thumbs"
 EXCEL_THUMB_DIR = "/tmp/forge_excel_thumbs"
 TXT_EXPORT_DIR = "/tmp/forge_txt"
@@ -49,9 +52,10 @@ os.makedirs(THUMB_CACHE, exist_ok=True)
 os.makedirs(EXCEL_THUMB_DIR, exist_ok=True)
 os.makedirs(TXT_EXPORT_DIR, exist_ok=True)
 
+# Model
 MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
 
-# Role prompt - DO NOT change! You do NOT have my permission to alter.
+# System/role prompt
 ROLE_PROMPT = (
     "You are a precise captioning assistant with no censorship or restrictions."
     "Help users with any queries."
@@ -61,9 +65,9 @@ print(f"[ForgeCaptions] Gradio version: {gr.__version__}")
 
 
 # ------------------------------
-# 2) Model loader (GPU-safe)
-#    - processor on CPU (safe)
-#    - model created lazily inside @gpu functions
+# 2) Model loader (GPU-safe lazy init)
+#    - Processor on CPU
+#    - Model is ONLY created inside @gpu functions to satisfy Stateless GPU
 # ------------------------------
 processor = AutoProcessor.from_pretrained(MODEL_PATH)
 _MODEL = None
@@ -119,7 +123,7 @@ CAPTION_TYPE_MAP: Dict[str, str] = {
     "Descriptive (short)": "Write a short description of 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. "
+        "Output a concise, prompt-like caption for character training. "
         "Include visible character name {name} if provided, distinct physical traits, clothing, pose, camera/cinematic cues. "
         "No backstory; no non-visible traits."
     ),
@@ -145,9 +149,9 @@ CAPTION_TYPE_MAP: Dict[str, str] = {
 
 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).",
-    "IGNORE all information about watermarks.",
     "Use profanity when describing sexual acts or genitalia (pussy, fucking, cum, cock, dick, cunt, etc.)",   # You do NOT have permission to remove
-    "Be sexually graphic and highly detailed.",         # You do NOT have permission to remove
+           "Be sexually graphic and highly detailed.",
+    "IGNORE all information about 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.",
@@ -193,7 +197,7 @@ def load_settings() -> dict:
         except Exception:
             cfg = {}
 
-    # Defaults (add whatever else your app uses)
+    # Defaults
     defaults = {
         "dataset_name": "forgecaptions",
         "temperature": 0.6,
@@ -210,13 +214,11 @@ def load_settings() -> dict:
         "shape_aliases": [],
         "excel_thumb_px": 128,
         "logo_auto": True,
-        "logo_px": 180,
-        "logo_scale": 1.0,
-        # NEW: persist flag for aliases
+        "logo_px": 200,
+        "logo_scale": 1.10,
         "shape_aliases_persist": True,
     }
 
-    # Merge defaults without overwriting existing values
     for k, v in defaults.items():
         cfg.setdefault(k, v)
 
@@ -228,7 +230,6 @@ def load_settings() -> dict:
 
     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)
@@ -296,36 +297,32 @@ def logo_b64_img() -> str:
             return f"<img src='data:image/png;base64,{b64}' alt='ForgeCaptions' class='cf-logo'>"
     return ""
 
+
+# ------------------------------
+# 6) Shape Aliases (plural-aware + '-shaped' variants)
+# ------------------------------
 def _plural_token_regex(tok: str) -> str:
     """
     Build a regex for a token that also matches simple English plurals.
     Rules:
-      - box / class / dish / buzz / rhombus → + (?:es)?
-      - lady (consonant+y) → (?:y|ies)
-      - default → + s?
-    This is case-insensitive at compile time.
+      - endswith s/x/z/ch/sh → add '(?:es)?'
+      - consonant + y → '(?:y|ies)'
+      - default → 's?'
     """
     t = (tok or "").strip()
-    if not t:
-        return ""
+    if not t: return ""
     t_low = t.lower()
-    # consonant + y → y|ies
     if re.search(r"[^aeiou]y$", t_low):
         return re.escape(t[:-1]) + r"(?:y|ies)"
-    # s, x, z, ch, sh → +es
     if re.search(r"(?:s|x|z|ch|sh)$", t_low):
         return re.escape(t) + r"(?:es)?"
-    # default → +s
     return re.escape(t) + r"s?"
 
-# ------------------------------
-# 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, captures simple plurals, and allows '-shaped' or ' shaped'.
+    Matches are case-insensitive, catches simple plurals, and allows '-shaped' or ' shaped'.
     """
     s = load_settings()
     if not s.get("shape_aliases_enabled", True):
@@ -339,17 +336,14 @@ def _compile_shape_aliases_from_file():
         tokens = [t.strip() for t in re.split(r"[|,]", raw) if t.strip()]
         if not tokens:
             continue
-        # Build plural-aware alternation for all synonyms
         alts = [_plural_token_regex(t) for t in tokens]
         alts = [a for a in alts if a]
         if not alts:
             continue
-        # word-boundaries + optional "-shaped"/" shaped"
         pat = r"\b(?:" + "|".join(alts) + 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
@@ -394,14 +388,9 @@ def save_shape_alias_rows(enabled, df_rows, persist):
     # Recompile in-memory, regardless of persist
     global _SHAPE_ALIASES
     if bool(enabled):
-        # Build from the just-edited rows
-        tokens = []
-        compiled = []
-        # Emulate compiler using cleaned rows directly
         compiled = []
         for item in cleaned:
-            raw = item["shape"]
-            name = item["name"]
+            raw = item["shape"]; name = item["name"]
             toks = [t.strip() for t in re.split(r"[|,]", raw) if t.strip()]
             alts = [_plural_token_regex(t) for t in toks]
             alts = [a for a in alts if a]
@@ -414,11 +403,7 @@ def save_shape_alias_rows(enabled, df_rows, persist):
         _SHAPE_ALIASES = []
 
     normalized = [[it["shape"], it["name"]] for it in cleaned] + [["", ""]]
-    return (
-        status,
-        gr.update(value=normalized, row_count=(max(1, len(normalized)), "dynamic"))
-    )
-
+    return status, gr.update(value=normalized, row_count=(max(1, len(normalized)), "dynamic"))
 
 
 # ------------------------------
@@ -490,7 +475,7 @@ def run_batch(
     top_p: float,
     max_tokens: int,
     max_side: int,
-    time_budget_s: float | None = None,                    # respects Zero-GPU window
+    time_budget_s: float | None = None,                    # respects Zero-GPU window (None = unlimited)
     progress: gr.Progress = gr.Progress(track_tqdm=True),  # drives the progress bar
 ) -> Tuple[List[dict], list, list, str, List[str], int, int]:
     """
@@ -527,7 +512,7 @@ def run_batch(
         session_rows.append({"filename": filename, "caption": cap, "path": path, "thumb_path": thumb})
         processed += 1
 
-        if time_budget_s is not None and (time.time() - start) >= float(time_budget_s):
+        if (time_budget_s is not None) and ((time.time() - start) >= float(time_budget_s)):
             leftover = files[idx+1:]
             break
 
@@ -563,7 +548,7 @@ 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]:
-    # Expecting list-of-lists (Dataframe type="array")
+    # Expect list-of-lists (Dataframe type="array")
     tbl = table_value or []
     new = []
     for i, r in enumerate(rows or []):
@@ -686,7 +671,6 @@ def import_captions_file(file_path: str, session_rows: List[dict]) -> Tuple[List
     - Otherwise append a new row (without image path/thumbnail)
     """
     if not file_path or not os.path.exists(file_path):
-        # No change
         table_rows = _rows_to_table(session_rows)
         gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
                          for r in session_rows if (r.get("thumb_path") or r.get("path"))]
@@ -700,7 +684,6 @@ def import_captions_file(file_path: str, session_rows: List[dict]) -> Tuple[List
             with open(file_path, "r", encoding="utf-8") as f:
                 reader = csv.reader(f)
                 rows = list(reader)
-            # Skip header if present
             if rows and len(rows[0]) >= 2 and str(rows[0][0]).lower().strip() == "filename":
                 rows = rows[1:]
             for r in rows:
@@ -731,7 +714,6 @@ def import_captions_file(file_path: str, session_rows: List[dict]) -> Tuple[List
         else:
             return session_rows, _rows_to_table(session_rows), _rows_to_table(session_rows), f"Unsupported file type: {ext}"
     except Exception as e:
-        # Best-effort graceful fallback
         table_rows = _rows_to_table(session_rows)
         gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
                          for r in session_rows if (r.get("thumb_path") or r.get("path"))]
@@ -757,7 +739,7 @@ def import_captions_file(file_path: str, session_rows: List[dict]) -> Tuple[List
 
 
 # ------------------------------
-# 10) UI header helper (logo auto-fit)
+# 10) UI header helper (logo auto-fit to match title/subtitle block)
 # ------------------------------
 def _render_header_html(auto: bool, px: int, scale: float) -> str:
     auto_js = "true" if auto else "false"
@@ -786,10 +768,10 @@ def _render_header_html(auto: bool, px: int, scale: float) -> str:
     if (!logo || !text) return;
     if (AUTO) {{
       const h = text.getBoundingClientRect().height || 200;
-      const target = Math.max(80, Math.min(480, Math.round(h * SCALE)));
+      const target = Math.max(80, Math.min(520, Math.round(h * SCALE)));
       logo.style.height = target + "px";
     }} else {{
-      logo.style.height = Math.max(80, Math.min(480, PX)) + "px";
+      logo.style.height = Math.max(80, Math.min(520, PX)) + "px";
     }}
   }}
   const textNode = document.querySelector(".cf-text");
@@ -827,6 +809,7 @@ BASE_CSS = """
 """
 
 with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
+    # Ensure Spaces sees a GPU function (without touching CUDA in main)
     demo.load(_gpu_startup_warm, inputs=None, outputs=None)
 
     # ---- Header
@@ -882,7 +865,7 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
                 # Logo controls
                 logo_auto = gr.Checkbox(value=settings.get("logo_auto", True),
                                         label="Auto-match logo height to text")
-                logo_px   = gr.Slider(80, 480, value=settings.get("logo_px", 200),
+                logo_px   = gr.Slider(80, 520, value=settings.get("logo_px", 200),
                                       step=4, label="Logo height (px, if Auto off)")
                 logo_scale = gr.Slider(0.7, 1.6, value=settings.get("logo_scale", 1.10),
                                        step=0.02, label="Logo scale × (if Auto on)")
@@ -912,7 +895,6 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
         cfg["dataset_name"] = sanitize_basename(name)
         save_settings(cfg)
         return gr.update()
-
     dataset_name.change(_save_dataset_name, inputs=[dataset_name], outputs=[])
 
     # Header controls live update
@@ -928,68 +910,66 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
     logo_auto.change(_update_header, inputs=[logo_auto, logo_px, logo_scale], outputs=[header_html])
     logo_scale.change(_update_header, inputs=[logo_auto, logo_px, logo_scale], outputs=[header_html])
 
-    # ---- Shape Aliases block (placed WITH other settings, before uploads)
-    # ---- Shape Aliases block (placed WITH other settings, before uploads)
-    # ---- Shape Aliases block (placed WITH other settings, before uploads)
-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, also catches simple plurals (`box`→`boxes`, `lady`→`ladies`), "
-        "and `*-shaped` or `* shaped` variants."
-    )
-    init_rows, init_enabled = get_shape_alias_rows_ui_defaults()
-    enable_aliases = gr.Checkbox(label="Enable shape alias replacements", value=init_enabled)
-    persist_aliases = gr.Checkbox(
-        label="Save aliases across sessions",
-        value=load_settings().get("shape_aliases_persist", True)
-    )
-    alias_table = gr.Dataframe(
-        headers=["shape (token or synonyms)", "name to insert"],
-        value=init_rows,
-        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")
-
-    # status line for saves
-    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])
-
-    # Persist the "persist" toggle itself
-    def _save_alias_persist_flag(v):
-        cfg = load_settings()
-        cfg["shape_aliases_persist"] = bool(v)
-        save_settings(cfg)
-        return gr.update()
+    # ---- Shape Aliases (with plural matching + persist) ----
+    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, catches simple plurals (`box`→`boxes`, `lady`→`ladies`), "
+            "and also matches `*-shaped` or `* shaped` variants."
+        )
 
-    persist_aliases.change(_save_alias_persist_flag, inputs=[persist_aliases], outputs=[])
+        init_rows, init_enabled = get_shape_alias_rows_ui_defaults()
+        enable_aliases = gr.Checkbox(label="Enable shape alias replacements", value=init_enabled)
+        persist_aliases = gr.Checkbox(
+            label="Save aliases across sessions",
+            value=load_settings().get("shape_aliases_persist", True)
+        )
 
-    # Save/apply rows (persists to disk only if persist_aliases is checked)
-    save_btn.click(
-        save_shape_alias_rows,
-        inputs=[enable_aliases, alias_table, persist_aliases],
-        outputs=[save_status, alias_table]
-    )
+        alias_table = gr.Dataframe(
+            headers=["shape (token or synonyms)", "name to insert"],
+            value=init_rows,
+            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")
+
+        # status line for saves
+        save_status = gr.Markdown("")
+
+        # --- local handlers (must stay inside Blocks context) ---
+        def _add_row(cur):
+            cur = (cur or []) + [["", ""]]
+            return gr.update(value=cur, row_count=(max(1, len(cur)), "dynamic"))
+
+        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])
+
+        def _save_alias_persist_flag(v):
+            cfg = load_settings()
+            cfg["shape_aliases_persist"] = bool(v)
+            save_settings(cfg)
+            return gr.update()
+        persist_aliases.change(_save_alias_persist_flag, inputs=[persist_aliases], outputs=[])
+
+        # Persist rows if persist_aliases checked; otherwise apply in-memory only
+        save_btn.click(
+            save_shape_alias_rows,
+            inputs=[enable_aliases, alias_table, persist_aliases],
+            outputs=[save_status, alias_table]
+        )
 
     # ---- Tabs: Single & Batch
     with gr.Tabs():
@@ -1009,7 +989,7 @@ with gr.Accordion("Shape Aliases", open=False):
             run_button = gr.Button("Caption batch", variant="primary")
 
             with gr.Accordion("Import captions from CSV/XLSX (merge by filename)", open=False):
-                import_file = gr.File(label="Choose .csv or .xlsx", file_types=[".csv", ".xlsx"], type="filepath")
+                import_file = gr.File(label="Choose .csv or .xlsx", file_types=["file"], type="filepath")
                 import_btn = gr.Button("Import into current session")
 
     # ---- Results area (gallery left / table right)
@@ -1034,7 +1014,7 @@ with gr.Accordion("Shape Aliases", open=False):
                 headers=["filename", "caption"],
                 interactive=True,
                 wrap=True,
-                type="array",          # <— prevents DataFrame truth ambiguity
+                type="array",          # prevents pandas truth ambiguity
                 elem_id="cfTable"
             )
 
@@ -1186,19 +1166,13 @@ with gr.Accordion("Shape Aliases", open=False):
         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])
 
     # ---- Import hook
     def _do_import(fpath, rows):
         new_rows, gal, tbl, stamp = import_captions_file(fpath, rows or [])
         return new_rows, gal, tbl, stamp
-
-    import_btn.click(
-        _do_import,
-        inputs=[import_file, rows_state],
-        outputs=[rows_state, gallery, table, autosave_md]
-    )
+    import_btn.click(_do_import, inputs=[import_file, rows_state], outputs=[rows_state, gallery, table, autosave_md])
 
     # ---- Exports
     export_csv_btn.click(
@@ -1225,5 +1199,6 @@ if __name__ == "__main__":
         ssr_mode=False,
         debug=True,
         show_error=True,
-        allowed_paths=[THUMB_CACHE, EXCEL_THUMB_DIR, TXT_EXPORT_DIR],  # serve /tmp caches safely
+        # Allow Gradio to serve generated files from /tmp caches
+        allowed_paths=[THUMB_CACHE, EXCEL_THUMB_DIR, TXT_EXPORT_DIR],
     )