| """`uvw` command line: environment, weight hashing and frozen-config generation.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from dataclasses import asdict |
| from pathlib import Path |
|
|
| from uvw import RUNS |
|
|
|
|
| def _cmd_env(args: argparse.Namespace) -> int: |
| from uvw.provenance import environment |
|
|
| print(json.dumps(environment(), indent=2)) |
| return 0 |
|
|
|
|
| def _cmd_hash(args: argparse.Namespace) -> int: |
| from uvw.models import load_pipeline, load_registry, weight_hashes |
|
|
| variant = load_registry()[args.variant] |
| print(json.dumps({"variant": variant.name, **weight_hashes(load_pipeline(variant))}, indent=2)) |
| return 0 |
|
|
|
|
| def _cmd_generate(args: argparse.Namespace) -> int: |
| """Generate for a JSONL file of {prompt_id, prompt, seed}; save images and a records file.""" |
| from uvw.generate import GenItem, generate, load_inference_config |
| from uvw.models import load_pipeline, load_registry, weight_hashes |
| from uvw.provenance import RunRecord, config_hash |
|
|
| cfg = load_inference_config() |
| if args.steps: |
| cfg["num_inference_steps"] = args.steps |
| variant = load_registry()[args.variant] |
| rows = [json.loads(line) for line in Path(args.prompts).read_text().splitlines() if line.strip()] |
| if args.limit: |
| rows = rows[: args.limit] |
|
|
| record = RunRecord(job="generate", variants=[variant.name], configs={"inference": config_hash(cfg)}) |
| out = Path(args.out or RUNS / record.run_id) |
| (out / "images").mkdir(parents=True, exist_ok=True) |
| pipe = load_pipeline(variant) |
| record.weights = {f"{variant.name}/{k}": v for k, v in weight_hashes(pipe).items()} |
|
|
| items = [GenItem(r["prompt_id"], r.get("text") or r["prompt"], int(r["seed"])) for r in rows] |
| with open(out / "generations.jsonl", "w") as f: |
| for g in generate(pipe, items, cfg, record.cost): |
| name = f"{g.item.prompt_id}_s{g.item.seed}.png" |
| if not args.no_images: |
| g.image.save(out / "images" / name) |
| f.write(json.dumps({**asdict(g.item), "image": name, "image_sha256": g.image_sha256, |
| "config_hash": g.config_hash, "variant": variant.name}) + "\n") |
| record.outputs = [str(out / "generations.jsonl")] |
| record.finish() |
| print(record.write(out)) |
| return 0 |
|
|
|
|
| def _cmd_splits(args: argparse.Namespace) -> int: |
| """Build sources → template-level splits → leakage report → data/splits/.""" |
| from uvw import ROOT |
| from uvw.data import leakage |
| from uvw.data.sources import build_all |
| from uvw.data.splits import assign, write_splits |
|
|
| out = Path(args.out or ROOT / "data" / "splits") |
| out.mkdir(parents=True, exist_ok=True) |
| df, excluded = build_all() |
| emb = leakage.clip_text_embeddings(df.text.tolist()) |
| df["group_id"] = leakage.group_templates(df, emb, threshold=args.threshold).values |
| df["_row"] = range(len(df)) |
| df = assign(df) |
| emb, df = emb[df._row.to_numpy()], df.drop(columns="_row") |
| overlap = leakage.template_overlap(df) |
| exact = leakage.exact_duplicates(df) |
| near = leakage.near_duplicates(df, emb, threshold=args.threshold) |
| near.to_csv(out / "near_duplicates.csv", index=False) |
| report = {"rows": len(df), "excluded": len(excluded), "template_overlap": len(overlap), |
| "exact_cross_split_duplicates": len(exact), "near_cross_split_duplicates": len(near), |
| "threshold": args.threshold} |
| if len(exact) or len(near): |
| raise SystemExit(f"Cross-split duplicates remain: {report}") |
| if len(overlap): |
| print(overlap.head()) |
| raise SystemExit("Template overlap across splits: refusing to write.") |
| manifest = write_splits(df, out) |
| (out / "excluded.jsonl").write_text("".join(json.dumps(e) + "\n" for e in excluded)) |
| exact.to_csv(out / "exact_duplicates.csv", index=False) |
| (out / "leakage_report.json").write_text(json.dumps(report, indent=2)) |
| print(json.dumps(report, indent=2)) |
| for name, info in sorted(manifest["files"].items()): |
| print(f"{name:28s} rows={info['rows']:6d} templates={info['templates']:5d} {info['by_role']}") |
| return 0 |
|
|
|
|
| def _cmd_evaluate(args: argparse.Namespace) -> int: |
| from uvw.evaluate import evaluate |
|
|
| for concept in args.concepts: |
| out = evaluate(args.variant, concept, args.split, n_seeds=args.seeds, limit=args.limit, |
| out=Path(args.out) / f"{args.variant}-{concept}-{args.split}" if args.out else None, |
| keep_images=args.keep_images, edits=Path(args.edits) if args.edits else None) |
| print(out) |
| return 0 |
|
|
|
|
| def _cmd_edit_agent(args: argparse.Namespace) -> int: |
| """Run the editor agent locally (ChatGPT plan via atbot) against a remote tool server.""" |
| from uvw import RUNS |
| from uvw.harness.agent import run_editor |
| from uvw.harness.server import RemoteSession |
|
|
| session = RemoteSession(args.url, RUNS / "edits-local" / args.run_id) |
| summary = run_editor(session, editor_model=args.editor, model_name=args.model_name) |
| print(json.dumps(summary, indent=2)) |
| return 0 |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(prog="uvw") |
| sub = parser.add_subparsers(dest="command", required=True) |
| sub.add_parser("env", help="Print the environment record") |
| p = sub.add_parser("hash", help="Hash a variant's weights") |
| p.add_argument("variant") |
| p = sub.add_parser("generate", help="Generate images with the frozen inference config") |
| p.add_argument("variant") |
| p.add_argument("prompts", help="JSONL with prompt_id, text (or prompt), seed") |
| p.add_argument("--out") |
| p.add_argument("--limit", type=int) |
| p.add_argument("--steps", type=int, help="Override steps (smoke tests only)") |
| p.add_argument("--no-images", action="store_true", help="Keep hashes and scores only") |
| p = sub.add_parser("splits", help="Build frozen dev/val/test prompt splits") |
| p.add_argument("--out") |
| p.add_argument("--threshold", type=float, default=0.95) |
| p = sub.add_parser("evaluate", help="Generate + score a variant on concept splits") |
| p.add_argument("variant") |
| p.add_argument("concepts", nargs="+", choices=["nudity", "vangogh", "church", "coco"]) |
| p.add_argument("--split", default="dev", choices=["dev", "val"]) |
| p.add_argument("--seeds", type=int, default=1) |
| p.add_argument("--limit", type=int, help="Max prompts per role (pilots)") |
| p.add_argument("--out") |
| p.add_argument("--keep-images", default="none", choices=["none", "benign", "all"]) |
| p.add_argument("--edits", help="edited_tensors.safetensors from a committed candidate") |
| p = sub.add_parser("edit-agent", help="Run the editor agent against a tool server") |
| p.add_argument("--url", default="http://127.0.0.1:8765") |
| p.add_argument("--run-id", required=True) |
| p.add_argument("--editor", default="gpt-6-astra") |
| p.add_argument("--model-name", default="Stable Diffusion v1.4") |
| args = parser.parse_args(argv) |
| commands = {"edit-agent": _cmd_edit_agent, "evaluate": _cmd_evaluate, "env": _cmd_env, "hash": _cmd_hash, "generate": _cmd_generate, "splits": _cmd_splits} |
| return commands[args.command](args) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
