File size: 5,279 Bytes
67f7f67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """Multi-reference a2a generation for the AKUSPACE LoRA — one model load, many sources.
ltx_a2a_generate.py pins a single --reference to every prompt, so an N-source
sweep costs N model loads per checkpoint. But validation samples carry their own
conditions[].audio, so one load can cover every (source, prompt) pair. Loading
dominates runtime, so this is the difference between 15 loads and 3.
Input is a TSV manifest (label<TAB>reference_wav<TAB>prompt), which avoids the
shell-quoting traps that bit us with long caption strings.
Outputs land as samples/step_000001_<n>.wav in prompt order; this renames them
to <label>.wav so a 60-clip sweep stays readable.
Trigger note: matches ltx_a2a_generate.py — AKUSPACE is prepended here because
validation_runner does not add it, while training captions had it.
"""
import argparse
import copy
import subprocess
import sys
from pathlib import Path
import yaml
DEFAULT_TRIGGER = "AKUSPACE"
NEVER_SAVE = 10**9
def read_manifest(path: Path) -> list[tuple[str, str, str]]:
rows = []
for n, line in enumerate(path.read_text().splitlines(), 1):
line = line.rstrip("\n")
if not line.strip() or line.lstrip().startswith("#"):
continue
parts = line.split("\t")
if len(parts) != 3:
raise SystemExit(f"{path}:{n}: expected 3 tab-separated fields, got {len(parts)}")
label, ref, prompt = (p.strip() for p in parts)
if not Path(ref).exists():
raise SystemExit(f"{path}:{n}: reference not found: {ref}")
rows.append((label, ref, prompt))
if not rows:
raise SystemExit(f"{path}: no rows")
return rows
def build_config(args, rows) -> dict:
cfg = yaml.safe_load(Path(args.base_config).read_text())
cfg["model"]["load_checkpoint"] = str(Path(args.checkpoint).resolve())
cfg.setdefault("checkpoints", {})["no_resume"] = True
cfg["checkpoints"]["interval"] = NEVER_SAVE
cfg["optimization"]["steps"] = 1
samples = []
for _label, ref, prompt in rows:
text = prompt if args.no_trigger else f"{args.trigger} {prompt}"
samples.append(
{
"prompt": text,
"conditions": [{"type": "reference", "audio": str(Path(ref).resolve())}],
}
)
validation = copy.deepcopy(cfg.get("validation", {}))
validation["samples"] = samples
validation["interval"] = 1
if args.inference_steps:
validation["inference_steps"] = args.inference_steps
if args.seed is not None:
validation["seed"] = args.seed
cfg["validation"] = validation
cfg["output_dir"] = str(Path(args.output_dir).resolve())
return cfg
def rename_outputs(out: Path, rows) -> int:
"""samples/step_000001_<n>.wav -> <label>.wav, n is 1-based prompt order."""
sample_dir = out / "samples"
renamed = 0
for idx, (label, _ref, _prompt) in enumerate(rows, 1):
matches = sorted(sample_dir.glob(f"step_*_{idx}.wav"))
if not matches:
print(f" MISSING output for #{idx} {label}")
continue
dest = out / f"{label}.wav"
dest.write_bytes(matches[-1].read_bytes())
renamed += 1
return renamed
def main() -> int:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--checkpoint", required=True)
p.add_argument("--manifest", required=True, help="TSV: label<TAB>reference<TAB>prompt")
p.add_argument("--output-dir", required=True)
p.add_argument("--base-config", required=True)
p.add_argument("--trigger", default=DEFAULT_TRIGGER)
p.add_argument("--no-trigger", action="store_true")
p.add_argument("--inference-steps", type=int, default=None)
p.add_argument("--seed", type=int, default=None)
p.add_argument("--ltx-repo", default="/workspace/LTX-2.5-repo")
p.add_argument("--dry-run", action="store_true")
args = p.parse_args()
for path in (args.checkpoint, args.base_config, args.manifest):
if not Path(path).exists():
print(f"missing: {path}")
return 2
rows = read_manifest(Path(args.manifest))
cfg = build_config(args, rows)
out = Path(args.output_dir)
out.mkdir(parents=True, exist_ok=True)
cfg_path = out / "generate_config.yaml"
cfg_path.write_text(yaml.safe_dump(cfg, sort_keys=False))
refs = sorted({r[1] for r in rows})
print(f"config : {cfg_path}")
print(f"checkpoint : {cfg['model']['load_checkpoint']}")
print(f"prompts : {len(rows)} across {len(refs)} reference(s), ONE model load")
print(f"trigger : {'off' if args.no_trigger else args.trigger}")
for i, (label, ref, _prompt) in enumerate(rows, 1):
print(f" {i:3d}. {label:38s} <- {Path(ref).name}")
if args.dry_run:
return 0
cmd = ["uv", "run", "python", "scripts/train.py", str(cfg_path)]
cwd = Path(args.ltx_repo) / "packages" / "ltx-trainer"
print(f"\nrunning: {' '.join(cmd)} (cwd={cwd})")
result = subprocess.run(cmd, cwd=cwd)
if result.returncode == 0:
n = rename_outputs(out, rows)
print(f"\nlabelled {n}/{len(rows)} clips in {out}")
return result.returncode
if __name__ == "__main__":
sys.exit(main())
|