Add original_shape parameter to transformer patch extra_options.

im-stuff · Jun 21, 2023 · 9fccf4a · 9fccf4a
1 parent 852cf4d
commit 9fccf4a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
@@ -529,6 +529,8 @@ def _forward(self, x, context=None, transformer_options={}):
             extra_options["transformer_index"] = transformer_options["current_index"]
         if "block_index" in transformer_options:
             extra_options["block_index"] = transformer_options["block_index"]
+        if "original_shape" in transformer_options:
+            extra_options["original_shape"] = transformer_options["original_shape"]
         if "patches" in transformer_options:
             transformer_patches = transformer_options["patches"]
         else: