""" Run VisEdit (VEAD) baseline for hallucination suppression. Uses the exact same data loading as DualEdit (run_dualedit.sh): build_requests(edit_set, use_eval_instances=True) → edit_set["eval_instances"]["bathroom_no_toilet"] (the fixed ~50 images) Only differences from DualEdit: prompts match run_dualedit.py and the requests are converted to VisEdit EIC JSON format for vead_train.py. NOTE: VisEdit loads TWO copies of LLaVA-1.5-7b simultaneously (one for training, one for data preprocessing). This requires ~32 GB VRAM. Use --proc_device to place the preprocessing model on a second GPU. Usage: python -m experiment.knowledge_editing.run_visedit \\ --edit_set experiment/data/edit_set.json \\ --output_dir step4_baseline_outputs/visedit \\ --device cuda:0 \\ --proc_device cuda:1 # Skip training (use existing checkpoint) python -m experiment.knowledge_editing.run_visedit \\ --edit_set experiment/data/edit_set.json \\ --output_dir step4_baseline_outputs/visedit \\ --skip_train """ import argparse import json import os import sys import subprocess from pathlib import Path # --------------------------------------------------------------------------- # Prompts — match run_dualedit.py exactly # --------------------------------------------------------------------------- EDIT_PROMPT = "Describe this image." REPHRASE_PROMPT = "What do you see in this image?" # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- VISEDIT_DIR = Path(__file__).resolve().parents[2] / "VisEdit" def requests_to_visedit_json(requests: list[dict], img_dir: Path) -> list[dict]: """Convert request list to VisEdit EIC JSON format. Saves PIL images to img_dir and records filenames. Fields: image, src, alt, rephrase, image_rephrase, loc, loc_ans, m_loc, m_loc_q, m_loc_a """ img_dir.mkdir(parents=True, exist_ok=True) def save_img(pil_img, filename: str) -> str: dst = img_dir / filename if not dst.exists(): pil_img.save(dst, format="JPEG", quality=95) return filename records = [] for req in requests: image_id = req["_image_id"] fname = save_img(req["image"], f"edit_{image_id}.jpg") reph_img = req.get("image_rephrase") reph_fname = save_img(reph_img, f"reph_{image_id}.jpg") if reph_img else fname loc_img = req.get("multimodal_locality_image") loc_fname = save_img(loc_img, f"loc_{image_id}.jpg") if loc_img else fname records.append({ "image": fname, "src": req["prompt"], "alt": req["target"], "rephrase": req["rephrase_prompt"], "image_rephrase": reph_fname, "loc": req["locality_prompt"], "loc_ans": req["locality_ground_truth"], "m_loc": loc_fname, "m_loc_q": req["multimodal_locality_prompt"], "m_loc_a": req["multimodal_locality_ground_truth"], }) return records def patch_global_py(model_name: str = "llava-hf/llava-1.5-7b-hf"): """Overwrite VisEdit/utils/GLOBAL.py with correct absolute paths.""" global_py = VISEDIT_DIR / "utils" / "GLOBAL.py" root = str(VISEDIT_DIR) content = ( f"ROOT_PATH = {root!r}\n" f"model_path_map = {{\n" f" 'llava-v1.5-7b': {model_name!r},\n" f" 'blip2-opt-2.7b': 'models/blip2-opt-2.7b',\n" f" 'minigpt-4-vicuna-7b': 'models/minigpt-4-vicuna-7b',\n" f"}}\n" ) global_py.write_text(content) print(f" Patched {global_py}") def find_latest_checkpoint(records_dir: Path) -> str | None: """Find the most recently modified checkpoint in records/vead/llava-v1.5-7b/.""" base = records_dir / "vead" / "llava-v1.5-7b" if not base.exists(): return None ckpts = sorted(base.rglob("epoch-*"), key=lambda p: p.stat().st_mtime, reverse=True) return str(ckpts[0]) if ckpts else None def run_training(device: str, proc_device: str, epochs: int, batch_size: int, save_per: int): """Run vead_train.py as subprocess from VISEDIT_DIR.""" proc_idx = proc_device.split(":")[-1] if ":" in proc_device else "0" cmd = [ sys.executable, "vead_train.py", "-mn", "llava", "-dna", "EIC", "-bs", str(batch_size), "-dvc", device, "-edvc", proc_idx, "-eps", str(epochs), "-sci", str(save_per), "-tnp", "bathroom-toilet", ] print(f"\n Running: {' '.join(cmd)}") env = os.environ.copy() env["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" result = subprocess.run(cmd, cwd=str(VISEDIT_DIR), env=env) if result.returncode != 0: raise RuntimeError(f"vead_train.py exited with code {result.returncode}") # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="Prepare data and train VisEdit (VEAD) baseline" ) parser.add_argument("--edit_set", type=str, default="experiment/data/edit_set.json", help="Path to edit_set.json — same file DualEdit uses") parser.add_argument("--output_dir", type=str, default="step4_baseline_outputs/visedit") parser.add_argument("--dataset_id", type=str, default="pbcong/bathroom-toilet", help="HF dataset ID for image loading fallback") parser.add_argument("--model_name", type=str, default="llava-hf/llava-1.5-7b-hf", help="HF model ID or local path for llava-v1.5-7b") parser.add_argument("--device", type=str, default="cuda:0", help="CUDA device for the main (training) model") parser.add_argument("--proc_device", type=str, default="cuda:1", help="CUDA device for data-preprocessing copy of model") parser.add_argument("--epochs", type=int, default=1000) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--save_per", type=int, default=500, help="Save checkpoint every N iterations") parser.add_argument("--skip_train", action="store_true", help="Skip training; look for existing checkpoint") args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) # ----------------------------------------------------------------------- # 1. Load data — identical to DualEdit (run_baselines.py use_eval_instances=True) # ----------------------------------------------------------------------- from experiment.knowledge_editing.run_baselines import build_requests print(f"Loading edit set from {args.edit_set}...") with open(args.edit_set) as f: edit_set = json.load(f) print(f" Stats: {edit_set['stats']}") print("\nBuilding requests (eval_instances.bathroom_no_toilet, same as DualEdit)...") requests = build_requests(edit_set, dataset_id=args.dataset_id, use_eval_instances=True) # Override prompts to match run_dualedit.py for req in requests: req["prompt"] = EDIT_PROMPT req["rephrase_prompt"] = REPHRASE_PROMPT req["multimodal_locality_prompt"] = EDIT_PROMPT print(f" {len(requests)} requests ready") # ----------------------------------------------------------------------- # 2. Convert to VisEdit EIC JSON format # ----------------------------------------------------------------------- visedit_img_dir = VISEDIT_DIR / "data" / "easy-edit-mm" / "images" visedit_cap_dir = VISEDIT_DIR / "data" / "easy-edit-mm" / "caption" visedit_cap_dir.mkdir(parents=True, exist_ok=True) print("\nConverting to VisEdit EIC format...") records = requests_to_visedit_json(requests, visedit_img_dir) train_json_path = visedit_cap_dir / "caption_train_edit.json" with open(train_json_path, "w") as f: json.dump(records, f, indent=2) print(f" Saved {len(records)} records → {train_json_path}") # VisEdit also expects a caption_eval_edit.json (same data for EIC mode) eval_json_path = visedit_cap_dir / "caption_eval_edit.json" with open(eval_json_path, "w") as f: json.dump(records, f, indent=2) print(f" Saved {len(records)} records → {eval_json_path}") # ----------------------------------------------------------------------- # 3. Patch GLOBAL.py # ----------------------------------------------------------------------- print("\nPatching VisEdit GLOBAL.py...") patch_global_py(args.model_name) # ----------------------------------------------------------------------- # 4. Training # ----------------------------------------------------------------------- if not args.skip_train: print("\n>>> Running VEAD training...") run_training( device=args.device, proc_device=args.proc_device, epochs=args.epochs, batch_size=args.batch_size, save_per=args.save_per, ) else: print("\n>>> Skipping training (--skip_train)") # ----------------------------------------------------------------------- # 5. Find checkpoint # ----------------------------------------------------------------------- records_dir = VISEDIT_DIR / "records" ckpt = find_latest_checkpoint(records_dir) if ckpt is None: print("WARNING: No checkpoint found. Run training first.") else: print(f"\n Found checkpoint: {ckpt}") # ----------------------------------------------------------------------- # 6. Save eval_targets.json and edit_image_ids.json # ----------------------------------------------------------------------- eval_targets = {req["_image_id"]: req["target"] for req in requests} eval_targets_path = os.path.join(args.output_dir, "eval_targets.json") with open(eval_targets_path, "w") as f: json.dump(eval_targets, f, indent=2) print(f" Saved {len(eval_targets)} eval targets → {eval_targets_path}") edit_image_ids = [req["_image_id"] for req in requests] edit_image_ids_path = os.path.join(args.output_dir, "edit_image_ids.json") with open(edit_image_ids_path, "w") as f: json.dump(edit_image_ids, f, indent=2) print(f" Saved {len(edit_image_ids)} edit image IDs → {edit_image_ids_path}") # ----------------------------------------------------------------------- # 7. Save run config # ----------------------------------------------------------------------- run_config = { "visedit_dir": str(VISEDIT_DIR), "checkpoint": ckpt, "model_name": args.model_name, "device": args.device, "n_train": len(records), "edit_set": args.edit_set, "dataset_id": args.dataset_id, "eval_targets": eval_targets_path, "edit_image_ids": edit_image_ids_path, } config_path = os.path.join(args.output_dir, "run_config.json") with open(config_path, "w") as f: json.dump(run_config, f, indent=2) print(f"Run config saved → {config_path}") if __name__ == "__main__": main()