support rectified flow

diffusion/rectified_flow.py

@@ -0,0 +1,85 @@
+import torch 
+
+class RectifiedFlow(torch.nn.Module):
+    def __init__(self, model, ln=True):
+        super().__init__()
+        self.model = model
+        self.ln = ln
+        self.stratified = False
+
+    def forward(self, x, cond):
+
+        b = x.size(0)
+        if self.ln:
+            if self.stratified:
+                # stratified sampling of normals
+                # first stratified sample from uniform
+                quantiles = torch.linspace(0, 1, b + 1).to(x.device)
+                z = quantiles[:-1] + torch.rand((b,)).to(x.device) / b
+                # now transform to normal
+                z = torch.erfinv(2 * z - 1) * math.sqrt(2)
+                t = torch.sigmoid(z)
+            else:
+                nt = torch.randn((b,)).to(x.device)
+                t = torch.sigmoid(nt)
+        else:
+            t = torch.rand((b,)).to(x.device)
+        texp = t.view([b, *([1] * len(x.shape[1:]))])
+        z1 = torch.randn_like(x)
+        zt = (1 - texp) * x + texp * z1
+
+        # make t, zt into same dtype as x
+        zt, t = zt.to(x.dtype), t.to(x.dtype)
+        vtheta = self.model(zt, t, cond) 
+        if self.model.learn_sigma == True: 
+            vtheta, _ = vtheta.chunk(2, dim=1) 
+        batchwise_mse = ((z1 - x - vtheta) ** 2).mean(dim=list(range(1, len(x.shape))))
+        tlist = batchwise_mse.detach().cpu().reshape(-1).tolist()
+        ttloss = [(tv, tloss) for tv, tloss in zip(t, tlist)]
+        return batchwise_mse.mean(), {"batchwise_loss": ttloss}
+
+    @torch.no_grad()
+    def sample(self, z, cond, null_cond=None, sample_steps=50, cfg=2.0):
+        b = z.size(0)
+        dt = 1.0 / sample_steps
+        dt = torch.tensor([dt] * b).to(z.device).view([b, *([1] * len(z.shape[1:]))])
+        images = [z]
+        for i in range(sample_steps, 0, -1):
+            t = i / sample_steps
+            t = torch.tensor([t] * b).to(z.device)
+
+            vc = self.model(z, t, cond) 
+            if self.model.learn_sigma == True: 
+                vc, _ = vc.chunk(2, dim=1) 
+            if null_cond is not None:
+                vu = self.model(z, t, null_cond) 
+                if self.model.learn_sigma == True: 
+                    vu, _ = vu.chunk(2, dim=1) 
+                vc = vu + cfg * (vc - vu)
+
+            z = z - dt * vc
+            images.append(z)
+        return images
+
+    @torch.no_grad()
+    def sample_with_xps(self, z, cond, null_cond=None, sample_steps=50, cfg=2.0):
+        b = z.size(0)
+        dt = 1.0 / sample_steps
+        dt = torch.tensor([dt] * b).to(z.device).view([b, *([1] * len(z.shape[1:]))])
+        images = [z]
+        for i in range(sample_steps, 0, -1):
+            t = i / sample_steps
+            t = torch.tensor([t] * b).to(z.device)
+
+            vc = self.model(z, t, cond) 
+            if self.model.learn_sigma == True: 
+                vc, _ = vc.chunk(2, dim=1)  
+            if null_cond is not None:
+                vu = self.model(z, t, null_cond) 
+                if self.model.learn_sigma == True: 
+                    vu, _ = vu.chunk(2, dim=1) 
+                vc = vu + cfg * (vc - vu)
+            x = z - i * dt * vc
+            z = z - dt * vc
+            images.append(x)
+        return images

sample.py

@@ -15,6 +15,7 @@ from diffusion import create_diffusion
 from diffusers.models import AutoencoderKL
 from download import find_model
 from models import DiT_models
+from diffusion.rectified_flow import RectifiedFlow 
 import argparse
 
 
@@ -60,17 +61,21 @@ def main(args):
         elif args.model == "DiT-B/2":
             ckpt_path = "dit_moe_b_8E2A.pt" 
         elif args.model == "DiT-XL/2": 
-            ckpt_path = "results/deepspeed-DiT-XL-2/checkpoints/ckpt.pt" 
+            ckpt_path = "results/deepspeed-DiT-XL-2-rf/checkpoints/tmp.pt" 
         else: 
