eval_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 os
import copy

import torch
import torch.backends.cudnn as cudnn
import numpy as np
import torch.distributed as dist
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from models.vae import AutoencoderKL, VAE_MODEL_DICT
from models import simflow
import util.misc as misc
from engine import 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()

    # 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,
        )
    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).eval()
    requires_grad(model, False)
    model_without_ddp = model

    # Load VAE and create an EMA of the VAE
    vae = AutoencoderKL(embed_dim=vae_embed_dim, ch_mult=ch_mult, use_variational=args.use_variational, fixed_std=args.fixed_std).to(device)
    if args.vae_path is not None:
        vae.init_from_ckpt(args.vae_path)
    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.eval()
    requires_grad(vae, False)
    vae_without_ddp = vae

    # resume training
    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_ema'])
        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()]

        missing, unexpected = model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
        print("missing model params: ", missing)
        print("unexpected model params: ", unexpected)
        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)

        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")

    # Evaluate
    evaluate(model_without_ddp, vae, model_ema_params, args, 0, batch_size=args.eval_bsz, cfg=args.cfg, use_ema=args.use_ema)



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("--resolution", type=int, default=256)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--class_num", type=int, default=1000)

    # optimization params
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--resume_step', default=-1, type=int, help='resume from checkpoint at a specific step')
    
    # Model
    parser.add_argument('--model_type', default="tarflow", type=str)
    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('--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')
    
    # 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("--vae_ln", default="True", type=misc.str_to_bool, help='Whether to use layer norm in the VAE')
    parser.add_argument('--fixed_std', default=1.0, type=float, help='Fixed standard deviation for the VAE')

    # Generation parameters
    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('--use_ema', default="True", type=misc.str_to_bool, help='Whether to use the EMA model')
    parser.add_argument('--label_drop_prob', default=0.1, type=float)
    parser.add_argument('--evaluate', action='store_true')
    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)
    parser.add_argument('--specific_cls_to_gen', default=-1, type=int, help='specific class to generate')

    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)