| """Offline: rewrite a training metadata file's raw prompts into structured JSON captions. |
| |
| Run this once before training if your dataset's prompt column holds free-form prose, then |
| train on the rewritten metadata. The rewriter is a separate VLM + stage-2 LoRA adapter |
| (not the DiT); provide it via --base/--adapter or REWRITER_BASE_MODEL/REWRITER_ADAPTER. |
| |
| Usage: |
| python rewrite_captions.py --metadata metadata.csv --output metadata_rewritten.csv \ |
| --base /path/to/rewriter-base --adapter /path/to/rewriter-step2-lora --duration 5 |
| |
| Supports .csv / .json / .jsonl. Rows whose stage-2 output fails to parse keep their |
| original prompt and are logged, so training never silently uses a broken row. |
| """ |
|
|
| import argparse |
| import json |
| import os |
|
|
| try: |
| from .prompt_rewriter import Rewriter, make_backend |
| except ImportError: |
| from prompt_rewriter import Rewriter, make_backend |
|
|
|
|
| def _load_rows(path): |
| if path.endswith(".json"): |
| with open(path, "r", encoding="utf-8") as f: |
| return json.load(f), "json" |
| if path.endswith(".jsonl"): |
| rows = [] |
| with open(path, "r", encoding="utf-8") as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| rows.append(json.loads(line)) |
| return rows, "jsonl" |
| import pandas |
| df = pandas.read_csv(path) |
| return [df.iloc[i].to_dict() for i in range(len(df))], "csv" |
|
|
|
|
| def _save_rows(rows, path, fmt): |
| if fmt == "json": |
| with open(path, "w", encoding="utf-8") as f: |
| json.dump(rows, f, ensure_ascii=False, indent=2) |
| elif fmt == "jsonl": |
| with open(path, "w", encoding="utf-8") as f: |
| for row in rows: |
| f.write(json.dumps(row, ensure_ascii=False) + "\n") |
| else: |
| import pandas |
| pandas.DataFrame(rows).to_csv(path, index=False) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--metadata", required=True, help="input metadata (.csv/.json/.jsonl)") |
| ap.add_argument("--output", required=True, help="output metadata with rewritten captions") |
| ap.add_argument("--prompt-column", default="prompt") |
| ap.add_argument("--mode", default="t2v", choices=["t2v", "ti2v", "t2i"]) |
| ap.add_argument("--duration", type=float, default=5) |
| ap.add_argument("--backend", default="transformers") |
| ap.add_argument("--base", default=None, help="rewriter base model (or REWRITER_BASE_MODEL)") |
| ap.add_argument("--adapter", default=None, help="rewriter stage-2 LoRA (or REWRITER_ADAPTER)") |
| ap.add_argument("--first-frame-column", default=None, |
| help="for ti2v: column holding the first-frame image path/URL") |
| args = ap.parse_args() |
|
|
| rows, fmt = _load_rows(args.metadata) |
| rewriter = Rewriter(make_backend(args.backend, args.base, args.adapter)) |
|
|
| ok, failed = 0, 0 |
| for i, row in enumerate(rows): |
| raw = row.get(args.prompt_column, "") |
| first_frame = row.get(args.first_frame_column) if args.first_frame_column else None |
| result = rewriter.rewrite(raw, mode=args.mode, first_frame=first_frame, duration=args.duration) |
| if result["json"] is not None: |
| row[args.prompt_column] = json.dumps(result["json"], ensure_ascii=False, separators=(",", ":")) |
| ok += 1 |
| else: |
| failed += 1 |
| print(f"[warn] row {i}: stage-2 JSON did not parse; keeping original prompt.") |
| if (i + 1) % 20 == 0: |
| print(f" ... {i + 1}/{len(rows)} rewritten") |
|
|
| _save_rows(rows, args.output, fmt) |
| print(f"[done] {ok} rewritten, {failed} kept-as-is -> {os.path.abspath(args.output)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