-            ckpt_path = "results/deepspeed-DiT-G-2/checkpoints/ckpt.pt" 
+            ckpt_path = "results/deepspeed-DiT-G-2-rf/checkpoints/tmp.pt" 
     else:
         ckpt_path = args.ckpt 
 
 
     state_dict = find_model(ckpt_path)
     model.load_state_dict(state_dict)
-    model.eval()  # important!
+    model.eval()  # important! 
-    diffusion = create_diffusion(str(args.num_sampling_steps))
+
+    if args.rf:
+        diffusion = RectifiedFlow(model)
+    else:
+        diffusion = create_diffusion(str(args.num_sampling_steps))
     vae = AutoencoderKL.from_pretrained(args.vae_path).to(device) 
     
     # Labels to condition the model with (feel free to change):
@@ -88,17 +93,28 @@ def main(args):
     model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
 
     if dtype == torch.float16: 
-        with torch.autocast(device_type='cuda'):
+        if args.rf: 
-            samples = diffusion.p_sample_loop(
+            with torch.autocast(device_type='cuda'):
-                model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+                STEPSIZE = 50
-            )
+                init_noise = torch.randn(n, 4, latent_size, latent_size, device=device) 
+                conds = torch.tensor(class_labels, device=device)
+                images = diffusion.sample_with_xps(init_noise, conds, null_cond = torch.tensor([1000] * n).cuda(), sample_steps = STEPSIZE, cfg = 7.0)
+                samples = vae.decode(images[-1] / 0.18215).sample
+        
+        else:
+            with torch.autocast(device_type='cuda'):
+                samples = diffusion.p_sample_loop(
+                    model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+                )
+            samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+            samples = vae.decode(samples / 0.18215).sample
+    
     else:
         samples = diffusion.p_sample_loop(
                 model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
             )
-    
+        samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
-    samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+        samples = vae.decode(samples / 0.18215).sample
-    samples = vae.decode(samples / 0.18215).sample
 
     # Save and display images: 
     if args.model == "DiT-S/2":
@@ -113,15 +129,16 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2")
+    parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-G/2")
     parser.add_argument("--vae-path", type=str, default="/maindata/data/shared/multimodal/zhengcong.fei/ckpts/sd-vae-ft-mse")
     parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
     parser.add_argument("--num-classes", type=int, default=1000)
     parser.add_argument("--cfg-scale", type=float, default=4.0)
-    parser.add_argument('--num_experts', default=8, type=int,) 
+    parser.add_argument('--num_experts', default=16, type=int,) 
     parser.add_argument('--num_experts_per_tok', default=2, type=int,) 
     parser.add_argument("--num-sampling-steps", type=int, default=250)
     parser.add_argument("--seed", type=int, default=2024)
     parser.add_argument("--ckpt", type=str, default=None, )
+    parser.add_argument("--rf", type=bool, default=True) 
     args = parser.parse_args()
-    main(args)
+    main(args) 

sample_ddp.py

