"
return ""
# ------------------------------
# 6) Shape Aliases (plural-aware + '-shaped' variants)
# ------------------------------
def _plural_token_regex(tok: str) -> str:
t = (tok or "").strip()
if not t: return ""
t_low = t.lower()
if re.search(r"[^aeiou]y$", t_low):
return re.escape(t[:-1]) + r"(?:y|ies)"
if re.search(r"(?:s|x|z|ch|sh)$", t_low):
return re.escape(t) + r"(?:es)?"
return re.escape(t) + r"s?"
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", []):
raw = (item.get("shape") or "").strip()
name = (item.get("name") or "").strip()
if not raw or not name:
continue
tokens = [t.strip() for t in re.split(r"[|,]", raw) if t.strip()]
if not tokens:
continue
alts = [_plural_token_regex(t) for t in tokens]
alts = [a for a in alts if a]
if not alts:
continue
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
_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
def get_shape_alias_rows_ui_defaults():
s = load_settings()
rows = [[it.get("shape",""), it.get("name","")] for it in s.get("shape_aliases", [])]
enabled = bool(s.get("shape_aliases_enabled", True))
if not rows:
rows = [["", ""]]
return rows, enabled
def save_shape_alias_rows(enabled, df_rows, persist):
cleaned = []
for r in (df_rows or []):
if not r:
continue
shape = (r[0] or "").strip()
name = (r[1] or "").strip()
if shape and name:
cleaned.append({"shape": shape, "name": name})
status = "✅ Applied for this session only."
if persist:
cfg = load_settings()
cfg["shape_aliases_enabled"] = bool(enabled)
cfg["shape_aliases"] = cleaned
save_settings(cfg)
status = "💾 Saved to disk (will persist across restarts)."
global _SHAPE_ALIASES
if bool(enabled):
compiled = []
for item in cleaned:
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]
if not alts:
continue
pat = r"\b(?:" + "|".join(alts) + r")(?:[-\s]?shaped)?\b"
compiled.append((re.compile(pat, flags=re.I), name))
_SHAPE_ALIASES = compiled
else:
_SHAPE_ALIASES = []
normalized = [[it["shape"], it["name"]] for it in cleaned] + [["", ""]]
return status, gr.update(value=normalized, row_count=(max(1, len(normalized)), "dynamic"))
# ------------------------------
# 7) Prompt builder
# ------------------------------
def final_instruction(style_list: List[str], extra_opts: List[str], name_value: str, length_choice: str = "long") -> str:
styles = style_list or ["Character training"]
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())
if "Aesthetic tags (comma-sep)" not in styles:
lh = _length_hint(length_choice or "any")
if lh:
core += " " + lh
return core
# ------------------------------
# 8) GPU caption functions (Spaces-only)
# ------------------------------
def _build_inputs(im: Image.Image, instr: str, dtype) -> Dict[str, Any]:
convo = [
{"role": "system", "content": ROLE_PROMPT},
{"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")
if "pixel_values" in inputs:
inputs["pixel_values"] = inputs["pixel_values"].to(dtype)
return inputs
@spaces.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_core(im, instr, s)
return cap
@spaces.GPU()
@torch.no_grad()
def run_batch(
files: List[Any],
session_rows: List[dict],
instr_text: str,
temp: float,
top_p: float,
max_tokens: int,
max_side: int,
time_budget_s: float | None = None,
progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> Tuple[List[dict], list, list, str, List[str], int, int]:
return run_batch_core(files, session_rows, instr_text, temp, top_p, max_tokens, max_side, time_budget_s, progress)
# Optional tiny probe to satisfy strict scanners (not called)
@spaces.GPU()
def _gpu_probe() -> str:
return "ok"
# ---- shared core routines used by both GPU functions ----
def caption_once_core(im: Image.Image, instr: str, settings: dict) -> str:
cap = caption_once(
im, instr,
settings.get("temperature", 0.6),
settings.get("top_p", 0.9),
settings.get("max_tokens", 256),
)
cap = apply_shape_aliases(cap)
cap = apply_prefix_suffix(cap, settings.get("trigger",""), settings.get("begin",""), settings.get("end",""))
return cap
@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()
inputs = _build_inputs(im, instr, dtype)
inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
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,
use_cache=True,
)
gen_ids = out[0, inputs["input_ids"].shape[1]:]
return processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
def run_batch_core(
files: List[Any],
session_rows: List[dict],
instr_text: str,
temp: float,
top_p: float,
max_tokens: int,
max_side: int,
time_budget_s: float | None,
progress: gr.Progress,
) -> Tuple[List[dict], list, list, str, List[str], int, int]:
session_rows = session_rows or []
files = [f for f in (files or []) if f and os.path.exists(f)]
total = len(files)
processed = 0
if total == 0:
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')}", [], 0, 0
start = time.time()
leftover: List[str] = []
for idx, path in enumerate(progress.tqdm(files, desc="Captioning")):
try:
im = Image.open(path).convert("RGB")
except Exception:
continue
im = resize_for_model(im, max_side)
cap = caption_once(im, instr_text, temp, top_p, max_tokens)
cap = apply_shape_aliases(cap)
s = load_settings()
cap = apply_prefix_suffix(cap, s.get("trigger",""), s.get("begin",""), s.get("end",""))
filename = os.path.basename(path)
thumb = ensure_thumb(path, 256)
session_rows.append({"filename": filename, "caption": cap, "path": path, "thumb_path": thumb})
processed += 1
if (time_budget_s is not None) and ((time.time() - start) >= float(time_budget_s)):
leftover = files[idx+1:]
break
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"))]
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')}",
leftover,
processed,
total,
)
# ------------------------------
# 9) Export helpers (CSV/XLSX/TXT ZIP)
# ------------------------------
def _rows_to_table(rows: List[dict]) -> list:
return [[r.get("filename",""), r.get("caption","")] for r in (rows or [])]
def _table_to_rows(table_value: Any, rows: List[dict]) -> List[dict]:
tbl = table_value or []
new = []
for i, r in enumerate(rows or []):
r = dict(r)
if i < len(tbl) and len(tbl[i]) >= 2:
r["filename"] = str(tbl[i][0]) if tbl[i][0] is not None else r.get("filename","")
r["caption"] = str(tbl[i][1]) if tbl[i][1] is not None else r.get("caption","")
new.append(r)
return new
def export_csv_from_table(table_value: Any, dataset_name: str) -> str:
data = table_value or []
name = sanitize_basename(dataset_name)
out = f"/tmp/{name}_{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)
return out
def _resize_for_excel(path: str, px: int) -> str:
try:
im = Image.open(path).convert("RGB")
except Exception:
return path
w, h = im.size
if max(w, h) > px:
s = px / max(w, h)
im = im.resize((int(w*s), int(h*s)), Image.LANCZOS)
base = os.path.basename(path)
out_path = os.path.join(EXCEL_THUMB_DIR, f"{os.path.splitext(base)[0]}_{px}px.jpg")
try:
im.save(out_path, "JPEG", quality=85, optimize=True)
return out_path
except Exception:
return path
def export_excel_with_thumbs(table_value: Any, session_rows: List[dict], thumb_px: int, dataset_name: 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
caption_by_file = {}
for row in (table_value or []):
if not row:
continue
fn = str(row[0]) if len(row) > 0 else ""
cap = str(row[1]) if len(row) > 1 and row[1] is not None else ""
if fn:
caption_by_file[fn] = cap
wb = Workbook(); ws = wb.active; ws.title = "ForgeCaptions"
ws.append(["image", "filename", "caption"])
ws.column_dimensions["A"].width = 24
ws.column_dimensions["B"].width = 42
ws.column_dimensions["C"].width = 100
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",""))
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:
resized = _resize_for_excel(img_path, int(thumb_px))
xlimg = XLImage(resized)
ws.add_image(xlimg, f"A{r_i}")
ws.row_dimensions[r_i].height = row_h
except Exception:
pass
r_i += 1
name = sanitize_basename(dataset_name)
out = f"/tmp/{name}_{int(time.time())}.xlsx"
wb.save(out)
return out
def export_txt_zip(table_value: Any, dataset_name: str) -> str:
"""
Create one .txt per caption, zip them.
"""
data = table_value or []
# wipe old
for fn in os.listdir(TXT_EXPORT_DIR):
try:
os.remove(os.path.join(TXT_EXPORT_DIR, fn))
except Exception:
pass
used: Dict[str,int] = {}
for row in data:
if not row:
continue
orig = (row[0] or "item").strip() if len(row) > 0 else "item"
stem = re.sub(r"\.[A-Za-z0-9]+$", "", orig)
stem = sanitize_basename(stem or "item")
if stem in used:
used[stem] += 1
stem = f"{stem}_{used[stem]}"
else:
used[stem] = 0
cap = (row[1] or "").strip() if len(row) > 1 and row[1] is not None else ""
with open(os.path.join(TXT_EXPORT_DIR, f"{stem}.txt"), "w", encoding="utf-8") as f:
f.write(cap)
name = sanitize_basename(dataset_name)
zpath = f"/tmp/{name}_{int(time.time())}_txt.zip"
with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as z:
for fn in os.listdir(TXT_EXPORT_DIR):
if fn.endswith(".txt"):
z.write(os.path.join(TXT_EXPORT_DIR, fn), arcname=fn)
return zpath
# ------------------------------
# 10) UI header helper (fixed logo size)
# ------------------------------
def _render_header_html(px: int) -> str:
return f"""