From d43390a199ee104999a1d37632af99316b2b9b85 Mon Sep 17 00:00:00 2001 From: Raphael Walker Date: Fri, 28 Feb 2025 11:02:22 +0100 Subject: [PATCH] working properly for flux --- __init__.py | 10 ++-- ras.py | 142 ++++++++++++++++++++++++++++++++++++---------------- 2 files changed, 106 insertions(+), 46 deletions(-) diff --git a/__init__.py b/__init__.py index e8b18a5..2710758 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,5 @@ from .ras import RASConfig, RASManager +from comfy.model_patcher import ModelPatcher class RegionalAdaptiveSampling: @@ -18,12 +19,15 @@ class RegionalAdaptiveSampling: FUNCTION = "apply_ras" CATEGORY = "ras" - def apply_ras(self, model, sample_ratio): + def apply_ras(self, model: ModelPatcher, sample_ratio: float): model = model.clone() + # unpatch the model + # this makes sure that we're wrapping the model "in a pure state" + # the model will repatch itself later + model.unpatch_model() config = RASConfig(sample_ratio=sample_ratio) manager = RASManager(config) - # todo: improve this, because this is a godawful way to go about it - manager.wrap_model(model.model.diffusion_model) + manager.wrap_model(model) return (model,) diff --git a/ras.py b/ras.py index c63f171..2a45eb7 100644 --- a/ras.py +++ b/ras.py @@ -11,6 +11,7 @@ from comfy.ldm.flux.model import Flux from comfy.ldm.hunyuan_video.model import HunyuanVideo from comfy.ldm.flux.layers import DoubleStreamBlock, SingleStreamBlock, LastLayer from comfy.ldm.modules.attention import optimized_attention +from comfy.model_patcher import ModelPatcher def apply_pe(x: Tensor, pe: Tensor) -> Tensor: @@ -25,6 +26,11 @@ def apply_pe(x: Tensor, pe: Tensor) -> Tensor: return x +def take_attributes_from(source, target, keys): + for x in keys: + setattr(target, x, getattr(source, x)) + + @dataclass class RASConfig: start_step: int = 4 @@ -58,6 +64,12 @@ class RASManager: def wrap_layer( self, layer: DoubleStreamBlock | SingleStreamBlock | LastLayer, first=False ): + if isinstance( + layer, + (DoubleStreamBlockWrapper, SingleStreamBlockWrapper, LastLayerWrapper), + ): + raise TypeError("Old wrapping wasn't removed!") + if isinstance(layer, DoubleStreamBlock): wrapped = DoubleStreamBlockWrapper(layer, self, first) elif isinstance(layer, SingleStreamBlock): @@ -68,7 +80,8 @@ class RASManager: raise TypeError(f"Can't wrap layer of type {layer.__class__.__name__}") return wrapped - def wrap_model(self, model: Flux | HunyuanVideo): + def wrap_model(self, patcher: ModelPatcher): + model = patcher.model.diffusion_model self.model = model if isinstance(model, Flux): self.patch_size = [model.patch_size, model.patch_size] @@ -80,11 +93,14 @@ class RASManager: for i, v in enumerate(model.double_blocks): if i == 0: self.flipped_img_txt = v.flipped_img_txt - model.double_blocks[i] = self.wrap_layer(v, first=True) + layer = self.wrap_layer(v, first=True) else: - model.double_blocks[i] = self.wrap_layer(v) + layer = self.wrap_layer(v) + patcher.add_object_patch(f"diffusion_model.double_blocks.{i}", layer) for i, v in enumerate(model.single_blocks): - model.single_blocks[i] = self.wrap_layer(v) + patcher.add_object_patch( + f"diffusion_model.single_blocks.{i}", self.wrap_layer(v) + ) # wrap the forward_orig method, to be able to get the timestep forward_orig = model.forward_orig @@ -101,10 +117,14 @@ class RASManager: self.drop_count = None return forward_orig(*args, **kwargs) - model.forward_orig = MethodType(new_forward, model) + patcher.add_object_patch( + "diffusion_model.forward_orig", MethodType(new_forward, model) + ) # wrap the last_layer, to be able to read the output, and replace the missing (cached) tokens - model.final_layer = self.wrap_layer(model.final_layer) + patcher.add_object_patch( + "diffusion_model.final_layer", self.wrap_layer(model.final_layer) + ) @staticmethod def timestep_from_sigmas(sigmas: Tensor, sample_sigmas: Tensor): @@ -183,19 +203,20 @@ class RASManager: ) -class LastLayerWrapper(nn.Module): +class LastLayerWrapper(LastLayer): """ Same as the LastLayer, but reports its output to a manager. """ def __init__(self, original: LastLayer, manager: RASManager): - super().__init__() - self.original = original + nn.Module.__init__(self) + take_attributes_from( + original, self, ["norm_final", "linear", "adaLN_modulation"] + ) self.manager = manager def forward(self, x, vec) -> Tensor: - # call the original first - output = self.original(x, vec) + output = super().forward(x, vec) # if we were fully hydrating, fill the cache if self.manager.live_indices() is None or self.manager.cached_output is None: self.manager.cached_output = output.clone() @@ -205,14 +226,32 @@ class LastLayerWrapper(nn.Module): return self.manager.cached_output -class DoubleStreamBlockWrapper(nn.Module): +class DoubleStreamBlockWrapper(DoubleStreamBlock): """ Same as the DoubleStreamBlock, but uses a RASManager and RASCache to do KV caching. """ def __init__(self, original: DoubleStreamBlock, manager: RASManager, first=False): - super().__init__() - self.block = original + nn.Module.__init__(self) + take_attributes_from( + original, + self, + [ + "num_heads", + "hidden_size", + "img_mod", + "img_norm1", + "img_attn", + "img_norm2", + "img_mlp", + "txt_mod", + "txt_norm1", + "txt_attn", + "txt_norm2", + "txt_mlp", + "flipped_img_txt", + ], + ) self.manager = manager self.k_cache: torch.Tensor self.v_cache: torch.Tensor @@ -231,26 +270,26 @@ class DoubleStreamBlockWrapper(nn.Module): img = img[..., img_idx, :] txt = txt[..., txt_idx, :] - img_mod1, img_mod2 = self.block.img_mod(vec) - txt_mod1, txt_mod2 = self.block.txt_mod(vec) + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) # prepare image for attention - img_modulated = self.block.img_norm1(img) + img_modulated = self.img_norm1(img) img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift - img_qkv = self.block.img_attn.qkv(img_modulated) + img_qkv = self.img_attn.qkv(img_modulated) img_q, img_k, img_v = img_qkv.view( - img_qkv.shape[0], img_qkv.shape[1], 3, self.block.num_heads, -1 + img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1 ).permute(2, 0, 3, 1, 4) - img_q, img_k = self.block.img_attn.norm(img_q, img_k, img_v) + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) # prepare txt for attention - txt_modulated = self.block.txt_norm1(txt) + txt_modulated = self.txt_norm1(txt) txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift - txt_qkv = self.block.txt_attn.qkv(txt_modulated) + txt_qkv = self.txt_attn.qkv(txt_modulated) txt_q, txt_k, txt_v = txt_qkv.view( - txt_qkv.shape[0], txt_qkv.shape[1], 3, self.block.num_heads, -1 + txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1 ).permute(2, 0, 3, 1, 4) - txt_q, txt_k = self.block.txt_attn.norm(txt_q, txt_k, txt_v) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) # RAS: KV Cache and Attention Call # select part of the PE @@ -258,7 +297,7 @@ class DoubleStreamBlockWrapper(nn.Module): pe = pe[:, :, idx] # create queries, keys, and values - if self.block.flipped_img_txt: + if self.flipped_img_txt: queries = apply_pe(torch.cat((img_q, txt_q), dim=2), pe) keys = apply_pe(torch.cat((img_k, txt_k), dim=2), pe) values = torch.cat((img_v, txt_v), dim=2) @@ -283,22 +322,22 @@ class DoubleStreamBlockWrapper(nn.Module): skip_reshape=True, mask=attn_mask, ) - if self.block.flipped_img_txt: + if self.flipped_img_txt: img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1] :] else: txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] # End of RAS code # calculate the img blocks - img = img + img_mod1.gate * self.block.img_attn.proj(img_attn) - img = img + img_mod2.gate * self.block.img_mlp( - (1 + img_mod2.scale) * self.block.img_norm2(img) + img_mod2.shift + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp( + (1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift ) # calculate the txt blocks - txt += txt_mod1.gate * self.block.txt_attn.proj(txt_attn) - txt += txt_mod2.gate * self.block.txt_mlp( - (1 + txt_mod2.scale) * self.block.txt_norm2(txt) + txt_mod2.shift + txt += txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt += txt_mod2.gate * self.txt_mlp( + (1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift ) if txt.dtype == torch.float16: @@ -307,31 +346,48 @@ class DoubleStreamBlockWrapper(nn.Module): return img, txt -class SingleStreamBlockWrapper(nn.Module): +class SingleStreamBlockWrapper(SingleStreamBlock): """ Same as the SingleStreamBlock, but uses a RASManager and RASCache to do KV caching. """ def __init__(self, original: SingleStreamBlock, manager: RASManager): - super().__init__() - self.block = original + nn.Module.__init__(self) + # steal all the attributes from the SingleStreamBlock + take_attributes_from( + original, + self, + [ + "hidden_dim", + "num_heads", + "scale", + "mlp_hidden_dim", + "linear1", + "linear2", + "norm", + "hidden_size", + "pre_norm", + "mlp_act", + "modulation", + ], + ) self.manager = manager self.k_cache: torch.Tensor self.v_cache: torch.Tensor def forward(self, x, vec, pe, attn_mask=None): idx = self.manager.live_indices() - mod, _ = self.block.modulation(vec) + mod, _ = self.modulation(vec) qkv, mlp = torch.split( - self.block.linear1((1 + mod.scale) * self.block.pre_norm(x) + mod.shift), - [3 * self.block.hidden_size, self.block.mlp_hidden_dim], + self.linear1((1 + mod.scale) * self.pre_norm(x) + mod.shift), + [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1, ) - q, k, v = qkv.view( - qkv.shape[0], qkv.shape[1], 3, self.block.num_heads, -1 - ).permute(2, 0, 3, 1, 4) - q, k = self.block.norm(q, k, v) + q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute( + 2, 0, 3, 1, 4 + ) + q, k = self.norm(q, k, v) # RAS: KV Cache if idx is not None: @@ -356,7 +412,7 @@ class SingleStreamBlockWrapper(nn.Module): ) # End of RAS code # compute activation in mlp stream, cat again and run second linear layer - output = self.block.linear2(torch.cat((attn, self.block.mlp_act(mlp)), 2)) + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) x += mod.gate * output if x.dtype == torch.float16: x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)