eval_vae.py

# Copyright (c) 2025 Jingfeng Yao
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT license
#
# This file has been modified by ByteDance Ltd. and/or its affiliates on 2025-9-10.
#
# Original file was released under MIT license, with the full license text
# available at https://spdx.org/licenses/MIT.html.
#
# This modified file is released under the Apache-2.0 License.
# 
# Reference:
#    https://github.com/hustvl/LightningDiT/blob/main/evaluate_tokenizer.py


import argparse
import os
import shutil

import cv2
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from torchmetrics import StructuralSimilarityIndexMeasure
import numpy as np
from tqdm import tqdm

from models.vae import AutoencoderKL, VAE_MODEL_DICT
from util.crop import center_crop_arr
import util.misc as misc
from loss.lpips import LPIPS
from util.tools import calculate_fid_given_paths, calculate_psnr_between_folders


def img2save(img_tensor):
    imgs = torch.clamp(127.5 * img_tensor + 128.0, 0, 255)
    imgs = imgs.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
    return imgs

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

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

    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    print("{}".format(args).replace(', ', ',\n'))

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

    # set up the logger and checkpoint dirs
    save_dir = os.path.join(args.output_dir, args.exp_name, "reconstructed")
    ref_dir = os.path.join(args.output_dir, args.exp_name, "reference")
    if global_rank == 0:
        os.makedirs(args.output_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
        os.makedirs(save_dir, exist_ok=True)
        os.makedirs(ref_dir, exist_ok=True)

    # Load VAE and create a EMA of the VAE
    vae_config = VAE_MODEL_DICT[args.vae]
    vae_embed_dim = vae_config["embed_dim"]
    ch_mult = vae_config["ch_mult"]
    vae = AutoencoderKL(embed_dim=vae_embed_dim, ch_mult=ch_mult, 
                        ln=args.vae_ln, use_variational=args.use_variational, fixed_std=args.fixed_std).to(device)
    if args.vae_ckpt is not None:
        vae_ckpt = torch.load(args.vae_ckpt, map_location=device)
        vae_ckpt = vae_ckpt['vae_ema']
        msg = vae.load_state_dict(vae_ckpt, strict=False)
        print("Loading pre-trained VAE")
        print("Missing keys:")
        print(msg.missing_keys)
        print("Unexpected keys:")
        print(msg.unexpected_keys)
        del vae_ckpt
    vae.eval()
    for param in vae.parameters():
        param.requires_grad = False

    # Setup data
    transform_val = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(256),
        transforms.CenterCrop(256),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    dataset_val = datasets.ImageFolder(os.path.join(args.data_path, 'val'), transform=transform_val)
    sampler_val = torch.utils.data.DistributedSampler(
        dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False
    )
    val_dataloader = torch.utils.data.DataLoader(
        dataset_val, sampler=sampler_val,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=False,
    )
    if global_rank == 0:
        print(f"Dataset contains {len(dataset_val):,} images")

    # Initialize metrics
    lpips_values = []
    ssim_values = []
    lpips = LPIPS().to(device).eval()
    ssim_metric = StructuralSimilarityIndexMeasure(data_range=(-1.0, 1.0)).to(device)

    # main training loop
    for i, (imgs, _) in enumerate(tqdm(val_dataloader)):
        imgs = imgs.to(device)
        with torch.amp.autocast("cuda") and torch.no_grad():
            _, _, z, recon_image = vae(imgs, disturb_latents="none")
            
            # Compute metrics
            lpips_values.append(lpips.forward_with_norm(recon_image, imgs).mean())
            ssim_values.append(ssim_metric(recon_image, imgs))
            
            # Save images
            recon_image = img2save(recon_image)
            ref_imgs = img2save(imgs)

            # distributed save
            for b_id in range(recon_image.shape[0]):
                gen_img = recon_image[b_id][:, :, ::-1]
                ref_img = ref_imgs[b_id][:, :, ::-1]
                save_id = args.batch_size * i + b_id
                cv2.imwrite(os.path.join(save_dir, '{}_{}.png'.format(args.gpu, str(save_id).zfill(5))), gen_img)
                cv2.imwrite(os.path.join(ref_dir, '{}_{}.png'.format(args.gpu, str(save_id).zfill(5))), ref_img)

        torch.distributed.barrier()

    lpips_values = torch.tensor(lpips_values).to(device)
    ssim_values = torch.tensor(ssim_values).to(device)
    dist.all_reduce(lpips_values, op=dist.ReduceOp.AVG)
    dist.all_reduce(ssim_values, op=dist.ReduceOp.AVG)
    
    avg_lpips = lpips_values.mean().item()
    avg_ssim = ssim_values.mean().item()

    if global_rank == 0:
        # Calculate FID
        print("Computing rFID...")
        fid = calculate_fid_given_paths([ref_dir, save_dir], batch_size=50, dims=2048, device=device, num_workers=16)

        # Calculate PSNR
        print("Computing PSNR...")
        psnr_values = calculate_psnr_between_folders(ref_dir, save_dir)
        avg_psnr = sum(psnr_values) / len(psnr_values)

        # Print final results
        print(f"Final Metrics:")
        print(f"rFID: {fid:.3f}")
        print(f"PSNR: {avg_psnr:.3f}")
        print(f"LPIPS: {avg_lpips:.3f}")
        print(f"SSIM: {avg_ssim:.3f}")

        shutil.rmtree(os.path.join(args.output_dir, args.exp_name))


def parse_args(input_args=None):
    parser = argparse.ArgumentParser(description="Evaluation")
    # 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("--exp-name", type=str, required=True)

    # dataset params
    parser.add_argument("--data_path", type=str, default="data")
    parser.add_argument("--resolution", type=int, choices=[256], default=256)
    parser.add_argument("--batch-size", type=int, default=256)

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

    # cpu params
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')

    # vae params
    parser.add_argument("--vae", type=str, default="vae_f8d4")
    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 normalization in the VAE')
    parser.add_argument("--noise-level", type=float, default=0.0)
    parser.add_argument("--vae-ckpt", type=str, default=None)

    if input_args is not None:
        args = parser.parse_args(input_args)
    else:
        args = parser.parse_args()
    return args


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