xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-26 21:37:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added EasyCache support to LTXV (not very good, but does not crash)

Browse Source
This commit is contained in:
Jedrzej Kosinski 2025-08-21 23:55:07 -07:00
parent d60e7a07cb
commit e1b94339b7
2 changed files with 20 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
comfy/ldm/lightricks/model.py
View File

@ -1,5 +1,6 @@
import torch import torch
from torch import nn from torch import nn
import comfy.patcher_extension
import comfy.ldm.modules.attention import comfy.ldm.modules.attention
import comfy.ldm.common_dit import comfy.ldm.common_dit
from einops import rearrange from einops import rearrange
@ -420,6 +421,13 @@ class LTXVModel(torch.nn.Module):
self.patchifier = SymmetricPatchifier(1) self.patchifier = SymmetricPatchifier(1)
def forward(self, x, timestep, context, attention_mask, frame_rate=25, transformer_options={}, keyframe_idxs=None, **kwargs): 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", {}) patches_replace = transformer_options.get("patches_replace", {})
orig_shape = list(x.shape) orig_shape = list(x.shape)

15
comfy_extras/nodes_easycache.py
View File

@ -15,6 +15,8 @@ def easycache_forward_wrapper(executor, *args, **kwargs):
transformer_options: dict[str] = args[-1] transformer_options: dict[str] = args[-1]
if not isinstance(transformer_options, dict): if not isinstance(transformer_options, dict):
transformer_options = kwargs.get("transformer_options") transformer_options = kwargs.get("transformer_options")
if not transformer_options:
transformer_options = args[-2]
easycache: EasyCacheHolder = transformer_options["easycache"] easycache: EasyCacheHolder = transformer_options["easycache"]
sigmas = transformer_options["sigmas"] sigmas = transformer_options["sigmas"]
uuids = transformer_options["uuids"] uuids = transformer_options["uuids"]
@ -186,6 +188,7 @@ class EasyCacheHolder:
self.approx_output_change_rates = [] self.approx_output_change_rates = []
self.total_steps_skipped = 0 self.total_steps_skipped = 0
# how to deal with mismatched dims # how to deal with mismatched dims
self.allow_mismatch = True
self.cut_from_start = True self.cut_from_start = True
def is_past_end_timestep(self, timestep: float) -> bool: def is_past_end_timestep(self, timestep: float) -> bool:
@ -230,7 +233,9 @@ class EasyCacheHolder:
batch_offset = x.shape[0] // len(uuids) batch_offset = x.shape[0] // len(uuids)
for i, uuid in enumerate(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 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 = [] slicing = []
skip_this_dim = True skip_this_dim = True
for dim_u, dim_x in zip(self.uuid_cache_diffs[uuid].shape, x.shape): 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]): 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 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 = [] slicing = []
skip_dim = True
for dim_o, dim_x in zip(output.shape, x.shape): 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: if self.cut_from_start:
slicing.append(slice(dim_x-dim_o, None)) slicing.append(slice(dim_x-dim_o, None))
else: else:
slicing.append(slice(None, dim_o)) slicing.append(slice(None, dim_o))
else: else:
slicing.append(slice(None)) slicing.append(slice(None))
skip_dim = False
x = x[slicing] x = x[slicing]
diff = output - x diff = output - x
batch_offset = diff.shape[0] // len(uuids) batch_offset = diff.shape[0] // len(uuids)