train_vae_w_nf.py

# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http:#www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import logging
import os
import json
import copy
from collections import OrderedDict

from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
import torch
import torch.backends.cudnn as cudnn
import numpy as np
import torch.distributed as dist
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.datasets as datasets
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from tqdm.auto import tqdm
from omegaconf import OmegaConf
import wandb

from util.crop import center_crop_arr
from loss.losses import ReconstructionLoss_Single_Stage
from models.vae import AutoencoderKL, ViTAsVAE, VAE_MODEL_DICT
from models import simflow
import util.misc as misc
from util.lr_sched import adjust_learning_rate
from util.encoder import load_encoders, preprocess_raw_image
from engine import update_ema, evaluate


def img2save(img):
    return (img * 0.5 + 0.5).clamp(0, 1)


def requires_grad(model, flag=True):
    """
    Set requires_grad flag for all parameters in a model.
    """
    for p in model.parameters():
        p.requires_grad = flag


def gather(tensor):
    """
    Gather tensors from all workers.
    """
    tensor = tensor.clone()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    tensor /= dist.get_world_size()

    return tensor.detach().item()

#################################################################################
#                                  Training Loop                                #
#################################################################################

def main(args):    
    misc.init_distributed_mode(args)

    device = torch.device(args.device)

    # fix the seed for reproducibility
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    cudnn.benchmark = True

    num_tasks = misc.get_world_size()
    global_rank = misc.get_rank()

    if global_rank == 0:
        print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
        print("{}".format(args).replace(', ', ',\n'))
        os.makedirs(args.log_dir, exist_ok=True)
        layers_per_block = misc.arg_to_str(args.layers_per_block)
        align_depth = misc.arg_to_str(args.repa_align_depth)
        wandb.init(
            project="SimFlow",
            name=f"{args.vae}_std{args.fixed_std}_w_NF_c{args.channels}_b{args.blocks}_l{layers_per_block}_bs{args.batch_size}_REPAenc_{args.enc_type}_align{align_depth}_lr{args.vae_learning_rate}_nf_coef{args.nf_loss_coef}_ln",
            config=args,
            dir=args.log_dir,
        )

    if args.enc_type != 'None':
        encoders, encoder_types, architectures = load_encoders(
            args.enc_type, device, args.resolution
        )
    external_feat_dims = [encoder.embed_dim for encoder in encoders] if args.enc_type != 'None' else [0]

    # Load model and create an EMA of the model
    vae_config = VAE_MODEL_DICT[args.vae]
    vae_embed_dim = vae_config["embed_dim"]
    vae_patch_size = vae_config["vae_patch_size"]
    ch_mult = vae_config["ch_mult"]
    img_size = args.resolution // vae_patch_size

    model = simflow.SimFlow(
            in_channels=vae_embed_dim,
            img_size=img_size,
            patch_size=1,
            channels=args.channels,
            num_blocks=args.blocks,
            layers_per_block=args.layers_per_block,
            num_heads=args.num_heads,
            nvp=args.nvp,
            num_classes=args.class_num,
            label_drop_prob=args.label_drop_prob,
            external_feat_dims=external_feat_dims,
            repa_align_depth=args.repa_align_depth,
            repa_loss_coef=args.repa_loss_coef,
        )
    if global_rank == 0:
        print("Model = %s" % str(model))
        # following timm: set wd as 0 for bias and norm layers
        n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print("Number of trainable parameters: {}M".format(n_params / 1e6))
    model.to(device)
    model_without_ddp = model

    # Load VAE and create an EMA of the VAE
    if args.vae.startswith('vae'):
        vae = AutoencoderKL(embed_dim=vae_embed_dim, ch_mult=ch_mult, 
                            use_variational=args.use_variational, fixed_std=args.fixed_std,
                            ln=args.vae_ln).to(device)
    elif args.vae.startswith('vit'):
        vae = ViTAsVAE(embed_dim=vae_embed_dim, img_size=args.resolution, 
        patch_size=vae_patch_size, use_variational=args.use_variational, fixed_std=args.fixed_std).to(device)
    else:
        raise ValueError(f"Invalid VAE type: {args.vae}")
    if args.vae_path is not None:
        vae.init_from_ckpt(args.vae_path)
    if args.vae_frozen_encoder:
        for param in vae.encoder.parameters():
            param.requires_grad = False
        for param in vae.quant_conv.parameters():
            param.requires_grad = False
        vae.encoder.eval()
    if global_rank == 0:
        print("VAE = %s" % str(vae))
        # following timm: set wd as 0 for bias and norm layers
        n_params = sum(p.numel() for p in vae.parameters() if p.requires_grad)
        print("Number of trainable parameters: {}M".format(n_params / 1e6))
    vae_without_ddp = vae

    loss_cfg = OmegaConf.load(args.loss_cfg_path)
    vae_loss_fn = ReconstructionLoss_Single_Stage(loss_cfg).to(device)

    # Define the optimizers for SiT, VAE, and VAE loss function separately
    optimizer_vae = torch.optim.AdamW(
        [
            {
                "params": model.parameters(),
                "lr": args.model_learning_rate,
            },
            {
                "params": [p for p in vae.parameters() if p.requires_grad],
                "lr": args.vae_learning_rate,
            },
        ],
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )
    optimizer_loss_fn = torch.optim.AdamW(
        vae_loss_fn.parameters(),
        lr=args.disc_learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    # Setup data
    # augmentation following DiT and ADM
    transform_train = transforms.Compose([
        transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.resolution)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])

    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
    sampler_train = torch.utils.data.DistributedSampler(
        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
    )
    if global_rank == 0:
        print(dataset_train)
        print("Sampler_train = %s" % str(sampler_train))

    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True,
    )

    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        model_without_ddp = model.module
        vae = torch.nn.parallel.DistributedDataParallel(vae, device_ids=[args.gpu])
        vae_without_ddp = vae.module
        vae_loss_fn = torch.nn.parallel.DistributedDataParallel(vae_loss_fn, device_ids=[args.gpu])
        vae_loss_fn_without_ddp = vae_loss_fn.module

    # resume training
    global_step = 0
    if args.resume and os.path.exists(os.path.join(args.resume, "checkpoint-last.pth")):
        if args.resume_step == -1:
            checkpoint = torch.load(os.path.join(args.resume, "checkpoint-last.pth"), map_location='cpu')
        else:
            checkpoint = torch.load(os.path.join(args.resume, f"checkpoint-{args.resume_step}.pth"), map_location='cpu')
        vae_without_ddp.load_state_dict(checkpoint['vae'])
        vae_params = list(vae_without_ddp.parameters())
        vae_ema_state_dict = checkpoint['vae_ema']
        vae_ema_params = [vae_ema_state_dict[name].cuda() for name, _ in vae_without_ddp.named_parameters()]

        miss_keys, unexpected_keys = model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
        print("miss_keys: ", miss_keys)
        print("unexpected_keys: ", unexpected_keys)
        model_params = list(model_without_ddp.parameters())
        ema_state_dict = checkpoint['model_ema']
        model_ema_params = [ema_state_dict[name].cuda() for name, _ in model_without_ddp.named_parameters()]

        if global_rank == 0:
            print("Resume checkpoint %s" % args.resume)

        vae_loss_fn_without_ddp.discriminator.load_state_dict(checkpoint['vae_disc'])

        optimizer_vae.load_state_dict(checkpoint['optimizer_vae'])
        optimizer_loss_fn.load_state_dict(checkpoint['optimizer_loss_fn'])
        global_step = checkpoint['steps']
        if global_rank == 0:
            print("With optim & sched!")
        del checkpoint
    else:
        vae_params = list(vae_without_ddp.parameters())
        vae_ema_params = copy.deepcopy(vae_params)
        model_params = list(model_without_ddp.parameters())
        model_ema_params = copy.deepcopy(model_params)
        if global_rank == 0:
            print("Training from scratch")

    # if args.disc_pretrained_ckpt is not None:
    #     # Load the discriminator from a pretrained checkpoint if provided
    #     disc_ckpt = torch.load(args.disc_pretrained_ckpt, map_location=device)
    #     vae_loss_fn.discriminator.load_state_dict(disc_ckpt)
    #     if global_rank == 0:
    #         print(f"Loaded discriminator from {args.disc_pretrained_ckpt}")

    progress_bar = tqdm(
        range(0, args.max_train_steps),
        initial=global_step,
        desc="Steps",
        # Only show the progress bar once on each machine.
        disable=not (global_rank == 0),
    )

    # main training loop
    for epoch in range(args.epochs):
        vae_lr = adjust_learning_rate(optimizer_vae, args.vae_learning_rate, epoch, args)
        disc_lr = adjust_learning_rate(optimizer_loss_fn, args.disc_learning_rate, epoch, args)

        for imgs, labels in data_loader_train:
            imgs = imgs.to(device)
            labels = labels.to(device)
            with torch.amp.autocast("cuda", dtype=torch.bfloat16):
                # 1). Forward pass: VAE
                vae_imgs = imgs * 2. - 1.
                posterior, z_ori, z, recon_image = vae(vae_imgs, disturb_latents=args.disturb_latents)
                z_ori_mean_log, z_ori_std_log, z_ori_min_log, z_ori_max_log = z_ori.mean(), z_ori.std(), z_ori.min(), z_ori.max()
                z_mean_log, z_std_log, z_min_log, z_max_log = z.mean(), z.std(), z.min(), z.max()

                # 2) Forwrad pass: NF
                # extract the encoder features
                external_features = []
                if args.enc_type != 'None':
                    with torch.no_grad():
                        for encoder, encoder_type, arch in zip(encoders, encoder_types, architectures):
                            raw_image_ = preprocess_raw_image(imgs, encoder_type)
                            ext = encoder.forward_features(raw_image_)
                            if 'mocov3' in encoder_type: ext = ext[:, 1:] 
                            if 'dinov2' in encoder_type: ext = ext['x_norm_patchtokens']
                            external_features.append(ext)

                loss_dict, track_dict = model(z, labels, external_features=external_features)
                nf_loss = sum(loss_dict.values())

                # 2). Backward pass: VAE, compute the VAE loss, backpropagate, and update the VAE; Then, compute the discriminator loss and update the discriminator
                extra_result_dict = {
                    "posterior": posterior,
                    "nf_loss": nf_loss,
                }
                vae_loss, vae_loss_dict = vae_loss_fn(vae_imgs, recon_image, extra_result_dict, global_step, "generator")
                vae_loss.backward()
                grad_norm_vae = torch.nn.utils.clip_grad_norm_(vae.parameters(), args.max_grad_norm)
                grad_norm_model = torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                optimizer_vae.step()
                optimizer_vae.zero_grad(set_to_none=True)
                torch.cuda.synchronize()

                # discriminator loss and update
                d_loss, d_loss_dict = vae_loss_fn(vae_imgs, recon_image, posterior, global_step, "discriminator")
                d_loss.backward()
                grad_norm_disc = torch.nn.utils.clip_grad_norm_(vae_loss_fn.parameters(), args.max_grad_norm)
                optimizer_loss_fn.step()
                optimizer_loss_fn.zero_grad(set_to_none=True)
                torch.cuda.synchronize()

                update_ema(vae_ema_params, vae_params, rate=args.ema_rate)
                update_ema(model_ema_params, model_params, rate=args.ema_rate)

            # Prepare the logs based on the current step
            loss_dict_reduced = {
                "nf_" + k: gather(v) for k, v in loss_dict.items()
            }
            track_dict_reduced = {
                "nf_" + k: gather(v) for k, v in track_dict.items()
            }
            logs = {
                "epoch": epoch,
                # VAE loss
                "vae_loss/vae_loss": gather(vae_loss),
                "vae_loss/reconstruction_loss": gather(vae_loss_dict["reconstruction_loss"]),
                "vae_loss/perceptual_loss": gather(vae_loss_dict["perceptual_loss"]),
                "vae_loss/ent_loss": gather(vae_loss_dict["ent_loss"]),
                "vae_loss/weighted_gan_loss": gather(vae_loss_dict["weighted_gan_loss"]),
                "vae_loss/discriminator_factor": gather(vae_loss_dict["discriminator_factor"]),
                "vae_loss/gan_loss": gather(vae_loss_dict["gan_loss"]),
                "vae_loss/d_weight": gather(vae_loss_dict["d_weight"]),
                # Statistics
                "stats/z_ori_mean": gather(z_ori_mean_log),
                "stats/z_ori_std": gather(z_ori_std_log),
                "stats/z_ori_min": gather(z_ori_min_log),
                "stats/z_ori_max": gather(z_ori_max_log),
                "stats/z_mean": gather(z_mean_log),
                "stats/z_std": gather(z_std_log),
                "stats/z_min": gather(z_min_log),
                "stats/z_max": gather(z_max_log),
                # NF loss
                "nf_loss/nf_loss": gather(nf_loss),
                # Gradient norm
                "grad_norm/grad_norm_vae": gather(grad_norm_vae),
                "grad_norm/grad_norm_disc": gather(grad_norm_disc),
                "grad_norm/grad_norm_model": gather(grad_norm_model),
                # Discriminator loss
                "d_loss/d_loss": gather(d_loss),
                "d_loss/logits_real": gather(d_loss_dict["logits_real"]),
                "d_loss/logits_fake": gather(d_loss_dict["logits_fake"]),
                "d_loss/lecam_loss": gather(d_loss_dict["lecam_loss"]),
                # Learning rate
                "lr/vae_lr": vae_lr,
                "lr/disc_lr": disc_lr,
            }
            logs.update(loss_dict_reduced)
            logs.update(track_dict_reduced)

            # Log the metrics to wandb
            if global_rank == 0:
                wandb.log(logs, step=global_step)

                if global_step % args.checkpointing_steps == 0 or global_step + 1 == args.max_train_steps:
                    vae_ema_state_dict = copy.deepcopy(vae_without_ddp.state_dict())
                    for i, (name, _value) in enumerate(vae_without_ddp.named_parameters()):
                        assert name in vae_ema_state_dict
                        vae_ema_state_dict[name] = vae_ema_params[i]
                    model_ema_state_dict = copy.deepcopy(model_without_ddp.state_dict())
                    for i, (name, _value) in enumerate(model_without_ddp.named_parameters()):
                        assert name in model_ema_state_dict
                        model_ema_state_dict[name] = model_ema_params[i]
                    checkpoint = {
                        "model": model_without_ddp.state_dict(),
                        "model_ema": model_ema_state_dict,
                        "vae_ema": vae_ema_state_dict,
                        "vae": vae_without_ddp.state_dict(),
                        "vae_disc": vae_loss_fn_without_ddp.discriminator.state_dict(),
                        "optimizer_vae": optimizer_vae.state_dict(),
                        "optimizer_loss_fn": optimizer_loss_fn.state_dict(),
                        "args": args,
                        "steps": global_step,
                    }
                    torch.save(checkpoint, os.path.join(args.output_dir, "checkpoint-last.pth"))
                    torch.save(checkpoint, os.path.join(args.output_dir, f"checkpoint-{global_step}.pth"))
                    print(f"Saved checkpoint to {args.output_dir}")

                if (global_step == 1 or (global_step % args.sampling_steps == 0 and global_step > 0)):
                    wandb.log({"Original images": wandb.Image(img2save(vae_imgs[:8]))}, step=global_step)
                    wandb.log({"Reconstructed": wandb.Image(img2save(recon_image[:8]))}, step=global_step)

                    labels_gen = torch.Tensor(torch.randint(0, args.class_num, (8,))).long().cuda()
                    vae.eval()
                    model.eval()
                    with torch.no_grad():
                        sampled_tokens = model_without_ddp.sample_tokens(bsz=8, num_iter=0, labels=labels_gen)
                        samples = vae_without_ddp.decode(sampled_tokens)
                    wandb.log({"Generated": wandb.Image(img2save(samples))}, step=global_step)
                    if args.vae_frozen_encoder:
                        vae.module.decoder.train()
                    else:
                        vae.train()
                    model.train()
            
            # Online evaluate
            if args.online_eval and global_step % args.eval_steps == 0:
                vae.eval()
                model.eval()
                fid, is_score = evaluate(model_without_ddp, vae_without_ddp, model_ema_params, args, 0, batch_size=args.eval_bsz, cfg=args.cfg, use_ema=True)
                if global_rank == 0:
                    wandb.log({"val/fid": fid, "val/is_score": is_score}, step=global_step)
                if args.vae_frozen_encoder:
                    vae.module.decoder.train()
                else:
                    vae.train()
                model.train()

            progress_bar.update(1)
            global_step += 1
            if global_step >= args.max_train_steps:
                break
            
        if global_step >= args.max_train_steps:
            break

    if global_rank == 0:
        print("Done!")
    wandb.finish()


