File size: 4,236 Bytes
788274d | 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 | """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()
|