05185634

context_unet_backup.py

@@ -127,6 +127,10 @@ class TimestepBlock(ABC, nn.Module):
         """
 
 class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    def __init__(self, *arg, use_checkpoint=False):
+        super().__init__(*args)
+        self.use_checkpoint = use_checkpoint
+
     def forward(self, x, emb, encoder_out=None):
         for layer in self:
             if isinstance(layer, TimestepBlock):
@@ -134,7 +138,11 @@ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
             elif isinstance(layer, AttentionBlock):
                 x = layer(x, encoder_out)
             else:
-                x = layer(x)
+                if self.use_checkpoint:
+                    print(f"TimestepEmbedSequential checkpoint working")
+                    x = checkpoint.checkpoint(layer, x)
+                else:
+                    x = layer(x)
         return x
 
 class ResBlock(TimestepBlock):