Add gradient checkpointing to the vae

README.md

@@ -66,6 +66,9 @@ shapes, so `weights_name="my_finetune.safetensors"` loads any Primal VAE checkpo
 `pixel_shuffle` by `compression / 2`, so `compression=8` is a 4x4 patch. `vae.token_channels(8)` and
 `vae.image_multiple` give you the numbers above.
 
+To backprop through `decode` (a pixel-space loss on a diffusion model's prediction), `vae.enable_gradient_checkpointing()`
+checkpoints every residual block of both U-Nets; inference under `torch.no_grad()` is unaffected.
+
 ## Architecture
 
 Both models are the same network, they only differ in output channels.

modeling_primal_vae.py

@@ -15,6 +15,7 @@ from typing import List, Optional
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.utils.checkpoint
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
 from huggingface_hub import hf_hub_download
@@ -66,6 +67,14 @@ class UStack(nn.Module):
         self.up_blocks = nn.ModuleList([nn.ModuleList([ResBlock(widths[i], DILATION_CYCLE[j % len(DILATION_CYCLE)]) for j in range(up_depths[i])]) for i in range(L - 1)])
         self.norm_out = ChannelRMSNorm(widths[0])
         self.conv_out = nn.Conv2d(widths[0], out_ch, 3, padding=1)
+        # set by ModelMixin.enable_gradient_checkpointing(); per ResBlock, only while grad is enabled
+        self.gradient_checkpointing = False
+        self._gradient_checkpointing_func = lambda blk, x: torch.utils.checkpoint.checkpoint(blk, x, use_reentrant=False)
+
+    def _block(self, blk: ResBlock, x: torch.Tensor) -> torch.Tensor:
+        if self.gradient_checkpointing and torch.is_grad_enabled():
+            return self._gradient_checkpointing_func(blk, x)
+        return blk(x)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv_in(x)
@@ -73,14 +82,14 @@ class UStack(nn.Module):
         L = len(self.down_blocks)
         for i in range(L):
             for blk in self.down_blocks[i]:
-                x = blk(x)
+                x = self._block(blk, x)
             if i < L - 1:
                 skips.append(x)
                 x = self.downs[i](F.pixel_unshuffle(x, 2))
         for i in reversed(range(L - 1)):
             x = F.pixel_shuffle(self.ups[i](x), 2) + skips[i]
             for blk in self.up_blocks[i]:
-                x = blk(x)
+                x = self._block(blk, x)
         return self.conv_out(F.silu(self.norm_out(x)))
 
 
@@ -88,6 +97,8 @@ class PrimalVAE(ModelMixin, ConfigMixin):
     """Primal VAE: `encode` -> [B, latent_channels, H/compression, W/compression] (~N(0,1)), `decode` back to
     [B, 3, H, W] in [-1, 1]. No posterior, no scaling/shift factor."""
 
+    _supports_gradient_checkpointing = True
+
     @register_to_config
     def __init__(
         self,