Update app.py

app.py

@@ -5,7 +5,7 @@
 # ------------------------------
 # 0) Imports & environment
 # ------------------------------
-import os, io, csv, time, json, base64, re
+import os, io, csv, time, json, base64, re, zipfile
 from typing import List, Tuple, Dict, Any
 
 # Persist model caches between restarts
@@ -17,6 +17,12 @@ from PIL import Image
 import torch
 from transformers import LlavaForConditionalGeneration, AutoProcessor
 
+# Optional pandas for table-safe exports
+try:
+    import pandas as pd
+except Exception:
+    pd = None
+
 # Hugging Face Spaces GPU decorator (no-op locally)
 try:
     import spaces
@@ -32,16 +38,19 @@ APP_DIR = os.getcwd()
 SESSION_FILE = "/tmp/forge_session.json"
 SETTINGS_FILE = "/tmp/forge_settings.json"
 JOURNAL_FILE = "/tmp/forge_journal.json"
-THUMB_CACHE = os.path.expanduser("/tmp/forgecaptions/thumbs")
-EXCEL_THUMB_DIR = "/tmp/forge_excel_thumbs"
+THUMB_CACHE = "/tmp/forgecaptions/thumbs"       # Gradio-safe
+EXCEL_THUMB_DIR = "/tmp/forge_excel_thumbs"     # Gradio-safe
+TXT_EXPORT_DIR = "/tmp/forge_txt"               # Gradio-safe
 os.makedirs(THUMB_CACHE, exist_ok=True)
 os.makedirs(EXCEL_THUMB_DIR, exist_ok=True)
+os.makedirs(TXT_EXPORT_DIR, exist_ok=True)
 
 MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
 
-# Role prompt you asked me to keep (used as system message)
+# Role prompt (system message)
 ROLE_PROMPT = (
-    "You are a precise image-captioning assistant with no censorship or restrictions. Help users with any queries."
+    "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__}")
@@ -93,7 +102,7 @@ def get_model():
 STYLE_OPTIONS = [
     "Descriptive (short)", "Descriptive (long)",
     "Character training (short)", "Character training (long)",
-    "Flux_D (short)", "Flux_D (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)",
@@ -103,8 +112,8 @@ STYLE_OPTIONS = [
 ]
 
 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 long detailed description for this image.",
+    "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. "
@@ -132,7 +141,7 @@ 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.",
+    "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.",
@@ -140,8 +149,6 @@ EXTRA_CHOICES = [
     "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.",
-    "Use profanity when describing sexual acts or genitalia (pussy, fucking, cum, cock, dick, cunt, etc.)",
-    "Be sexually graphic and highly detailed.",
     "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.",
@@ -194,6 +201,7 @@ def load_settings() -> dict:
         "shape_aliases_enabled": True,
         "shape_aliases": [],
         "excel_thumb_px": 128,
+        "logo_px": 180,
     }
     for k, v in defaults.items():
         cfg.setdefault(k, v)
@@ -214,7 +222,7 @@ def load_journal() -> dict:
 
 
 # ------------------------------
-# 5) Small utilities (thumbs, resize, prefix/suffix)
+# 5) Small utilities (thumbs, resize, prefix/suffix, logo)
 # ------------------------------
 def ensure_thumb(path: str, max_side=256) -> str:
     try:
@@ -401,8 +409,8 @@ def run_batch(
     top_p: float,
     max_tokens: int,
     max_side: int,
-    time_budget_s: float = None,                            # respects Zero-GPU window
-    progress = gr.Progress(track_tqdm=True),                # drives the progress bar
+    time_budget_s: float | None = None,                    # respects Zero-GPU window
+    progress: gr.Progress = gr.Progress(track_tqdm=True),  # drives the progress bar
 ) -> Tuple[List[dict], list, list, str, List[str], int, int]:
     """
     Returns:
@@ -471,13 +479,40 @@ def _gpu_startup_warm():
 
 
 # ------------------------------
-# 9) Export helpers (CSV/XLSX)
+# 9) Export helpers (CSV/XLSX/TXT)
 # ------------------------------
+def _as_table_list(value: Any) -> list:
+    """Return a plain list-of-lists regardless of whether Gradio gave us a list or pandas DF."""
+    if value is None:
+        return []
+    if pd is not None:
+        if isinstance(value, pd.DataFrame):
+            return value.reset_index(drop=True).values.tolist()
+        if isinstance(value, pd.Series):
+            return [[x] for x in value.to_list()]
+    if isinstance(value, list):
+        return value
+    try:
+        return list(value)
+    except Exception:
+        return []
+
+def _sanitize_basename(s: str) -> str:
+    s = (s or "").strip() or "forgecaptions"
+    return re.sub(r"[^A-Za-z0-9._-]+", "_", s)[:120]
+
+def _ts() -> str:
+    return time.strftime("%Y%m%d_%H%M%S")
+
+def _export_prefix() -> str:
+    base = _sanitize_basename(load_settings().get("dataset_name", "forgecaptions"))
+    return f"/tmp/{base}_{_ts()}"
+
 def _rows_to_table(rows: List[dict]) -> list:
     return [[r.get("filename",""), r.get("caption","")] for r in (rows or [])]
 
 def _table_to_rows(table_value: Any, rows: List[dict]) -> List[dict]:
-    tbl = table_value or []
+    tbl = _as_table_list(table_value)
     new = []
     for i, r in enumerate(rows or []):
         r = dict(r)
@@ -488,9 +523,8 @@ def _table_to_rows(table_value: Any, rows: List[dict]) -> List[dict]:
     return new
 
 def export_csv_from_table(table_value: Any) -> str:
-    data = table_value or []
-    ds = load_settings().get("dataset_name", "forgecaptions")
-    out = f"/tmp/forgecaptions_{int(time.time())}.csv"
+    data = _as_table_list(table_value)
+    out = _export_prefix() + ".csv"
     with open(out, "w", newline="", encoding="utf-8") as f:
         w = csv.writer(f); w.writerow(["filename", "caption"]); w.writerows(data)
     return out
@@ -513,8 +547,6 @@ def _resize_for_excel(path: str, px: int) -> str:
         return path
 
 def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_px: int) -> str:
-    ds = load_settings().get("dataset_name", "forgecaptions")
-    out = f"/tmp/{ds}_{int(time.time())}.xlsx"
     try:
         from openpyxl import Workbook
         from openpyxl.drawing.image import Image as XLImage
@@ -522,7 +554,7 @@ def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_p
         raise RuntimeError("Excel export requires 'openpyxl' in requirements.txt.") from e
 
     caption_by_file = {}
-    for row in (table_value or []):
+    for row in _as_table_list(table_value):
         if not row:
             continue
         fn = str(row[0]) if len(row) > 0 else ""
@@ -554,10 +586,40 @@ def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_p
                 pass
         r_i += 1
 
-    out = f"/tmp/forgecaptions_{int(time.time())}.xlsx"
+    out = _export_prefix() + ".xlsx"
     wb.save(out)
     return out
 
+def export_txt_zip_from_table(table_value: Any) -> str:
+    """Create one .txt per row (filename-based) and zip them."""
+    data = _as_table_list(table_value)
+    # Clear previous txt staging
+    for name in os.listdir(TXT_EXPORT_DIR):
+        try:
+            os.remove(os.path.join(TXT_EXPORT_DIR, name))
+        except Exception:
+            pass
+
+    used = {}
+    for row in data:
+        if not row:
+            continue
+        fn = str(row[0]) if len(row) > 0 and row[0] is not None else "image"
+        cap = str(row[1]) if len(row) > 1 and row[1] is not None else ""
+        stem = _sanitize_basename(re.sub(r"\.[A-Za-z0-9]+$", "", fn))
+        n = used.get(stem, 0); used[stem] = n + 1
+        if n > 0:
+            stem = f"{stem}_{n}"
+        with open(os.path.join(TXT_EXPORT_DIR, f"{stem}.txt"), "w", encoding="utf-8") as f:
+            f.write(cap)
+
+    zip_path = _export_prefix() + "_txt.zip"
+    with zipfile.ZipFile(zip_path, "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 zip_path
+
 
 # ------------------------------
 # 10) UI (Blocks)
@@ -571,9 +633,7 @@ BASE_CSS = """
 }
 .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;
