pipeline.py

import yaml
import subprocess
from pathlib import Path
from argparse import ArgumentParser


def parse_args():
    parser = ArgumentParser()

    parser.add_argument(
        '--pretrained_model_or_path',
        default='stabilityai/stable-diffusion-2-base',
        help='hf model or path to pipeline',
    )
    parser.add_argument(
        '--base_dir',
        required=True,
        type=lambda arg: Path(arg).absolute().resolve(),
        help='path to directory with experiments',
    )
    parser.add_argument(
        '--prompts_dir',
        required=True,
        type=lambda arg: Path(arg).absolute().resolve(),
        help='path to directory with prompt files (named `{i}.json`)',
    )
    parser.add_argument(
        '--pairs_dir',
        required=True,
        type=lambda arg: Path(arg).absolute().resolve(),
        help='path to directory with pairs files (named `{i}`.json), with 0.json being present from base model'
    )

    parser.add_argument(
        '--concept',
        required=True,
        help='concept (e.g. dog6)',
    )
    parser.add_argument(
        '--concept_class',
        required=True,
        help='concept class (e.g. dog for dog6)',
    )

    parser.add_argument(
        '--exp_id',
        required=True,
        type=int,
        help='experiment ID (will be used in output path)',
    )
    parser.add_argument(
        '--exp_name',
        required=True,
        help='experiment name (will be used in output path)',
    )

    parser.add_argument(
        '--n_steps',
        required=True,
        type=int,
        help='number of DPO steps',
    )
    parser.add_argument(
        '--train_steps_per_dpo_step',
        required=True,
        type=int,
        help='total number of train steps'
    )
    parser.add_argument(
        '--checkpoint_steps',
        required=True,
        type=int,
        help='period for checkpointing & validation',
    )
    parser.add_argument(
        '--global_threshold',
        required=True,
        type=float,
        help='threshold for collect_pairs.py'
    )
    parser.add_argument(
        '--dry_run',
        action='store_true',
        help='do not execute, just print commands'
    )
    parser.add_argument(
        '--start_step',
        type=int,
        default=0,
        help='start dpo step (0-based)',
    )
    parser.add_argument(
        '--atan_min',
        type=float,
        help='lower bound for arctan in grad'
    )
    parser.add_argument(
        '--atan_max',
        type=float,
        help='upper bound for arctan in grad'
    )

    args = parser.parse_args()

    # Validation
    assert Path(args.prompts_dir).exists(), "prompts_dir is not found"
    assert args.n_steps > 0, "n_steps should be positive"
    assert len(list(Path(args.prompts_dir).iterdir())) == args.n_steps, "n_steps and data_files mismatch"
    for i in range(args.n_steps):
        assert (Path(args.prompts_dir) / f"{i}.json").exists(), f"`{i}.json` is not found in prompts_dir"
    assert Path(args.pairs_dir).exists(), "pairs_dir is not found"
    assert (Path(args.pairs_dir) / "0.json").exists(), "pairs_dir/0.json is not found"

    assert args.train_steps_per_dpo_step % args.checkpoint_steps == 0, "checkpoint_steps should divide train steps"

    print(f"Total of {args.n_steps} DPO steps")
    print(f"{args.train_steps_per_dpo_step} train steps per DPO step")

    return args


def get_full_exp_name(step: int, args) -> str:
    step_suffix = f"-step{step+1}of{args.n_steps}" if args.n_steps > 1 else ""
    return f"{args.exp_id:05d}-{step+1:04d}-{args.exp_name}{step_suffix}"


def get_workdir(step: int, args) -> Path:
    return args.base_dir / get_full_exp_name(step, args)


def get_all_checkpoints(workdir: Path):
    ckpt_dirs = list(workdir.glob("checkpoint-*"))
    return [int(ckpt.name.removeprefix("checkpoint-")) for ckpt in ckpt_dirs]


def train_ddpo(step: int, args):
    prev_workdir = get_workdir(step - 1, args) if step != 0 else args.pretrained_model_or_path
    workdir = get_workdir(step, args)

    cmd_args = [
        "accelerate", "launch", "train_ddpo.py",
        "--pretrained_model_name_or_path", f"{prev_workdir}",
        "--dataset_path", f"{args.pairs_dir}/{step}.json",
        "--output_dir", f"{workdir}",
        "--beta_dpo", "5000",
        "--max_train_steps", f"{args.train_steps_per_dpo_step}",
        "--train_batch_size", "8",
        "--gradient_accumulation_steps", "32",
        "--learning_rate", "1e-8",
        "--scale_lr",
        "--lr_scheduler", "constant_with_warmup",
        "--lr_warmup_steps", f"{int(0.2 * args.train_steps_per_dpo_step)}",
        "--checkpointing_steps", f"{args.checkpoint_steps}",
        "--dataloader_num_workers", "4",
        "--hard_skip_resume",
        "--report_to", "wandb",
        "--tracker_project_name", "persgen",
        "--validation_steps", f"{args.checkpoint_steps}",
        "--validation_concept", f"sks {args.concept_class}",
        "--validation_concept_image_dir", f"dreambooth/dataset/{args.concept}",
        "--val_batch_size", "64",
    ]

    print("Executing...")
    print('\n'.join(cmd_args))
    print(flush=True)

    if not args.dry_run:
        subprocess.run(cmd_args)


