inference_rot.py

"""
Sampling Scripts of LightningDiT.

by Maple (Jingfeng Yao) from HUST-VL
"""

import os, math, json, pickle, logging, argparse, yaml, torch, numpy as np
from time import time, strftime
from glob import glob
from copy import deepcopy
from collections import OrderedDict
from PIL import Image
from tqdm import tqdm
import torch.distributed as dist
from accelerate import Accelerator
from torch.utils.data import DataLoader
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.tensorboard import SummaryWriter
import torchvision
# local imports
from tokenizer.vavae import VA_VAE
from models.lightningdit_rot import LightningDiT_models
from transport import create_transport, Sampler
from datasets.img_latent_dataset import ImgLatentDataset
from visualize_attention import visualize_attention_matrix

# sample function
def do_sample(train_config, accelerator, ckpt_path=None, cfg_scale=None, model=None, vae=None, demo_sample_mode=False):
    """
    Run sampling.
    """

    folder_name = f"{train_config['model']['model_type'].replace('/', '-')}-ckpt-{ckpt_path.split('/')[-1].split('.')[0]}-{train_config['sample']['sampling_method']}-{train_config['sample']['num_sampling_steps']}".lower()
    # folder_name = "test_speed"
    if cfg_scale is None:
        cfg_scale = train_config['sample']['cfg_scale']
    cfg_interval_start = train_config['sample']['cfg_interval_start'] if 'cfg_interval_start' in train_config['sample'] else 0
    timestep_shift = train_config['sample']['timestep_shift'] if 'timestep_shift' in train_config['sample'] else 0
    if cfg_scale > 1.0:
        folder_name += f"-interval{cfg_interval_start:.2f}"+f"-cfg{cfg_scale:.2f}"
        folder_name += f"-shift{timestep_shift:.2f}"

    if demo_sample_mode:
        cfg_interval_start = 0
        timestep_shift = 0
        cfg_scale = 9.0

    sample_folder_dir = os.path.join(train_config['train']['output_dir'], train_config['train']['exp_name'], folder_name)
    if accelerator.process_index == 0:
        if not demo_sample_mode:
            print_with_prefix('Sample_folder_dir=', sample_folder_dir)
        print_with_prefix('ckpt_path=', ckpt_path)
        print_with_prefix('cfg_scale=', cfg_scale)
        print_with_prefix('cfg_interval_start=', cfg_interval_start)
        print_with_prefix('timestep_shift=', timestep_shift)

    if not os.path.exists(sample_folder_dir):
        if accelerator.process_index == 0:
            os.makedirs(sample_folder_dir, exist_ok=True) 
    else:
        png_files = [f for f in os.listdir(sample_folder_dir) if f.endswith('.png')]
        png_count = len(png_files)
        if png_count > train_config['sample']['fid_num']:
            if accelerator.process_index == 0:
                print_with_prefix(f"Found {png_count} PNG files in {sample_folder_dir}, skip sampling.")
            return sample_folder_dir

    torch.backends.cuda.matmul.allow_tf32 = True  # True: fast but may lead to some small numerical differences
    assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage"
    torch.set_grad_enabled(False)

    # Setup accelerator:
    device = accelerator.device

    # Setup DDP:
    device = accelerator.device
    seed = train_config['train']['global_seed'] * accelerator.num_processes + accelerator.process_index
    torch.manual_seed(seed)
    # torch.cuda.set_device(device)
    print_with_prefix(f"Starting rank={accelerator.local_process_index}, seed={seed}, world_size={accelerator.num_processes}.")
    rank = accelerator.local_process_index

    # Load model:
    if 'downsample_ratio' in train_config['vae']:
        downsample_ratio = train_config['vae']['downsample_ratio']
    else:
        downsample_ratio = 16
    latent_size = train_config['data']['image_size'] // downsample_ratio

    checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
    if "ema" in checkpoint:  # supports checkpoints from train.py
        checkpoint = checkpoint["ema"]
    model.load_state_dict(checkpoint)
    model.eval()  # important!
    model.to(device)

    transport = create_transport(
        train_config['transport']['path_type'],
        train_config['transport']['prediction'],
        train_config['transport']['loss_weight'],
        train_config['transport']['train_eps'],
        train_config['transport']['sample_eps'],
        use_cosine_loss = train_config['transport']['use_cosine_loss'] if 'use_cosine_loss' in train_config['transport'] else False,
        use_lognorm = train_config['transport']['use_lognorm'] if 'use_lognorm' in train_config['transport'] else False,
    )  # default: velocity;
    sampler = Sampler(transport)
    mode = train_config['sample']['mode']
    if mode == "ODE":
        sample_fn = sampler.sample_ode(
            sampling_method=train_config['sample']['sampling_method'],
            num_steps=train_config['sample']['num_sampling_steps'],
            atol=train_config['sample']['atol'],
            rtol=train_config['sample']['rtol'],
            reverse=train_config['sample']['reverse'],
            timestep_shift=timestep_shift,
        )
    else:
        raise NotImplementedError(f"Sampling mode {mode} is not supported.")
    
    if vae is None:
        vae = VA_VAE(
            f'tokenizer/configs/{train_config["vae"]["model_name"]}.yaml',
        )
        if accelerator.process_index == 0:
            print_with_prefix('Loaded VAE model')

    using_cfg = cfg_scale > 1.0
    if using_cfg:
        if accelerator.process_index == 0:
            print_with_prefix('Using cfg:', using_cfg)

    if rank == 0:
        os.makedirs(sample_folder_dir, exist_ok=True)
        if accelerator.process_index == 0 and not demo_sample_mode:
            print_with_prefix(f"Saving .png samples at {sample_folder_dir}")
    accelerator.wait_for_everyone()

    # Figure out how many samples we need to generate on each GPU and how many iterations we need to run:
    n = train_config['sample']['per_proc_batch_size']
    global_batch_size = n * accelerator.num_processes
    # To make things evenly-divisible, we'll sample a bit more than we need and then discard the extra samples:
    num_samples = len([name for name in os.listdir(sample_folder_dir) if (os.path.isfile(os.path.join(sample_folder_dir, name)) and ".png" in name)])
    total_samples = int(math.ceil(train_config['sample']['fid_num'] / global_batch_size) * global_batch_size)
    if rank == 0:
        if accelerator.process_index == 0:
            print_with_prefix(f"Total number of images that will be sampled: {total_samples}")
    assert total_samples % accelerator.num_processes == 0, "total_samples must be divisible by world_size"
    samples_needed_this_gpu = int(total_samples // accelerator.num_processes)
    assert samples_needed_this_gpu % n == 0, "samples_needed_this_gpu must be divisible by the per-GPU batch size"
    iterations = int(samples_needed_this_gpu // n)
    done_iterations = int( int(num_samples // accelerator.num_processes) // n)
    pbar = range(iterations)
    if not demo_sample_mode:
        pbar = tqdm(pbar) if rank == 0 else pbar
    total = 0
    
    if accelerator.process_index == 0:
        print_with_prefix("Using latent normalization")
    dataset = ImgLatentDataset(
        data_dir=train_config['data']['data_path'],
        latent_norm=train_config['data']['latent_norm'] if 'latent_norm' in train_config['data'] else False,
        latent_multiplier=train_config['data']['latent_multiplier'] if 'latent_multiplier' in train_config['data'] else 0.18215,
    )
    latent_mean, latent_std = dataset.get_latent_stats()
    latent_multiplier = train_config['data']['latent_multiplier'] if 'latent_multiplier' in train_config['data'] else 0.18215
    # move to device
    latent_mean = latent_mean.clone().detach().to(device)
    latent_std = latent_std.clone().detach().to(device)


    # if demo_sample_mode:
    #     if accelerator.process_index == 0:
    #         images = []
    #         for label in tqdm([975, 3, 207, 387, 388, 88, 979, 279], desc="Generating Demo Samples"):
    #             z = torch.randn(1, model.in_channels, latent_size, latent_size, device=device)
    #             y = torch.tensor([label], device=device)
    #             z = torch.cat([z, z], 0)
    #             y_null = torch.tensor([1000] * 1, device=device)
    #             y = torch.cat([y, y_null], 0)
    #             model_kwargs = dict(y=y, cfg_scale=cfg_scale, cfg_interval=False, cfg_interval_start=cfg_interval_start)
    #             model_fn = model.forward_with_cfg
    #             samples = sample_fn(z, model_fn, **model_kwargs)[-1]
    #             samples = (samples * latent_std) / latent_multiplier + latent_mean
    #             samples = vae.decode_to_images(samples)
    #             images.append(samples)
    #         # Combine 8 images into a 2x4 grid
    #         os.makedirs('demo_images', exist_ok=True)
    #         # Stack all images into a large numpy array
    #         all_images = np.stack([img[0] for img in images])  # Take first image from each batch            
    #         # Rearrange into 2x4 grid
    #         h, w = all_images.shape[1:3]
    #         grid = np.zeros((2 * h, 4 * w, 3), dtype=np.uint8)
    #         for idx, image in enumerate(all_images):
    #             i, j = divmod(idx, 4)  # Calculate position in 2x4 grid
    #             grid[i*h:(i+1)*h, j*w:(j+1)*w] = image
                
    #         # Save the combined image
    #         Image.fromarray(grid).save('demo_images/demo_samples.png')

    #         return None
    if demo_sample_mode:
        # Demo mode: sample exactly one image and save it, then return.
        # Choose a demo label (can be changed or made an argument)
        demo_label = 975
        # create latent noise for one sample
        z = torch.randn(1, model.in_channels, latent_size, latent_size, device=device)
        y = torch.tensor([demo_label], device=device)

        # Setup classifier-free guidance if needed
        if using_cfg:
            z = torch.cat([z, z], 0)
            y_null = torch.tensor([1000], device=device)
            y = torch.cat([y, y_null], 0)
            model_kwargs = dict(y=y, cfg_scale=cfg_scale, cfg_interval=False, cfg_interval_start=cfg_interval_start)
            model_fn = model.forward_with_cfg
        else:
            model_kwargs = dict(y=y)
            model_fn = model.forward

        # Run sampling (single batch)
        samples = sample_fn(z, model_fn, **model_kwargs)[-1]

        if using_cfg:
            # samples contains [cond; uncond] stacked, keep the conditional output
            samples, _ = samples.chunk(2, dim=0)

        # un-normalize and decode
        samples = (samples * latent_std) / latent_multiplier + latent_mean
        images = vae.decode_to_images(samples)

        # save first image
        if accelerator.process_index == 0:
            os.makedirs(sample_folder_dir, exist_ok=True)
            Image.fromarray(images[0]).save(os.path.join(sample_folder_dir, 'demo_sample.png'))
            print_with_prefix(f"Saved demo sample to {os.path.join(sample_folder_dir, 'demo_sample.png')}")

        return sample_folder_dir
    
    else:
        # 初始化时间统计变量
        total_sampling_time = 0
        total_vae_decode_time = 0
        total_images_generated = 0
        batch_times = []
        
        for i in pbar:
            batch_start_time = time()
            # print("starting batch ", i)

            # Sample inputs:
            z = torch.randn(n, model.in_channels, latent_size, latent_size, device=device)
            y = torch.randint(0, train_config['data']['num_classes'], (n,), device=device)
            
            # Setup classifier-free guidance:
            if using_cfg:
                z = torch.cat([z, z], 0)
                y_null = torch.tensor([1000] * n, device=device)
                y = torch.cat([y, y_null], 0)
                model_kwargs = dict(y=y, cfg_scale=cfg_scale, cfg_interval=True, cfg_interval_start=cfg_interval_start)
                model_fn = model.forward_with_cfg
            else:
                model_kwargs = dict(y=y)
                model_fn = model.forward

            # 记录采样开始时间
            sampling_start_time = time()
            # print("starting sampling batch ", i)
            samples = sample_fn(z, model_fn, **model_kwargs)[-1]
            sampling_end_time = time()
            
            if using_cfg:
                samples, _ = samples.chunk(2, dim=0)  # Remove null class samples

            samples = (samples * latent_std) / latent_multiplier + latent_mean
            
            # 记录VAE解码开始时间
            vae_decode_start_time = time()
            # print("starting VAE decode batch ", i)
            samples = vae.decode_to_images(samples)
            vae_decode_end_time = time()

            # Save samples to disk as individual .png files
            # print("start saving batch ", i)
            for j, sample in enumerate(samples):
                index = j * accelerator.num_processes + accelerator.process_index + total
                Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")
            
            # 统计时间
            batch_end_time = time()
            batch_time = batch_end_time - batch_start_time
            sampling_time = sampling_end_time - sampling_start_time
            vae_decode_time = vae_decode_end_time - vae_decode_start_time
            
            batch_times.append(batch_time)
            total_sampling_time += sampling_time
            total_vae_decode_time += vae_decode_time
            total_images_generated += len(samples)
            
            # 每10个batch输出一次统计信息
            if accelerator.process_index == 0 and (i + 1) % 10 == 0:
                avg_sampling_time_per_image = total_sampling_time / total_images_generated
                avg_vae_time_per_image = total_vae_decode_time / total_images_generated
                avg_total_time_per_image = sum(batch_times) / total_images_generated
                # print_with_prefix(f"Batch {i+1}/{iterations}: Avg sampling time per image: {avg_sampling_time_per_image:.3f}s, "
                #                 f"Avg VAE decode time per image: {avg_vae_time_per_image:.3f}s, "
                #                 f"Avg total time per image: {avg_total_time_per_image:.3f}s")
            
            total += global_batch_size
            accelerator.wait_for_everyone()
        
        # 输出最终统计结果
        if accelerator.process_index == 0 and total_images_generated > 0:
            avg_sampling_time_per_image = total_sampling_time / total_images_generated
            avg_vae_time_per_image = total_vae_decode_time / total_images_generated
            avg_total_time_per_image = sum(batch_times) / total_images_generated
            print_with_prefix("=" * 60)
            print_with_prefix("FINAL TIMING STATISTICS:")
            print_with_prefix(f"Total images generated: {total_images_generated}")
            print_with_prefix(f"Average sampling time per image: {avg_sampling_time_per_image:.3f} seconds")
            print_with_prefix(f"Average VAE decode time per image: {avg_vae_time_per_image:.3f} seconds") 
            print_with_prefix(f"Average total time per image: {avg_total_time_per_image:.3f} seconds")
            print_with_prefix(f"Total sampling throughput: {total_images_generated / sum(batch_times):.2f} images/second")
            print_with_prefix("=" * 60)

    return sample_folder_dir

# some utils
def print_with_prefix(*messages):
    prefix = f"\033[34m[LightningDiT-Sampling {strftime('%Y-%m-%d %H:%M:%S')}]\033[0m"
    combined_message = ' '.join(map(str, messages))
    print(f"{prefix}: {combined_message}")

def load_config(config_path):
    with open(config_path, "r") as file:
        config = yaml.safe_load(file)
    return config

if __name__ == "__main__":

    # read config
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='configs/lightningdit_b_ldmvae_f16d16.yaml')
    parser.add_argument('--demo', action='store_true', default=False)
    args = parser.parse_args()
    accelerator = Accelerator()
    train_config = load_config(args.config)

    # get ckpt_dir
    assert 'ckpt_path' in train_config, "ckpt_path must be specified in config"
    if accelerator.process_index == 0:
        print_with_prefix('Using ckpt:', train_config['ckpt_path'])
    ckpt_dir = train_config['ckpt_path']

    if 'downsample_ratio' in train_config['vae']:
        latent_size = train_config['data']['image_size'] // train_config['vae']['downsample_ratio']
    else:
        latent_size = train_config['data']['image_size'] // 16

    # get model
    model = LightningDiT_models[train_config['model']['model_type']](
        input_size=latent_size,
        num_classes=train_config['data']['num_classes'],
        use_qknorm=train_config['model']['use_qknorm'],
        use_swiglu=train_config['model']['use_swiglu'] if 'use_swiglu' in train_config['model'] else False,
        use_rope=train_config['model']['use_rope'] if 'use_rope' in train_config['model'] else False,
        use_rmsnorm=train_config['model']['use_rmsnorm'] if 'use_rmsnorm' in train_config['model'] else False,
        wo_shift=train_config['model']['wo_shift'] if 'wo_shift' in train_config['model'] else False,
        in_channels=train_config['model']['in_chans'] if 'in_chans' in train_config['model'] else 4,
        learn_sigma=train_config['model']['learn_sigma'] if 'learn_sigma' in train_config['model'] else False,
        degree=train_config['model']['degree'] if 'degree' in train_config['model'] else '180',
    )

    # naive sample
    sample_folder_dir = do_sample(train_config, accelerator, ckpt_path=ckpt_dir, model=model, demo_sample_mode=args.demo)

    #visualize attention map
    # attn_weights_list = []
    # for block in model.blocks:
    #     attn = block.attn
    #     attn_weights = attn.attn_weights[:1] #cfg情况下取条件分支的注意力权重
    #     # print("attn_weights.shape:", attn_weights.shape)
    #     attn_weights_list.append(attn_weights)
    # attn_weights = torch.cat(attn_weights_list, dim=0)  # (num_layers, num_heads, N, N)
    # # print("Concatenated attn_weights shape:", attn_weights.shape)
    # print("visualizing attention maps...")
    # visualize_attention_matrix(attn_weights, 
    #                            save_path="/home/jiayou.zhang/hom/personal/jinyuan/LightningDiT/attention_maps_new/xl800",
    #                            model_name="xl800-cfg-cond", 
    #                            method="heatmap")
    
    if not args.demo:
        # calculate FID
        # Important: FID is only for reference, please use ADM evaluation for paper reporting
        if accelerator.process_index == 0:
            from tools.calculate_fid import calculate_fid_given_paths
            print_with_prefix('Calculating FID with {} number of samples'.format(train_config['sample']['fid_num']))
            assert 'fid_reference_file' in train_config['data'], "fid_reference_file must be specified in config"
            fid_reference_file = train_config['data']['fid_reference_file']
            fid = calculate_fid_given_paths(
                [fid_reference_file, sample_folder_dir],
                batch_size=50,
                dims=2048,
                device='cuda',
                num_workers=8,
                sp_len = train_config['sample']['fid_num']
            )
            print_with_prefix('fid=',fid)