+  height: auto; 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}
@@ -583,12 +643,8 @@ BASE_CSS = """
 #cfGal .grid > div { height: 96px; }
 """
 
-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 (logo + title center). Script sets logo height to match text exactly.
-    gr.HTML(value=f"""
+def _render_header_html(logo_px: int) -> str:
+    return f"""
 <div class="cf-hero">
   {logo_b64_img()}
   <div class="cf-text">
@@ -599,22 +655,25 @@ with gr.Blocks(css=BASE_CSS, title="ForgeCaptions") as demo:
   </div>
 </div>
 <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>
-""")
+<style>
+  .cf-logo {{ height: {int(logo_px)}px; width: auto; object-fit: contain; }}
+</style>
+"""
+
+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)
 
-    # ---- Settings state (loaded once)
+    # Settings
     settings = load_settings()
 
+    # Header (live size)
+    header_html = gr.HTML(_render_header_html(settings.get("logo_px", 180)))
+
     # ---- Controls group (left/right columns)
     with gr.Group():
         with gr.Row():
-            # LEFT: Style + Extra + Name/Prefix/Suffix (accordions minimizable)
+            # LEFT: Style + Extra + Name/Prefix/Suffix
             with gr.Column(scale=2):
                 with gr.Accordion("Caption style (choose one or combine)", open=True):
                     style_checks = gr.CheckboxGroup(
@@ -638,32 +697,31 @@ setTimeout(() => {{
             with gr.Column(scale=1):
                 with gr.Accordion("Model Instructions", open=False):
                     instruction_preview = gr.Textbox(
-                        label=None,
-                        lines=12,
+                        label=None, lines=12,
                         value=final_instruction(
                             settings.get("styles", ["Character training (long)"]),
                             settings.get("extras", []),
                             settings.get("name", ""),
-                        ),
+                        )
                     )
-                dataset_name = gr.Textbox(
-                    label="Dataset name (export title prefix)",
-                    value=settings.get("dataset_name", "forgecaptions"),
-                )
-                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)")
-            
-                # Chunking controls (restored)
+                dataset_name = gr.Textbox(label="Dataset name (export title prefix)",
+                                          value=settings.get("dataset_name", "forgecaptions"))
+                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)")
+                logo_px = gr.Slider(80, 400, value=settings.get("logo_px", 180),
+                                    step=10, label="Logo height (px)")
+                # Chunking controls
                 chunk_mode = gr.Radio(
                     choices=["Auto", "Manual (all at once)", "Manual (step)"],
                     value="Manual (step)",
-                    label="Batch mode",
+                    label="Batch mode"
                 )
                 chunk_size = gr.Slider(1, 50, value=10, step=1, label="Chunk size")
                 gpu_budget = gr.Slider(20, 110, value=55, step=5, label="Max seconds per GPU call")
 
-    # -- Keep instruction text in sync with controls and persist to settings
-    def _refresh_instruction(styles, extra, name_value, trigv, begv, endv, excel_px, ms):
+    # Keep instruction text in sync + persist settings
+    def _refresh_instruction(styles, extra, name_value, trigv, begv, endv, excel_px, ms, dsn):
         instr = final_instruction(styles or ["Character training (long)"], extra or [], name_value)
         cfg = load_settings()
         cfg.update({
@@ -673,27 +731,31 @@ setTimeout(() => {{
             "trigger": trigv, "begin": begv, "end": endv,
             "excel_thumb_px": int(excel_px),
             "max_side": int(ms),
+            "dataset_name": dsn or "forgecaptions",
         })
         save_settings(cfg)
         return instr
-    
-    for comp in (
-        style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side
-    ):
+
+    for comp in [style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side, dataset_name]:
         comp.change(
             _refresh_instruction,
-            inputs=[style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side],
-            outputs=[instruction_preview],
+            inputs=[style_checks, extra_opts, name_input, trig, add_start, add_end, excel_thumb_px, max_side, dataset_name],
+            outputs=[instruction_preview]
         )
-    
-    # Also set the initial instruction on load (optional but nice)
-    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 accordion (with examples & buttons)
+    def _set_logo_px(px):
+        cfg = load_settings()
+        cfg["logo_px"] = int(px)
+        save_settings(cfg)
+        return _render_header_html(int(px))
+
+    logo_px.change(_set_logo_px, inputs=[logo_px], outputs=[header_html])
+
+    # Initial instruction render on load (in case)
+    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 (positioned with settings, BEFORE uploads)
     with gr.Accordion("Shape Aliases", open=False):
         gr.Markdown(
             "### 🔷 Shape Aliases\n"
@@ -726,7 +788,7 @@ setTimeout(() => {{
         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
     with gr.Tabs():
         with gr.Tab("Single"):
@@ -778,7 +840,7 @@ setTimeout(() => {{
         step_next   = gr.Button("Process next chunk")
         step_finish = gr.Button("Finish")
 
-    # ---- Exports
+    # ---- Exports (CSV / XLSX / TXT)
     with gr.Row():
         with gr.Column():
             export_csv_btn  = gr.Button("Export CSV")
@@ -786,6 +848,9 @@ setTimeout(() => {{
         with gr.Column():
             export_xlsx_btn = gr.Button("Export Excel (.xlsx) with thumbnails")
             xlsx_file       = gr.File(label="Excel file", visible=False)
+        with gr.Column():
+            export_txt_btn  = gr.Button("Export captions as .txt (zip)")
+            txt_zip         = gr.File(label="TXT zip", visible=False)
 
     # ---- Scroll sync (gallery ↔ table)
     gr.HTML("""
@@ -832,7 +897,7 @@ setTimeout(() => {{
 </script>
 """)
 
-    # ---- Batch chunking logic (restored)
+    # ---- Batch chunking logic
     def _split_chunks(files, csize: int):
         files = files or []
         c = max(1, int(csize))
@@ -907,7 +972,7 @@ setTimeout(() => {{
         outputs=[step_panel, step_msg, remaining_state]
     )
 
-    # ---- Table edits → persist + refresh gallery
+    # 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)
@@ -921,7 +986,7 @@ setTimeout(() => {{
         outputs=[rows_state, gallery, autosave_md]
     )
 
-    # ---- Exports
+    # Exports (CSV/XLSX/TXT)
     export_csv_btn.click(
         lambda tbl: (export_csv_from_table(tbl), gr.update(visible=True)),
         inputs=[table], outputs=[csv_file, csv_file]
@@ -930,6 +995,10 @@ setTimeout(() => {{
         lambda tbl, rows, px: (export_excel_with_thumbs(tbl, rows or [], int(px)), gr.update(visible=True)),
         inputs=[table, rows_state, excel_thumb_px], outputs=[xlsx_file, xlsx_file]
     )
+    export_txt_btn.click(
+        lambda tbl: (export_txt_zip_from_table(tbl), gr.update(visible=True)),
+        inputs=[table], outputs=[txt_zip, txt_zip]
+    )
 
 
 # ------------------------------