From ed27436f243330e6090db0d1c732a88c2e0abb3c Mon Sep 17 00:00:00 2001 From: Raphael Walker Date: Wed, 26 Feb 2025 17:18:58 +0100 Subject: [PATCH] initial commit --- README.md | 1 + __init__.py | 0 nodes.py | 0 ras.py | 200 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 201 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 ras.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..6539206 --- /dev/null +++ b/README.md @@ -0,0 +1 @@ +# RegionalAdaptiveSampling diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..e69de29 diff --git a/ras.py b/ras.py new file mode 100644 index 0000000..cf5dcf3 --- /dev/null +++ b/ras.py @@ -0,0 +1,200 @@ +from uuid import uuid4 as uuid + +import torch +from torch import nn +from torch import Tensor + +from comfy.ldm.flux.layers import DoubleStreamBlock, SingleStreamBlock +from comfy.ldm.modules.attention import optimized_attention + + +def apply_pe(x: Tensor, pe: Tensor) -> Tensor: + """ + The PE application from flux.math.attention, removed, so that we can cache the keys post-PE + """ + shape = x.shape + dtype = x.dtype + x = x.float().reshape(*x.shape[:-1], -1, 1, 2) + x = (pe[..., 0] * x[..., 0] + pe[..., 1] * x[..., 1]).reshape(*shape).to(dtype) + + return x + + +class RASCache(nn.Module): + """ + KV Cache container. + TODO: at the end of the project, decide if this class can be removed (and the *Wrappers given simply a k_cache and v_cache) + """ + + def __init__(self, hidden_size): + self.hidden_size = hidden_size + self.k_cache: Tensor + self.v_cache: Tensor + + +class RASManager: + """ + Coordinates the live indices, metrics, and model wrapping. + """ + + def __init__(self): + self.caches = {} + self.live_indices: torch.IntTensor + + def wrap_layer(self, layer: DoubleStreamBlock | SingleStreamBlock): + layer_id = uuid() + self.caches[layer_id] = RASCache(layer.hidden_size) + if isinstance(layer, DoubleStreamBlock): + wrapped = DoubleStreamBlockWrapper(layer, self, self.caches[layer_id]) + elif isinstance(layer, SingleStreamBlock): + wrapped = SingleStreamBlockWrapper(layer, self, self.caches[layer_id]) + else: + raise TypeError(f"Can't wrap layer of type {layer.__class__.__name__}") + + return wrapped + + +class DoubleStreamBlockWrapper(nn.Module): + """ + Same as the DoubleStreamBlock, but uses a RASManager and RASCache to do KV caching. + """ + + def __init__( + self, original: DoubleStreamBlock, manager: RASManager, cache: RASCache + ): + self.block = original + self.manager = manager + self.cache = cache + + def forward(self, args, _): + # we get passed the original args + img = args["img"] + txt = args["txt"] + vec = args["vec"] + pe = args["pe"] + attn_mask = args["attention_mask"] + + img_mod1, img_mod2 = self.block.img_mod(vec) + txt_mod1, txt_mod2 = self.block.txt_mod(vec) + + # prepare image for attention + img_modulated = self.block.img_norm1(img) + img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift + img_qkv = self.block.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 + ).permute(2, 0, 3, 1, 4) + img_q, img_k = self.block.img_attn.norm(img_q, img_k, img_v) + + # prepare txt for attention + txt_modulated = self.block.txt_norm1(txt) + txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift + txt_qkv = self.block.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 + ).permute(2, 0, 3, 1, 4) + txt_q, txt_k = self.block.txt_attn.norm(txt_q, txt_k, txt_v) + + # RAS: KV Cache and Attention Call + if self.block.flipped_img_txt: + self.cache.k_cache[..., self.manager.live_indices, :] = apply_pe( + torch.cat((img_k, txt_k), dim=2), pe + ) + self.cache.v_cache[..., self.manager.live_indices, :] = torch.cat( + (img_v, txt_v), dim=2 + ) + attn = optimized_attention( + apply_pe(torch.cat((img_q, txt_q), dim=2), pe), + self.k_cache, + self.v_cache, + img_q.shape[1], + skip_reshape=True, + mask=attn_mask, + ) + img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1] :] + else: + self.cache.k_cache[..., self.manager.live_indices, :] = apply_pe( + torch.cat((txt_k, img_k), dim=2), pe + ) + self.cache.v_cache[..., self.manager.live_indices, :] = torch.cat( + (txt_v, img_v), dim=2 + ) + attn = optimized_attention( + apply_pe(torch.cat((txt_q, img_q), dim=2), pe), + self.k_cache, + self.v_cache, + img_q.shape[1], + skip_reshape=True, + mask=attn_mask, + ) + + 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 + ) + + # calculate the txt bloks + 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 + ) + + if txt.dtype == torch.float16: + txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504) + + return img, txt + + +class SingleStreamBlockWrapper(nn.Module): + """ + Same as the SingleStreamBlock, but uses a RASManager and RASCache to do KV caching. + """ + + def __init__( + self, original: SingleStreamBlock, manager: RASManager, cache: RASCache + ): + self.block = original + self.manager = manager + self.cache = cache + + def forward(self, args, _): + # we get passed the original args + x = args["img"] + vec = args["vec"] + pe = args["pe"] + attn_mask = args["attention_mask"] + + mod, _ = self.block.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], + 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) + + # RAS: KV Cache + self.cache.k_cache[..., self.manager.live_indices, :] = apply_pe(k, pe) + self.cache.v_cache[..., self.manager.live_indices, :] = v + attn = optimized_attention( + apply_pe(q, pe), + self.k_cache, + self.v_cache, + q.shape[1], + skip_reshape=True, + mask=attn_mask, + ) + # 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)) + x += mod.gate * output + if x.dtype == torch.float16: + x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504) + return x