Add EasyCache support to Chroma

This commit is contained in:
Jedrzej Kosinski 2025-08-21 20:32:40 -07:00
parent d1ea8c0a58
commit 84609ca160

View File

@ -5,6 +5,7 @@ from dataclasses import dataclass
import torch import torch
from torch import Tensor, nn from torch import Tensor, nn
from einops import rearrange, repeat from einops import rearrange, repeat
import comfy.patcher_extension
import comfy.ldm.common_dit import comfy.ldm.common_dit
from comfy.ldm.flux.layers import ( from comfy.ldm.flux.layers import (
@ -253,6 +254,13 @@ class Chroma(nn.Module):
return img return img
def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs): def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **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, guidance, control, transformer_options, **kwargs)
def _forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
bs, c, h, w = x.shape bs, c, h, w = x.shape
x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size)) x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))