From 6c0fe7c491d68a3c149cdcb35057addbf8d17c24 Mon Sep 17 00:00:00 2001
From: Raphael Walker <slickytail.mc@gmail.com>
Date: Thu, 3 Apr 2025 15:50:33 +0200
Subject: [PATCH] Add activations_shape info in UNet models

---
 comfy/ldm/modules/attention.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
index ede506463..56f7e5844 100644
--- a/comfy/ldm/modules/attention.py
+++ b/comfy/ldm/modules/attention.py
@@ -847,6 +847,7 @@ class SpatialTransformer(nn.Module):
         if not isinstance(context, list):
             context = [context] * len(self.transformer_blocks)
         b, c, h, w = x.shape
+        transformer_options["activations_shape"] = x.shape
         x_in = x
         x = self.norm(x)
         if not self.use_linear:
@@ -962,6 +963,7 @@ class SpatialVideoTransformer(SpatialTransformer):
         transformer_options={}
     ) -> torch.Tensor:
         _, _, h, w = x.shape
+        transformer_options["activations_shape"] = x.shape
         x_in = x
         spatial_context = None
         if exists(context):