def generate_val(step: int, args):
    workdir = get_workdir(step, args)

    cmd_args = [
        "python", "generate.py",
        "--pretrained_model_or_path", f"{args.pretrained_model_or_path}",
        "--exp_dir", f"{workdir}",
        "--concept", f"sks {args.concept_class}",
        "--prompt_source", "eval_set",
        "--eval_set", "medium_live_all" if args.concept_class in ['cat', 'dog'] else "medium_object_all",
        "--samples_per_prompt", "10",
        "--prompts_per_batch", "8",
    ]

    print("Executing...")
    print('\n'.join(cmd_args))
    print(flush=True)

    if not args.dry_run:
        subprocess.run(cmd_args)


def generate_prompts(step: int, args):
    workdir = get_workdir(step, args)

    cmd_args = [
        "python", "generate.py",
        "--pretrained_model_or_path", f"{args.pretrained_model_or_path}",
        "--ckpt_path", f"{workdir}/checkpoint-{args.train_steps_per_dpo_step}",
        "--concept", f"sks {args.concept_class}",
        "--prompt_source", "json",
        "--prompts_json_path", f"{args.prompts_dir}/{step + 1}.json",
        "--samples_per_prompt", "10",
        "--prompts_per_batch", "8",
    ]

    print("Executing...")
    print('\n'.join(cmd_args))
    print(flush=True)

    if not args.dry_run:
        subprocess.run(cmd_args)


def log_hpararms(step: int, args):
    workdir = get_workdir(step, args)
    hparams = {
        "class_name": args.concept_class,
        "exp_name": workdir.name,
        "output_dir": str(workdir),
        "placeholder_token": "sks",
        "resolution": 512,
        "test_data_dir": f"dreambooth/dataset/{args.concept}",
    }

    if not args.dry_run:
        (workdir / "logs").mkdir(exist_ok=True)
        with open(workdir / "logs" / "hparams.yml", "w") as f:
            yaml.dump(hparams, f)


def evaluate_exp(step: int, args):
    workdir = get_workdir(step, args)
    ckpts = get_all_checkpoints(workdir)

    cmd_args = [
        "python", "evaluate_exp.py",
        "--base_path",  f"{args.base_dir}",
        "--exp_names", f"{workdir.name}",
        "--checkpoints_idxs"] + [str(c) for c in ckpts] + [
        "--cache_files_template", f"{args.base_dir}/*/eval.json",
    ]

    print("Executing...")
    print('\n'.join(cmd_args))
    print(flush=True)

    if not args.dry_run:
        subprocess.run(cmd_args)


def collect_pairs(step: int, args):
    workdir = get_workdir(step, args)
    ckpt = args.train_steps_per_dpo_step

    cmd_args = [
        "python",
        "collect_pairs.py",
        "--meta_path", f"{workdir}/eval.json",
        "--img_root", f"{workdir}/checkpoint-{ckpt}/samples/ns50_gs7.5/version_0",
        "--prompts_path", f"{args.prompts_dir}/{step + 1}.json",
        "--concept", f"sks {args.concept_class}",
        "--out_path", f"{args.pairs_dir}/{step + 1}.json",
        "--exp_key", f"('{workdir.name}', ('{ckpt}', '50', '7.5'))",
        "--score_type", "text",
        "--text_ratio", "0.75",
        "--threshold", f"{args.global_threshold}",]

    if args.atan_min is not None:
        cmd_args += ["--atan_min", str(args.atan_min)]
    if args.atan_max is not None:
        cmd_args += ["--atan_max", str(args.atan_max)]

    print("Executing...")
    print('\n'.join(cmd_args))
    print(flush=True)

    if not args.dry_run:
        subprocess.run(cmd_args)


def pipeline(args):
    for i_step in range(args.start_step, args.n_steps):
        print(f"Step {i_step + 1} / {args.n_steps}")
        is_last = (i_step + 1 == args.n_steps)

        train_ddpo(i_step, args)
        log_hpararms(i_step, args)
        generate_val(i_step, args)
        if not is_last:
            generate_prompts(i_step, args)
        evaluate_exp(i_step, args)
        if not is_last:
            collect_pairs(i_step, args)

        print()
        print("-" * 50)
        print()

    print("Done")


if __name__ == "__main__":
    args = parse_args()
    pipeline(args)