Upload inference_rot_head.py with huggingface_hub

inference_rot_head.py

@@ -0,0 +1,407 @@
+"""
+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_head_vis 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,
+        num_rot=train_config['model']['num_rot'] if 'num_rot' in train_config['model'] else 4
+    )
+
+    # 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)