add tread sampling

Author: flixmk

README.md

@@ -11,20 +11,18 @@
   <b>CompVis Group @ LMU Munich</b> <br/>
 </p>
 
-[![Paper](https://img.shields.io/badge/arXiv-PDF-b31b1b)](https://arxiv.org/abs/2501.04765)
+[![Paper](https://img.shields.io/badge/arXiv-PDF-b31b1b)](https://arxiv.org/pdf/2501.04765)
 [![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://compvis.github.io/tread/)
 
 This repository contains the official implementation of the paper "TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training".
 
 We propose TREAD, a new method to increase the efficiency of diffusion training by improving upon iteration speed and performance at the same time. For this, we use uni-directional token transportation to modulate the information flow in the network.
 
 <div align="center">
-  <img src="./docs/images/teaser.png" alt="teaser" style="width:50%;">
+  <img src="./docs/static/images/teaser.png" alt="teaser" style="width:50%;">
 </div>
 
-## 🚀 Usage
-
-### Training
+## 🚀 Training
 
 In order to train a diffusion model, we offer a minimalistic training script in `train.py`. In its simplest form it can be started using:
 
@@ -44,9 +42,20 @@ Under `model` one can decide between `dit` and `tread` which are the preconfigur
 
 In our paper, we show that TREAD can also work on other architectures. In practice, one needs to be more careful with the routing process in order to adhere to the characteristics of the specific architecture as some have a spatial bias (RWKV, Mamba, etc.). For simplicity, we only provide code for the Transformer architecture as it is the most widely used while being robust and easy to work with.
 
-### Sampling
+## 🖼️ Sampling
+
+For most experiments we use the [EDM](https://github.com/NVlabs/edm) training and sampling to stay consistent with prior art, and the FID calculation is done via the [ADM](https://github.com/openai/guided-diffusion) evaluation suite. We provide a `fid.py` to evaluate our models during training using the same reference batches as ADM.
+
+## 💥 Guiding TREAD
+
+TREAD works great during _training_! How about _inference_? \
+It turns out TREAD can be applied during guided inference as well to gain additional performance and reduce FLOPS at the same time! \
+Instead of dropping the class label (CFG), we can guide with a selection rate delta. Since TREAD's selection rate (0.5) generalizes to other rates, this can be tuned in inference-time only.
+
+We demonstrate this in `rf.py` which contains minimal flow matching code for training and sampling:
 
-For sampling, we use the [EDM](https://github.com/NVlabs/edm) sampling, and the FID calculation is done via the [ADM](https://github.com/openai/guided-diffusion) evaluation suite. We provide a `fid.py` to evaluate our models during training using the same reference batches as ADM.
+`sample`: normal sampling\
+`sample_tread`: TREAD sampling 🔥
 
 ## 🎓 Citation
 

dit.py

@@ -228,20 +228,22 @@ def unpatchify(self, x):
 
     def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]:
         class_drop_prob = kwargs.get('class_drop_prob', 0.0)
+        force_routing = kwargs.get('force_routing', False)
+        overwrite_selection_ratio = kwargs.get('overwrite_selection_ratio', None)
 
         x = self.x_embedder(x) + self.pos_embed
         t = self.t_embedder(t)
         y = self.y_embedder(y, class_drop_prob)
         c = t + y
 
-        use_routing = self.training and self.enable_routing and self.routes
+        use_routing = (self.training and self.enable_routing and self.routes) or force_routing
         route_count = 0 if use_routing else None
         fp32_next = False
 
         for idx, block in enumerate(self.blocks):
             if use_routing and idx == self.routes[route_count]['start_layer_idx']:
                 x_D_last = x.clone()
-                mask_info = self.router.get_mask(x, mask_ratio=self.routes[route_count]['selection_ratio'])
+                mask_info = self.router.get_mask(x, mask_ratio=self.routes[route_count]['selection_ratio'] if overwrite_selection_ratio is None else overwrite_selection_ratio)
                 x = self.router.start_route(x, mask_info)
 
             if fp32_next:

inference.py

@@ -0,0 +1,52 @@
+import torch
+import hydra
+from omegaconf import DictConfig
+
+@hydra.main(config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    diffuser = hydra.utils.instantiate(cfg.diffuser).to(device)
+
+    model = hydra.utils.instantiate(cfg.model).to(device)
+    # model = load_checkpoint(model, cfg.ckpt_path, device)
+    model.eval()
+    latents = torch.randn(1, 4, 32, 32, device=device) # for ImageNet-256 + SD-VAE
+    label = torch.randint(0, 1000, (1,), device=device)
+    
+    # normal sampling
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            diffuser.sample(
+                unet=model,
+                z=model,
+                sample_steps=40,
+                cfg_scale=1.5,
+            )
+            
+    # TREAD sampling
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            # Models trained with TREAD extrapolate to unseen selection ratios (here gamma).
+            # We can use this for GUIDED sampling.
+            # For example, gamma1=0.3 and gamma2=0.7 are not seen during training.
+            # The model can still sample with these values.
+            # This enables models trained with TREAD to achieve a on-the-fly trade-off between
+            # quality and FLOPS.
+            # One can use the delta between gamma1 and gamma2 to guide the sampling,
+            # but it can also be combined with CFG (class dropout) or AutoGuidance (unet1 != unet2).
+            diffuser.sample_tread(
+                unet1=model,
+                gamma1=0.3,
+                unet2=model,
+                gamma2=0.7,
+                z=latents,
+                y=label,
+                sample_steps=40,
+                cfg_scale=1.5,
+            )
+
+if __name__ == "__main__":
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+    main()

rf.py

@@ -0,0 +1,112 @@
+import re
+import torch
+from torch import nn
+
+"""
+The class structure is inspired by: https://github.com/cloneofsimo/minRF/blob/main/rf.py
+"""
+
+def parse_int_list(s):
+    if isinstance(s, list): return s
+    ranges = []
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    for p in s.split(','):
+        m = range_re.match(p)
+        if m:
+            ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
+        else:
+            ranges.append(int(p))
+    return ranges
+
+class RF(nn.Module):
+    def __init__(self, train_timestep_sampling: str = "logit_sigmoid", immiscible: bool = False, dts_lambda: float = 0.0):
+        super().__init__()
+        self.train_timestep_sampling = train_timestep_sampling
+
+    def forward(self, unet: nn.Module, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        B = x.size(0)
+        if self.train_timestep_sampling == "logit_sigmoid":
+            t = torch.sigmoid(torch.randn(B, device=x.device))
+        elif self.train_timestep_sampling == "uniform":
+            t = torch.rand(B, device=x.device)
+        else:
+            raise ValueError(f'Unknown train timestep sampling method "{self.train_timestep_sampling}".')
+        t_exp = t.view([B] + [1] * (x.ndim - 1))
+        z1 = torch.randn_like(x)
+        zt = (1 - t_exp) * x + t_exp * z1
+        vtheta = unet(zt, t, **kwargs)
+        loss = ((z1 - x - vtheta) ** 2).mean(dim=list(range(1, x.ndim)))
+        return loss
+
+    def sample(
+        self,
+        unet: nn.Module,
+        z: torch.Tensor,
+        sample_steps: int = 50,
+        cfg_scale: float = 1.0,
+        **kwargs,
+    ):
+        B = z.size(0)
+        dt_value = 1.0 / sample_steps
+        dt = torch.full((B,), dt_value, device=z.device, dtype=z.dtype) \
+                .view([B] + [1] * (z.ndim - 1))
+
+        for i in range(sample_steps, 0, -1):
+            t = torch.full((B,), i / sample_steps, device=z.device, dtype=z.dtype)
+
+            do_guidance = cfg_scale > 1.0
+            do_unconditional_only = cfg_scale == 0.0
+            
+            if do_unconditional_only:
+                vtheta = unet(z, t, class_drop_prob=1.0, **kwargs)
+            else:
+                vtheta = unet(z, t, class_drop_prob=0.0, force_dfs=do_guidance, **kwargs)
+            
+            if do_guidance:
+                vtheta_uncond = unet(z, t, class_drop_prob=1.0, force_dfs=do_guidance, **kwargs)
+                vtheta = vtheta_uncond + cfg_scale * (vtheta - vtheta_uncond)
+
+            z = z - dt * vtheta
+        return z
+    
+    def sample_tread(
+        self, 
+        unet1: nn.Module, 
+        gamma1: float, 
+        unet2: nn.Module, 
+        gamma2: float, 
+        z: torch.Tensor,
+        sample_steps: int = 50,
+        cfg_scale: float = 1.5,
+        **kwargs
+    ):
+        B = z.size(0)
+        dt_value = 1.0 / sample_steps
+        dt = torch.full((B,), dt_value, device=z.device, dtype=z.dtype) \
+                .view([B] + [1] * (z.ndim - 1))
+
+        for i in range(sample_steps, 0, -1):
+            t = torch.full((B,), i / sample_steps, device=z.device, dtype=z.dtype)
+
+            assert cfg_scale > 1.0, "We assume cfg_scale > 1.0 for this TREAD sampling."          
+            vtheta1 = unet1(z, t, class_drop_prob=1.0, force_routing=True, overwrite_selection_ratio=gamma1, **kwargs)
+            vtheta2 = unet2(z, t, class_drop_prob=1.0, force_routing=True, overwrite_selection_ratio=gamma2, **kwargs)
+            vtheta = vtheta2 + cfg_scale * (vtheta1 - vtheta2)
+            z = z - dt * vtheta
+            
+        return z
+
+class LatentRF(RF):
+    def __init__(self, ae: nn.Module, **kwargs):
+        super().__init__(**kwargs)
+        self.ae = ae
+
+    def forward(self, unet: nn.Module, x: torch.Tensor, precomputed: bool = False, **kwargs) -> torch.Tensor:
+        if not precomputed:
+            with torch.no_grad():
+                x = self.ae.encode(x)
+        return super().forward(unet, x, **kwargs)
+
+    def sample(self, unet: nn.Module, z: torch.Tensor, sample_steps: int = 50, return_list: bool = False, **kwargs):
+        latent = super().sample(unet, z, sample_steps=sample_steps, return_list=return_list, **kwargs)
+        return self.ae.decode(latent)