nodes/Slickytail-ComfyUI-RegionalAdaptiveSampling
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

working properly for flux

Browse Source
This commit is contained in:
Raphael Walker
2025-02-28 11:02:22 +01:00
parent b2466a6454
commit d43390a199
2 changed files with 106 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+7 -3
View File
@@ -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,)
ras.py
+99 -43
View File
@@ -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)