run.py

#!/usr/bin/env python3
import os
import json
from pathlib import Path
from argparse import ArgumentParser
from typing import Dict, List, Optional
from network import ldm_stable, NullInversion, make_controller, text2image_ldm_stable_TR, NUM_DDIM_STEPS, GUIDANCE_SCALE
from auto_mask import generate_mask

def load_identity_map(path: Path) -> Dict[str, List[str]]:
    if not path.exists():
        raise SystemExit(f"--ref_map_json not found: {path}")
    with open(path, "r") as f:
        return json.load(f)

def pick_ref_from_identity(filename_stem: str, identity_map: Dict[str, List[str]], ref_base_dir: Path, ext: str) -> Path:
    if not ref_base_dir.exists():
        raise SystemExit(f"--ref_base_dir not found: {ref_base_dir}")
    ident = None
    for k, images in identity_map.items():
        if filename_stem in images:
            ident = k
            break
    if ident is None:
        raise RuntimeError(f"No identity entry for '{filename_stem}' in ref map.")
    for candidate in identity_map[ident]:
        if candidate != filename_stem:
            p = ref_base_dir / f"{candidate}{ext}"
            if p.exists():
                return p
    raise RuntimeError(f"No alternate ref image found for '{filename_stem}' under {ref_base_dir}.")

def run_single(img_path: Path, mask_path: Path, result_dir: Path, source_prompt: str, target_prompt: str, target_word_list: List[str], max_iteration: int, scale: float, max_adv_iter: int, outside_grad_factor: float, outside_update_factor: float, target_model: str, tau_edit: int, tau_adv: int, lambda_cosine: float, lambda_lpips: float, identity_map: Dict[str, List[str]], ref_base_dir: Path, img_ext: str):
    result_dir.mkdir(parents=True, exist_ok=True)
    filename = img_path.stem
    ref_img_path = pick_ref_from_identity(filename, identity_map, ref_base_dir, f".{img_ext}")
    null_inversion = NullInversion(ldm_stable)
    (image_gt, _image_enc), x_t, uncond_embeddings = null_inversion.invert(str(img_path), source_prompt, offsets=(0, 0, 0, 0), verbose=True)
    prompts = [source_prompt, target_prompt]
    cross_replace_steps = {"default_": 0.8}
    self_replace_steps = 0.5
    positive_word = [[str(mask_path) for _ in target_word_list], list(target_word_list)]
    controller = make_controller(prompts, False, cross_replace_steps, self_replace_steps, None, name=str(result_dir), word=positive_word, adj_mask=str(mask_path), word_ng=None)
    images, _ = text2image_ldm_stable_TR(
        result_dir=str(result_dir),
        filename=filename,
        img_gt=image_gt,
        ref_image_path=str(ref_img_path),
        model=ldm_stable,
        prompt=prompts,
        controller=controller,
        latent=x_t,
        num_inference_steps=NUM_DDIM_STEPS,
        guidance_scale=GUIDANCE_SCALE,
        generator=None,
        uncond_embeddings=uncond_embeddings,
        max_iter=max_iteration,
        scale=scale,
        target_model=target_model,
        adv_max_iter=max_adv_iter,
        outside_grad_factor=outside_grad_factor,
        outside_update_factor=outside_update_factor,
        tau_edit=tau_edit,
        tau_adv=tau_adv,
        lambda_cosine=lambda_cosine,
        lambda_lpips=lambda_lpips,
    )
    return images

def run_batch(images_dir: Path, result_dir: Path, source_prompt: str, target_prompt: str, target_word_list: List[str], max_iteration: int, scale: float, max_adv_iter: int, outside_grad_factor: float, outside_update_factor: float, target_model: str, tau_edit: int, tau_adv: int, lambda_cosine: float, lambda_lpips: float, img_ext: str, mask_suffix: str, identity_map: Dict[str, List[str]], ref_base_dir: Path, skip_existing: bool, device_id: Optional[str], top_k: int):
    if device_id is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = device_id
    result_dir.mkdir(parents=True, exist_ok=True)
    if not images_dir.exists():
        raise SystemExit(f"--images_dir not found: {images_dir}")
    if not ref_base_dir.exists():
        raise SystemExit(f"--ref_base_dir not found: {ref_base_dir}")
    img_paths = sorted([p for p in images_dir.glob(f"*.{img_ext}")])
    auto_mask_dir = result_dir / "_masks"
    auto_mask_dir.mkdir(parents=True, exist_ok=True)
    for img_path in img_paths:
        stem = img_path.stem
        out_file = result_dir / f"{stem}.png"
        if skip_existing and out_file.exists():
            print(f"[SKIP] {img_path.name} -> exists")
            continue
        mask_path = auto_mask_dir / f"{stem}{mask_suffix}"
        if not mask_path.exists():
            generate_mask(target_prompt, str(img_path), str(mask_path), top_k=top_k, device="cuda" if device_id is None or device_id != "" else "cpu")
        try:
            ref_img_path = pick_ref_from_identity(stem, identity_map, ref_base_dir, f".{img_ext}")
        except RuntimeError as e:
            print(f"[SKIP] {img_path.name} -> {e}")
            continue
        print(f"[RUN ] {img_path.name}  mask={mask_path.name}  ref={ref_img_path.name}")
        run_single(img_path=img_path, mask_path=mask_path, result_dir=result_dir, source_prompt=source_prompt, target_prompt=target_prompt, target_word_list=target_word_list, max_iteration=max_iteration, scale=scale, max_adv_iter=max_adv_iter, outside_grad_factor=outside_grad_factor, outside_update_factor=outside_update_factor, target_model=target_model, tau_edit=tau_edit, tau_adv=tau_adv, lambda_cosine=lambda_cosine, lambda_lpips=lambda_lpips, identity_map=identity_map, ref_base_dir=ref_base_dir, img_ext=img_ext)

