| """Paired example images for the public dataset (non-sexual concepts only). |
| |
| Predeclared selection, not cherry-picked: the first 12 target and 12 labelled-neighbour test prompts |
| by prompt_id, and the first 12 indirect probes, each with its first frozen seed, rendered by every |
| model listed below under the frozen inference config. Output is HF imagefolder format: |
| <out>/<model>/<prompt_id>.png plus metadata.jsonl with prompt, seed, checkpoint and detector scores. |
| |
| python scripts/make_examples.py --variant tiny --concept vangogh --out runs/examples/tiny-vangogh |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import pandas as pd |
| from huggingface_hub import snapshot_download |
|
|
| from uvw import ROOT |
| from uvw.generate import GenItem, generate, load_inference_config |
| from uvw.models import apply_edited_tensors, load_pipeline, load_registry |
| from uvw.score import load_scoring_config, scorer_for |
|
|
| MODELS = { |
| "original": None, |
| "C2-UCE": "tiny/vangogh/C2/c2-uce-tiny-vangogh/uce-preserve-x2", |
| "C3-ESD": "tiny/vangogh/C3/c3-esd-tiny-vangogh/esd-esd-x", |
| "C5-r2": "tiny/vangogh/C5/c5-tiny-vangogh-r2/vg_value_remap_bounded", |
| "C6-r1": "tiny/vangogh/C6/c6-tiny-vangogh-r1/vg_value_remap_strong", |
| "C6-r2": "tiny/vangogh/C6/c6-tiny-vangogh-r2/vg_preserve_regularized", |
| "C6-r3": "tiny/vangogh/C6/c6-tiny-vangogh-r3/vg_value_rewrite_r10", |
| } |
| NEIGHBOUR_ARTIST_CATEGORIES = ("subject-", "style-") |
|
|
|
|
| def select_prompts(heldout: Path, concept: str, n: int) -> pd.DataFrame: |
| test = pd.read_json(heldout / "test" / f"{concept}.jsonl", lines=True).sort_values("prompt_id") |
| probes = pd.read_json(heldout / "probes" / f"{concept}_indirect.jsonl", lines=True).sort_values("prompt_id") |
| targets = test[test.role == "target"].head(n).assign(set="test-target") |
| nb = test[(test.role == "neighbour") & test.category.str.startswith(NEIGHBOUR_ARTIST_CATEGORIES)].head(n) |
| return pd.concat([targets, nb.assign(set="test-neighbour"), probes.head(n).assign(set="indirect-probe")]) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--variant", required=True) |
| p.add_argument("--concept", required=True) |
| p.add_argument("--n", type=int, default=12) |
| p.add_argument("--out", required=True) |
| a = p.parse_args() |
| if a.concept == "nudity": |
| raise SystemExit("no example images are produced for nudity") |
|
|
| heldout = Path(snapshot_download("javadtaghia/unlearning-via-weight-heldout", repo_type="dataset", |
| local_dir=ROOT / "heldout")) |
| ckroot = Path(snapshot_download("javadtaghia/unlearning-via-weight", |
| allow_patterns=[f"checkpoints/{c}/*" for c in MODELS.values() if c], |
| local_dir=ROOT / "hub-checkpoints")) / "checkpoints" |
| prompts = select_prompts(heldout, a.concept, a.n) |
| cfg, scorer = load_inference_config(), scorer_for(a.concept, load_scoring_config()) |
| out = Path(a.out) |
| rows = [] |
| for model, ckpt in MODELS.items(): |
| pipe = load_pipeline(load_registry()[a.variant]) |
| if ckpt: |
| apply_edited_tensors(pipe, ckroot / ckpt / "edited_tensors.safetensors") |
| items = [GenItem(r.prompt_id, r.text, int(r.seed)) for r in prompts.itertuples()] |
| images = [g.image for g in generate(pipe, items, cfg)] |
| scores = scorer.score(images) |
| (out / model).mkdir(parents=True, exist_ok=True) |
| for r, img, s in zip(prompts.itertuples(), images, scores): |
| name = f"{model}/{r.prompt_id}.png" |
| img.save(out / name) |
| rows.append({"file_name": name, "model": model, "checkpoint": ckpt or "original", "set": r.set, |
| "prompt_id": r.prompt_id, "category": r.category, "prompt": r.text, "seed": int(r.seed), |
| "style_top1_clip_l14": s.get("clip_l14_top1"), "style_top1_clip_h14": s.get("clip_h14_top1"), |
| "van_gogh_detected": bool(s.get("target_hit"))}) |
| print("rendered", model, flush=True) |
| with open(out / "metadata.jsonl", "w") as f: |
| for row in rows: |
| f.write(json.dumps(row) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
