From 4bc3e217d6e4ea8bb6e4cb16245fa8b5f97cbada Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Thu, 21 Aug 2025 21:46:58 -0700
Subject: [PATCH] Make EasyCache not crash with Cosmos Predict ImagToVideo
 latents, but does not work well at all

---
 comfy/ldm/cosmos/model.py       | 38 +++++++++++++++++++++++++++++++++
 comfy_extras/nodes_easycache.py | 33 +++++++++++++++++++++++++++-
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/comfy/ldm/cosmos/model.py b/comfy/ldm/cosmos/model.py
index 4836e0b69..53698b758 100644
--- a/comfy/ldm/cosmos/model.py
+++ b/comfy/ldm/cosmos/model.py
@@ -27,6 +27,8 @@ from torchvision import transforms
 from enum import Enum
 import logging
 
+import comfy.patcher_extension
+
 from .blocks import (
     FinalLayer,
     GeneralDITTransformerBlock,
@@ -435,6 +437,42 @@ class GeneralDIT(nn.Module):
         latent_condition_sigma: Optional[torch.Tensor] = None,
         condition_video_augment_sigma: Optional[torch.Tensor] = None,
         **kwargs,
+    ):
+        return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+            self._forward,
+            self,
+            comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, kwargs.get("transformer_options", {}))
+        ).execute(x,
+                timesteps,
+                context,
+                attention_mask,
+                fps,
+                image_size,
+                padding_mask,
+                scalar_feature,
+                data_type,
+                latent_condition,
+                latent_condition_sigma,
+                condition_video_augment_sigma,
+                **kwargs)
+
+    def _forward(
+        self,
+        x: torch.Tensor,
+        timesteps: torch.Tensor,
+        context: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        # crossattn_emb: torch.Tensor,
+        # crossattn_mask: Optional[torch.Tensor] = None,
+        fps: Optional[torch.Tensor] = None,
+        image_size: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        scalar_feature: Optional[torch.Tensor] = None,
+        data_type: Optional[DataType] = DataType.VIDEO,
+        latent_condition: Optional[torch.Tensor] = None,
+        latent_condition_sigma: Optional[torch.Tensor] = None,
+        condition_video_augment_sigma: Optional[torch.Tensor] = None,
+        **kwargs,
     ):
         """
         Args:
diff --git a/comfy_extras/nodes_easycache.py b/comfy_extras/nodes_easycache.py
index 80e7e8fc3..4139b1f4b 100644
--- a/comfy_extras/nodes_easycache.py
+++ b/comfy_extras/nodes_easycache.py
@@ -185,6 +185,8 @@ class EasyCacheHolder:
         self.output_change_rates = []
         self.approx_output_change_rates = []
         self.total_steps_skipped = 0
+        # how to deal with mismatched dims
+        self.cut_from_start = True
 
     def is_past_end_timestep(self, timestep: float) -> bool:
         return not (timestep[0] > self.end_t).item()
@@ -227,10 +229,39 @@ class EasyCacheHolder:
             self.total_steps_skipped += 1
         batch_offset = x.shape[0] // len(uuids)
         for i, uuid in enumerate(uuids):
-            x[i*batch_offset:(i+1)*batch_offset, ...] += self.uuid_cache_diffs[uuid].to(x.device)
+            # if cached dims don't match x dims, cut off excess and hope for the best (cosmos world2video)
+            if x.shape != self.uuid_cache_diffs[uuid].shape:
+                slicing = []
+                skip_this_dim = True
+                for dim_u, dim_x in zip(self.uuid_cache_diffs[uuid].shape, x.shape):
+                    if skip_this_dim:
+                        skip_this_dim = False
+                        continue
+                    if dim_u != dim_x:
+                        if self.cut_from_start:
+                            slicing.append(slice(dim_x-dim_u, None))
+                        else:
+                            slicing.append(slice(None, dim_u))
+                    else:
+                        slicing.append(slice(None))
+                slicing = [slice(i*batch_offset,(i+1)*batch_offset)] + slicing
+                x = x[slicing]
+            x += self.uuid_cache_diffs[uuid].to(x.device)
         return x
 
     def update_cache_diff(self, output: torch.Tensor, x: torch.Tensor, uuids: list[UUID]):
+        # if output dims don't match x dims, cut off excess and hope for the best (cosmos world2video)
+        if output.shape != x.shape:
+            slicing = []
+            for dim_o, dim_x in zip(output.shape, x.shape):
+                if dim_o != dim_x:
+                    if self.cut_from_start:
+                        slicing.append(slice(dim_x-dim_o, None))
+                    else:
+                        slicing.append(slice(None, dim_o))
+                else:
+                    slicing.append(slice(None))
+            x = x[slicing]
         diff = output - x
         batch_offset = diff.shape[0] // len(uuids)
         for i, uuid in enumerate(uuids):