def parse_args():
    p = ArgumentParser()
    p.add_argument("--mode", choices=["single", "batch"], default="single")
    p.add_argument("--source_prompt", required=True)
    p.add_argument("--target_prompt", required=True)
    p.add_argument("--target_word", nargs="+", required=True)
    p.add_argument("--result_dir", required=True)
    p.add_argument("--max_iteration", type=int, default=15)
    p.add_argument("--scale", type=float, default=2.5)
    p.add_argument("--max_adv_iter", type=int, default=20)
    p.add_argument("--outside_grad_factor", type=float, default=1.0)
    p.add_argument("--outside_update_factor", type=float, default=1.0)
    p.add_argument("--target_model", choices=["mobile_face", "facenet", "ir152", "irse50"], default="mobile_face")
    p.add_argument("--tau_edit", type=int, default=10)
    p.add_argument("--tau_adv", type=int, default=45)
    p.add_argument("--lambda_cosine", type=float, default=0.25)
    p.add_argument("--lambda_lpips", type=float, default=2.5)
    p.add_argument("--device_id", type=str)
    p.add_argument("--img_path")
    p.add_argument("--images_dir")
    p.add_argument("--img_ext", choices=["jpg", "png"], default="jpg")
    p.add_argument("--mask_suffix", default="_mask.png")
    p.add_argument("--ref_map_json", required=True)
    p.add_argument("--ref_base_dir", required=True)
    p.add_argument("--skip_existing", action="store_true")
    p.add_argument("--topk_regions", type=int, default=1)
    return p.parse_args()

def main():
    args = parse_args()
    ref_map_json = Path(args.ref_map_json)
    ref_base_dir = Path(args.ref_base_dir)
    identity_map = load_identity_map(ref_map_json)
    if args.mode == "single":
        if not args.img_path:
            raise SystemExit("--img_path is required for mode=single")
        result_dir = Path(args.result_dir)
        result_dir.mkdir(parents=True, exist_ok=True)
        auto_mask_dir = result_dir / "_masks"
        auto_mask_dir.mkdir(parents=True, exist_ok=True)
        img_path = Path(args.img_path)
        stem = img_path.stem
        mask_path = auto_mask_dir / f"{stem}{args.mask_suffix}"
        generate_mask(args.target_prompt, str(img_path), str(mask_path), top_k=args.topk_regions, device="cuda" if not args.device_id else "cuda")
        run_single(img_path=img_path, mask_path=mask_path, result_dir=result_dir, source_prompt=args.source_prompt, target_prompt=args.target_prompt, target_word_list=args.target_word, max_iteration=args.max_iteration, scale=args.scale, max_adv_iter=args.max_adv_iter, outside_grad_factor=args.outside_grad_factor, outside_update_factor=args.outside_update_factor, target_model=args.target_model, tau_edit=args.tau_edit, tau_adv=args.tau_adv, lambda_cosine=args.lambda_cosine, lambda_lpips=args.lambda_lpips, identity_map=identity_map, ref_base_dir=ref_base_dir, img_ext=args.img_ext)
    else:
        if not args.images_dir:
            raise SystemExit("--images_dir is required for mode=batch")
        run_batch(images_dir=Path(args.images_dir), result_dir=Path(args.result_dir), source_prompt=args.source_prompt, target_prompt=args.target_prompt, target_word_list=args.target_word, max_iteration=args.max_iteration, scale=args.scale, max_adv_iter=args.max_adv_iter, outside_grad_factor=args.outside_grad_factor, outside_update_factor=args.outside_update_factor, target_model=args.target_model, tau_edit=args.tau_edit, tau_adv=args.tau_adv, lambda_cosine=args.lambda_cosine, lambda_lpips=args.lambda_lpips, img_ext=args.img_ext, mask_suffix=args.mask_suffix, identity_map=identity_map, ref_base_dir=ref_base_dir, skip_existing=args.skip_existing, device_id=args.device_id, top_k=args.topk_regions)

if __name__ == "__main__":
    main()