diff --git a/comfy/ldm/lightricks/model.py b/comfy/ldm/lightricks/model.py
index ad9a7daea..aa2ea62b1 100644
--- a/comfy/ldm/lightricks/model.py
+++ b/comfy/ldm/lightricks/model.py
@@ -1,5 +1,6 @@
 import torch
 from torch import nn
+import comfy.patcher_extension
 import comfy.ldm.modules.attention
 import comfy.ldm.common_dit
 from einops import rearrange
@@ -420,6 +421,13 @@ class LTXVModel(torch.nn.Module):
         self.patchifier = SymmetricPatchifier(1)
 
     def forward(self, x, timestep, context, attention_mask, frame_rate=25, transformer_options={}, keyframe_idxs=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, transformer_options)
+        ).execute(x, timestep, context, attention_mask, frame_rate, transformer_options, keyframe_idxs, **kwargs)
+
+    def _forward(self, x, timestep, context, attention_mask, frame_rate=25, transformer_options={}, keyframe_idxs=None, **kwargs):
         patches_replace = transformer_options.get("patches_replace", {})
 
         orig_shape = list(x.shape)
diff --git a/comfy_extras/nodes_easycache.py b/comfy_extras/nodes_easycache.py
index 4139b1f4b..45e2e6838 100644
--- a/comfy_extras/nodes_easycache.py
+++ b/comfy_extras/nodes_easycache.py
@@ -15,6 +15,8 @@ def easycache_forward_wrapper(executor, *args, **kwargs):
     transformer_options: dict[str] = args[-1]
     if not isinstance(transformer_options, dict):
         transformer_options = kwargs.get("transformer_options")
+        if not transformer_options:
+            transformer_options = args[-2]
     easycache: EasyCacheHolder = transformer_options["easycache"]
     sigmas = transformer_options["sigmas"]
     uuids = transformer_options["uuids"]
@@ -186,6 +188,7 @@ class EasyCacheHolder:
         self.approx_output_change_rates = []
         self.total_steps_skipped = 0
         # how to deal with mismatched dims
+        self.allow_mismatch = True
         self.cut_from_start = True
 
     def is_past_end_timestep(self, timestep: float) -> bool:
@@ -230,7 +233,9 @@ class EasyCacheHolder:
         batch_offset = x.shape[0] // len(uuids)
         for i, uuid in enumerate(uuids):
             # 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:
+            if x.shape[1:] != self.uuid_cache_diffs[uuid].shape[1:]:
+                if not self.allow_mismatch:
+                    raise ValueError(f"Cached dims {self.uuid_cache_diffs[uuid].shape} don't match x dims {x.shape} - this is no good")
                 slicing = []
                 skip_this_dim = True
                 for dim_u, dim_x in zip(self.uuid_cache_diffs[uuid].shape, x.shape):
@@ -251,16 +256,20 @@ class EasyCacheHolder:
 
     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:
+        if output.shape[1:] != x.shape[1:]:
+            if not self.allow_mismatch:
+                raise ValueError(f"Output dims {output.shape} don't match x dims {x.shape} - this is no good")
             slicing = []
+            skip_dim = True
             for dim_o, dim_x in zip(output.shape, x.shape):
-                if dim_o != dim_x:
+                if not skip_dim and 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))
+                skip_dim = False
             x = x[slicing]
         diff = output - x
         batch_offset = diff.shape[0] // len(uuids)