def parse_args(input_args=None):
    parser = argparse.ArgumentParser(description="Training")
    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')

    # logging params
    parser.add_argument("--output_dir", type=str, default="exps")
    parser.add_argument("--log_dir", type=str, default="logs")
    parser.add_argument("--sampling-steps", type=int, default=1000)
    parser.add_argument("--continue-train-exp-dir", type=str, default=None)

    # dataset params
    parser.add_argument("--data_path", type=str, default="data")
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--class_num", type=int, default=1000)
    parser.add_argument("--resolution", type=int, default=256)

    # optimization params
    parser.add_argument("--epochs", type=int, default=1400)
    parser.add_argument("--max-train-steps", type=int, default=400000)
    parser.add_argument("--checkpointing-steps", type=int, default=50000)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
    parser.add_argument("--learning-rate", type=float, default=1e-4)
    parser.add_argument("--adam-beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam-beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument("--adam-weight-decay", type=float, default=0., help="Weight decay to use.")
    parser.add_argument("--adam-epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max-grad-norm", default=1.0, type=float, help="Max gradient norm.")
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--resume_step', type=int, default=-1, help='resume from specific step checkpoint')
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--lr_schedule', type=str, default='constant',
                        help='learning rate schedule')
    parser.add_argument('--warmup_epochs', type=int, default=10, metavar='N',
                        help='epochs to warmup LR')
    parser.add_argument('--hold_epochs', type=int, default=0, metavar='N',
                        help='epochs to hold LR')
    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                        help='lower lr bound for cyclic schedulers that hit 0')
    
    # Model
    parser.add_argument('--channels', default=512, type=int, help='Model width')
    parser.add_argument('--num_heads', default=8, type=int, help='Number of attention heads')
    parser.add_argument('--blocks', default=4, type=int, help='Number of autoregressive flow blocks')
    parser.add_argument('--layers_per_block', type=misc.parse_int_or_list, default="8", help='Depth per flow block')
    parser.add_argument('--label_drop_prob', default=0.1, type=float)
    parser.add_argument('--noise_std', default=0.05, type=float, help='Input noise standard deviation')
    parser.add_argument('--noise_type', default='gaussian', choices=['gaussian', 'uniform'], type=str)
    parser.add_argument('--nvp', default=True, action=argparse.BooleanOptionalAction, help='Whether to use the non volume preserving version')
    
    # REPA
    parser.add_argument("--enc_type", type=str, default='None')
    parser.add_argument("--repa_align_depth", type=misc.parse_int_or_list, default="-1,-1,-1,-1")

    # seed params
    parser.add_argument("--seed", type=int, default=0)

    # cpu params
    parser.add_argument("--num-workers", type=int, default=4)

    # vae params
    parser.add_argument("--vae", type=str, default="vae_f8d4")
    parser.add_argument("--vae_path", type=str, default=None)
    parser.add_argument('--use_variational', default="True", type=misc.str_to_bool, help='Whether to use the variational version of the VAE')
    parser.add_argument('--fixed_std', default=1.0, type=float, help='Fixed standard deviation for the VAE')
    parser.add_argument("--vae_ln", default="True", type=misc.str_to_bool, help='Whether to use layer norm in the VAE')
    parser.add_argument("--vae_frozen_encoder", default=False, action=argparse.BooleanOptionalAction, help='Whether to freeze the encoder of the VAE')

    # vae loss params
    parser.add_argument("--disc-pretrained-ckpt", type=str, default=None)
    parser.add_argument("--loss-cfg-path", type=str, default="configs/l1_lpips_kl_gan.yaml")

    # training params
    parser.add_argument("--model-learning-rate", type=float, default=1e-4)
    parser.add_argument("--vae-learning-rate", type=float, default=1e-4)
    parser.add_argument("--disc-learning-rate", type=float, default=1e-4)
    parser.add_argument("--nf_loss_coef", type=float, default=1.0)
    parser.add_argument("--repa_loss_coef", type=float, default=1.0)
    parser.add_argument('--ema_rate', default=0.9999, type=float)
    parser.add_argument('--disturb_latents', type=str, default='none')

    # eval params
    parser.add_argument('--num_iter', default=64, type=int,
                        help='number of autoregressive iterations to generate an image')
    parser.add_argument('--num_images', default=50000, type=int,
                        help='number of images to generate')
    parser.add_argument('--cfg', default=0.0, type=float, help="classifier-free guidance")
    parser.add_argument('--cfg_schedule', default="linear", type=str)
    parser.add_argument('--cfg_method', default="starflow", type=str)
    parser.add_argument('--temperature', default=1.0, type=float, help="temperature for sampling")
    parser.add_argument('--online_eval', action='store_true')
    parser.add_argument('--evaluate', action='store_true')
    parser.add_argument("--eval_steps", type=int, default=10000)
    parser.add_argument('--eval_bsz', type=int, default=256, help='generation batch size')
    parser.add_argument('--num_sampling_steps', type=int, default=50, help='number of sampling steps')
    parser.add_argument('--denoising_lr', default=0.0, type=float)

    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()
    os.makedirs(args.output_dir, exist_ok=True)
    args.log_dir = args.output_dir
    main(args)