Anonymous-123/ImageNet-Editing

1"""2Train a diffusion model on images.3"""4 5import argparse6 7from guided_diffusion import dist_util, logger8from guided_diffusion.image_datasets import load_data9from guided_diffusion.resample import create_named_schedule_sampler10from guided_diffusion.script_util import (11    model_and_diffusion_defaults,12    create_model_and_diffusion,13    args_to_dict,14    add_dict_to_argparser,15)16from guided_diffusion.train_util import TrainLoop17 18 19def main():20    args = create_argparser().parse_args()21 22    dist_util.setup_dist()23    logger.configure()24 25    logger.log("creating model and diffusion...")26    model, diffusion = create_model_and_diffusion(27        **args_to_dict(args, model_and_diffusion_defaults().keys())28    )29    model.to(dist_util.dev())30    schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion)31 32    logger.log("creating data loader...")33    data = load_data(34        data_dir=args.data_dir,35        batch_size=args.batch_size,36        image_size=args.image_size,37        class_cond=args.class_cond,38    )39 40    logger.log("training...")41    TrainLoop(42        model=model,43        diffusion=diffusion,44        data=data,45        batch_size=args.batch_size,46        microbatch=args.microbatch,47        lr=args.lr,48        ema_rate=args.ema_rate,49        log_interval=args.log_interval,50        save_interval=args.save_interval,51        resume_checkpoint=args.resume_checkpoint,52        use_fp16=args.use_fp16,53        fp16_scale_growth=args.fp16_scale_growth,54        schedule_sampler=schedule_sampler,55        weight_decay=args.weight_decay,56        lr_anneal_steps=args.lr_anneal_steps,57    ).run_loop()58 59 60def create_argparser():61    defaults = dict(62        data_dir="",63        schedule_sampler="uniform",64        lr=1e-4,65        weight_decay=0.0,66        lr_anneal_steps=0,67        batch_size=1,68        microbatch=-1,  # -1 disables microbatches69        ema_rate="0.9999",  # comma-separated list of EMA values70        log_interval=10,71        save_interval=10000,72        resume_checkpoint="",73        use_fp16=False,74        fp16_scale_growth=1e-3,75    )76    defaults.update(model_and_diffusion_defaults())77    parser = argparse.ArgumentParser()78    add_dict_to_argparser(parser, defaults)79    return parser80 81 82if __name__ == "__main__":83    main()84