#!/usr/bin/env bash
# C2 baseline: official UCE (rohitgandikota/unified-concept-editing @ pinned commit, MIT), unmodified,
# then converted to our candidate format (edited_tensors.safetensors with `unet.` prefixes).
# usage: scripts/run_uce.sh <variant: tiny|sd14> <hf model id>
set -euo pipefail
VARIANT=$1; MODEL_ID=$2
COMMIT=7c724d9d7d19190e9bd33fe62fb656f1df87cd8e
mkdir -p third_party/uce && cd third_party/uce
[ -f uce_sd_erase.py ] || curl -sfLO "https://raw.githubusercontent.com/rohitgandikota/unified-concept-editing/$COMMIT/trainscripts/uce_sd_erase.py"
cd ../..
NEIGHBOURS="Claude Monet;Pablo Picasso;Rembrandt;Andy Warhol;Caravaggio;Paul Cezanne;Paul Gauguin"
OUT=runs/edits/c2-uce-$VARIANT-vangogh
python third_party/uce/uce_sd_erase.py --model_id "$MODEL_ID" --edit_concepts "Van Gogh" --concept_type art \
  --save_dir $OUT/raw --exp_name uce-default
python third_party/uce/uce_sd_erase.py --model_id "$MODEL_ID" --edit_concepts "Van Gogh" --concept_type art \
  --preserve_concepts "$NEIGHBOURS" --save_dir $OUT/raw --exp_name uce-preserve
for SCALE in 2 4 8; do
  python third_party/uce/uce_sd_erase.py --model_id "$MODEL_ID" --edit_concepts "Van Gogh" --concept_type art \
    --preserve_concepts "$NEIGHBOURS" --erase_scale $SCALE --save_dir $OUT/raw --exp_name uce-preserve-x$SCALE
done
python - "$OUT" <<'PY'
import json, sys
from pathlib import Path
from safetensors.torch import load_file, save_file
from uvw.provenance import tensor_digest
out = Path(sys.argv[1])
for f in sorted((out / "raw").glob("*.safetensors")):
    t = {f"unet.{k}": v.contiguous() for k, v in load_file(str(f)).items()}
    d = out / "candidates" / f.stem
    d.mkdir(parents=True, exist_ok=True)
    save_file(t, str(d / "edited_tensors.safetensors"))
    (d / "meta.json").write_text(json.dumps({"name": f.stem, "condition": "C2", "method": "UCE (official, unmodified)",
                                             "commit": "7c724d9d7d19190e9bd33fe62fb656f1df87cd8e",
                                             "tensors": {k: tensor_digest(v) for k, v in t.items()}}))
    print(f.stem, len(t), "tensors")
PY