@@ -0,0 +1,166 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Samples a large number of images from a pre-trained DiT model using DDP.
+Subsequently saves a .npz file that can be used to compute FID and other
+evaluation metrics via the ADM repo: https://github.com/openai/guided-diffusion/tree/main/evaluations
+
+For a simple single-GPU/CPU sampling script, see sample.py.
+"""
+import torch
+import torch.distributed as dist
+from models import DiT_models
+from download import find_model
+from diffusion import create_diffusion
+from diffusers.models import AutoencoderKL
+from tqdm import tqdm
+import os
+from PIL import Image
+import numpy as np
+import math
+import argparse
+
+
+def create_npz_from_sample_folder(sample_dir, num=50_000):
+    """
+    Builds a single .npz file from a folder of .png samples.
+    """
+    samples = []
+    for i in tqdm(range(num), desc="Building .npz file from samples"):
+        sample_pil = Image.open(f"{sample_dir}/{i:06d}.png")
+        sample_np = np.asarray(sample_pil).astype(np.uint8)
+        samples.append(sample_np)
+    samples = np.stack(samples)
+    assert samples.shape == (num, samples.shape[1], samples.shape[2], 3)
+    npz_path = f"{sample_dir}.npz"
+    np.savez(npz_path, arr_0=samples)
+    print(f"Saved .npz file to {npz_path} [shape={samples.shape}].")
+    return npz_path
+
+
+def main(args):
+    """
+    Run sampling.
+    """
+    torch.backends.cuda.matmul.allow_tf32 = args.tf32  # 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 DDP:
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    device = rank % torch.cuda.device_count()
+    seed = args.global_seed * dist.get_world_size() + rank
+    torch.manual_seed(seed)
+    torch.cuda.set_device(device)
+    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
+
+    if args.ckpt is None:
+        assert args.model == "DiT-XL/2", "Only DiT-XL/2 models are available for auto-download."
+        assert args.image_size in [256, 512]
+        assert args.num_classes == 1000
+
+    # Load model:
+    latent_size = args.image_size // 8
+    model = DiT_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes
+    ).to(device)
+    # Auto-download a pre-trained model or load a custom DiT checkpoint from train.py:
+    ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt"
+    state_dict = find_model(ckpt_path)
+    model.load_state_dict(state_dict)
+    model.eval()  # important!
+    diffusion = create_diffusion(str(args.num_sampling_steps))
+    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+    assert args.cfg_scale >= 1.0, "In almost all cases, cfg_scale be >= 1.0"
+    using_cfg = args.cfg_scale > 1.0
+
+    # Create folder to save samples:
+    model_string_name = args.model.replace("/", "-")
+    ckpt_string_name = os.path.basename(args.ckpt).replace(".pt", "") if args.ckpt else "pretrained"
+    folder_name = f"{model_string_name}-{ckpt_string_name}-size-{args.image_size}-vae-{args.vae}-" \
+                  f"cfg-{args.cfg_scale}-seed-{args.global_seed}"
+    sample_folder_dir = f"{args.sample_dir}/{folder_name}"
+    if rank == 0:
+        os.makedirs(sample_folder_dir, exist_ok=True)
+        print(f"Saving .png samples at {sample_folder_dir}")
+    dist.barrier()
+
+    # Figure out how many samples we need to generate on each GPU and how many iterations we need to run:
+    n = args.per_proc_batch_size
+    global_batch_size = n * dist.get_world_size()
+    # To make things evenly-divisible, we'll sample a bit more than we need and then discard the extra samples:
+    total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size)
+    if rank == 0:
+        print(f"Total number of images that will be sampled: {total_samples}")
+    assert total_samples % dist.get_world_size() == 0, "total_samples must be divisible by world_size"
+    samples_needed_this_gpu = int(total_samples // dist.get_world_size())
+    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)
+    pbar = range(iterations)
+    pbar = tqdm(pbar) if rank == 0 else pbar
+    total = 0
+    for _ in pbar:
+        # Sample inputs:
+        z = torch.randn(n, model.in_channels, latent_size, latent_size, device=device)
+        y = torch.randint(0, args.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=args.cfg_scale)
+            sample_fn = model.forward_with_cfg
+        else:
+            model_kwargs = dict(y=y)
+            sample_fn = model.forward
+
+        # Sample images:
+        samples = diffusion.p_sample_loop(
+            sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=False, device=device
+        )
+        if using_cfg:
+            samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+
+        samples = vae.decode(samples / 0.18215).sample
+        samples = torch.clamp(127.5 * samples + 128.0, 0, 255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+
+        # Save samples to disk as individual .png files
+        for i, sample in enumerate(samples):
+            index = i * dist.get_world_size() + rank + total
+            Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")
+        total += global_batch_size
+
+    # Make sure all processes have finished saving their samples before attempting to convert to .npz
+    dist.barrier()
+    if rank == 0:
+        create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
+        print("Done.")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2")
+    parser.add_argument("--vae",  type=str, choices=["ema", "mse"], default="ema")
+    parser.add_argument("--sample-dir", type=str, default="samples")
+    parser.add_argument("--per-proc-batch-size", type=int, default=32)
+    parser.add_argument("--num-fid-samples", type=int, default=50_000)
+    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--cfg-scale",  type=float, default=1.5)
+    parser.add_argument("--num-sampling-steps", type=int, default=250)
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True,
+                        help="By default, use TF32 matmuls. This massively accelerates sampling on Ampere GPUs.")
+    parser.add_argument("--ckpt", type=str, default=None,
+                        help="Optional path to a DiT checkpoint (default: auto-download a pre-trained DiT-XL/2 model).")
+    args = parser.parse_args()
+    main(args)

test.py

@@ -106,12 +106,12 @@ def image_class_expert_ratio():
 
 def ckpts_clean(): 
     # only save ema ckpts for ckpt uploading 
-    ckpt_path = 'results/003-DiT-B-2/checkpoints/ckpt.pt' 
+    ckpt_path = 'results/003-DiT-B-2/checkpoints/ckpt2.pt' 
     checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage) 
     new_checkpoint = {
                         "ema": checkpoint['ema'],
                     }
-    torch.save(new_checkpoint, 'ckpt_clean.pt')
+    torch.save(new_checkpoint, 'dit_moe_b_8E2A.pt')
 
 
 

train.py

@@ -275,7 +275,7 @@ if __name__ == "__main__":
     parser.add_argument("--num-classes", type=int, default=1000)
     parser.add_argument("--epochs", type=int, default=1400) 
     parser.add_argument("--global-batch-size", type=int, default=64)
-    parser.add_argument("--global-seed", type=int, default=1234) 
+    parser.add_argument("--global-seed", type=int, default=2024) 
     parser.add_argument("--num-workers", type=int, default=4)
     parser.add_argument("--log-every", type=int, default=100)
     parser.add_argument('--accum_iter', default=8, type=int,)  

train_deepspeed.py

@@ -29,6 +29,7 @@ import os
 
 from models import DiT_models
 from diffusion import create_diffusion
+from diffusion.rectified_flow import RectifiedFlow
 from diffusers.models import AutoencoderKL
 from download import find_model
 import deepspeed
@@ -115,7 +116,10 @@ def main(args):
 
     # Setup an experiment folder 
     model_string_name = args.model.replace("/", "-")  # e.g., DiT-XL/2 --> DiT-XL-2 (for naming folders)
-    experiment_dir = f"{args.results_dir}/deepspeed-{model_string_name}"  # Create an experiment folder
+    if args.rf: 
+        experiment_dir = f"{args.results_dir}/deepspeed-{model_string_name}-rf"
+    else:
+        experiment_dir = f"{args.results_dir}/deepspeed-{model_string_name}"  # Create an experiment folder
     checkpoint_dir = f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
     
     if rank == 0:
@@ -142,9 +146,12 @@ def main(args):
         print('load from: ', args.resume) 
         state_dict = find_model(args.resume)
         model.load_state_dict(state_dict)
-
+    
-
+    if args.rf: 
-    diffusion = create_diffusion(timestep_respacing="")  # default: 1000 steps, linear noise schedule 
+        logger.info("train with rectified flow")
+        diffusion = RectifiedFlow(model)
+    else:
+        diffusion = create_diffusion(timestep_respacing="")  # default: 1000 steps, linear noise schedule 
     vae = AutoencoderKL.from_pretrained(args.vae_path).to(device)
     logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
 
@@ -194,11 +201,17 @@ def main(args):
             with torch.no_grad():
                 # Map input images to latent space + normalize latents:
                 x = vae.encode(x).latent_dist.sample().mul_(0.18215)
-            t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
+            if args.rf: 
-            model_kwargs = dict(y=y)
+                with torch.autocast(device_type='cuda'): 
-            with torch.autocast(device_type='cuda'):
+                    loss, _ = diffusion.forward(x, y)
-                loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
+            
-            loss = loss_dict["loss"].mean() 
+            else:
+                t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
+                model_kwargs = dict(y=y)
+                with torch.autocast(device_type='cuda'):
+                    loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
+                loss = loss_dict["loss"].mean() 
+            
             model_engine.backward(loss) 
             model_engine.step()
             
@@ -223,9 +236,13 @@ def main(args):
                 start_time = time()
 
             # Save DiT checkpoint:
-            if train_steps % args.ckpt_every == 0 and train_steps > 0:                 
+            if train_steps % args.ckpt_every == 0 and train_steps > 0:    
-                checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}"
+                try:             
-                model_engine.save_checkpoint(checkpoint_path) 
+                    checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}"
+                    model_engine.save_checkpoint(checkpoint_path) 
+                except Exception as e: 
+                    print(e)
+                
                 dist.barrier()
 
     # model.eval()  # important! This disables randomized embedding dropout
@@ -244,14 +261,15 @@ if __name__ == "__main__":
     parser.add_argument("--num-classes", type=int, default=1000)
     parser.add_argument("--epochs", type=int, default=1400) 
     parser.add_argument("--train_batch_size", type=int, default=2)
-    parser.add_argument("--global-seed", type=int, default=1234) 
+    parser.add_argument("--global-seed", type=int, default=2024) 
-    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--num-workers", type=int, default=0)
     parser.add_argument("--log-every", type=int, default=100)
     parser.add_argument('--accum_iter', default=8, type=int,)  
     parser.add_argument('--num_experts', default=8, type=int,) 
     parser.add_argument('--num_experts_per_tok', default=2, type=int,) 
     parser.add_argument("--ckpt-every", type=int, default=10_000) 
     parser.add_argument('--local-rank', type=int, default=-1, help='local rank passed from distributed launcher') 
+    parser.add_argument("--rf", type=bool, default=False) 
     parser = deepspeed.add_config_arguments(parser)
     args = parser.parse_args() 
     print(args)