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

Experimental support for Wan (VACE coming soon...)

Browse Source
This commit is contained in:
Yael Walker
2025-09-23 12:18:34 +02:00
parent 08f21c68f7
commit af455e091f
1 changed files with 435 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ras.py
+435 -37
View File
@@ -10,8 +10,20 @@ from einops import rearrange
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.wan.model import (
WanModel,
VaceWanModel,
CameraWanModel,
WanModel_S2V,
HumoWanModel,
WanAttentionBlock,
VaceWanAttentionBlock,
Head,
)
from comfy.ldm.flux.math import apply_rope1
from comfy.ldm.modules.attention import optimized_attention
from comfy.model_patcher import ModelPatcher
import comfy.model_management
def apply_pe(x: Tensor, pe: Tensor) -> Tensor:
@@ -61,23 +73,31 @@ class RASManager:
self.config.high_ratio >= 0 and self.config.high_ratio <= 1
), "High ratio should be in the range of [0, 1]"
def wrap_layer(
self,
layer: DoubleStreamBlock | SingleStreamBlock | LastLayer,
first_or_last=False,
):
def wrap_layer(self, layer, first: bool = False, last: bool = False):
if isinstance(
layer,
(DoubleStreamBlockWrapper, SingleStreamBlockWrapper, LastLayerWrapper),
(
DoubleStreamBlockWrapper,
SingleStreamBlockWrapper,
LastLayerWrapper,
WanAttentionBlockWrapper,
VaceWanAttentionBlockWrapper,
),
):
raise TypeError("Old wrapping wasn't removed!")
if isinstance(layer, DoubleStreamBlock):
wrapped = DoubleStreamBlockWrapper(layer, self, first_or_last)
wrapped = DoubleStreamBlockWrapper(layer, self, first)
elif isinstance(layer, SingleStreamBlock):
wrapped = SingleStreamBlockWrapper(layer, self, first_or_last)
wrapped = SingleStreamBlockWrapper(layer, self, last)
elif isinstance(layer, LastLayer):
wrapped = LastLayerWrapper(layer, self)
elif isinstance(layer, WanAttentionBlock):
wrapped = WanAttentionBlockWrapper(layer, self, first, last)
elif isinstance(layer, VaceWanAttentionBlock):
wrapped = VaceWanAttentionBlockWrapper(layer, self, first, last)
elif isinstance(layer, Head):
wrapped = HeadWrapper(layer, self)
else:
raise TypeError(f"Can't wrap layer of type {layer.__class__.__name__}")
return wrapped
@@ -89,29 +109,52 @@ class RASManager:
self.patch_size = [model.patch_size, model.patch_size]
elif isinstance(model, HunyuanVideo):
self.patch_size = model.patch_size
elif isinstance(
model, (WanModel, VaceWanModel, CameraWanModel, WanModel_S2V, HumoWanModel)
):
self.patch_size = model.patch_size
else:
raise TypeError(f"Can't wrap model of type {model.__class__.__name__}")
# wrap the single and double blocks to have caching
for i, v in enumerate(model.double_blocks):
# first block has the special responsibility of removing tokens
if i == 0:
self.flipped_img_txt = v.flipped_img_txt
layer = self.wrap_layer(v, first_or_last=True)
else:
layer = self.wrap_layer(v)
patcher.add_object_patch(f"diffusion_model.double_blocks.{i}", layer)
for i, v in enumerate(model.single_blocks):
# last block will put them back
patcher.add_object_patch(
f"diffusion_model.single_blocks.{i}",
self.wrap_layer(v, i == (len(model.single_blocks) - 1)),
)
# Handle different model architectures
if isinstance(model, (Flux, HunyuanVideo)):
# wrap the single and double blocks to have caching
for i, v in enumerate(model.double_blocks):
# first block has the special responsibility of removing tokens
if i == 0:
self.flipped_img_txt = v.flipped_img_txt
layer = self.wrap_layer(v, first=True)
else:
layer = self.wrap_layer(v)
patcher.add_object_patch(f"diffusion_model.double_blocks.{i}", layer)
for i, v in enumerate(model.single_blocks):
# last block will put them back
patcher.add_object_patch(
f"diffusion_model.single_blocks.{i}",
self.wrap_layer(v, last=(i == (len(model.single_blocks) - 1))),
)
elif isinstance(
model, (WanModel, VaceWanModel, CameraWanModel, WanModel_S2V, HumoWanModel)
):
# Wan models have a different structure with just 'blocks'
self.flipped_img_txt = False # Wan doesn't use the flipped pattern
for i, v in enumerate(model.blocks):
# first block has the special responsibility of removing tokens
# last block will put them back
layer = self.wrap_layer(
v, first=(i == 0), last=(i == (len(model.blocks) - 1))
)
patcher.add_object_patch(f"diffusion_model.blocks.{i}", layer)
# wrap the forward_orig method, to be able to get the timestep
forward_orig = model.forward_orig
def new_forward(_self, *args, **kwargs):
transformer_options = args[-1]
# Get transformer_options from kwargs for Wan models, or from args for Flux/Hunyuan
if "transformer_options" in kwargs:
transformer_options = kwargs["transformer_options"]
else:
transformer_options = args[-1]
self.timestep = self.timestep_from_sigmas(
transformer_options["sigmas"], transformer_options["sample_sigmas"]
)
@@ -126,9 +169,17 @@ class RASManager:
"diffusion_model.forward_orig", MethodType(new_forward, model)
)
# wrap the last_layer, to be able to read the output and calculate the metric
patcher.add_object_patch(
"diffusion_model.final_layer", self.wrap_layer(model.final_layer)
)
if isinstance(model, (Flux, HunyuanVideo)):
patcher.add_object_patch(
"diffusion_model.final_layer", self.wrap_layer(model.final_layer)
)
elif isinstance(
model, (WanModel, VaceWanModel, CameraWanModel, WanModel_S2V, HumoWanModel)
):
# Wan models use 'head' instead of 'final_layer'
patcher.add_object_patch(
"diffusion_model.head", self.wrap_layer(model.head)
)
@staticmethod
def timestep_from_sigmas(sigmas: Tensor, sample_sigmas: Tensor):
@@ -141,12 +192,14 @@ class RASManager:
def skip_ratio(self, timestep: int) -> float:
if timestep < self.config.warmup_steps:
return 0
if self.config.hydrate_every:
if (
1 + timestep - self.config.warmup_steps
) % self.config.hydrate_every == 0:
return 0
return 1.0 - self.config.sample_ratio
result = 1.0 - self.config.sample_ratio
return result
def select_indices(self, diff: Tensor, timestep: int):
if isinstance(self.model, Flux):
@@ -158,8 +211,19 @@ class RASManager:
pw=self.patch_size[1],
)
metric = torch.std(metric, dim=-1).mean((-1, -2))
elif isinstance(self.model, HunyuanVideo):
elif isinstance(
self.model,
(
HunyuanVideo,
WanModel,
VaceWanModel,
CameraWanModel,
WanModel_S2V,
HumoWanModel,
),
):
# b (h w) (c ph pw) = model_out.shape
# Both HunyuanVideo and Wan models are video models with 3D patches
metric = rearrange(
diff,
"b s (c pt ph pw ) -> b s pt ph pw c",
@@ -170,8 +234,9 @@ class RASManager:
metric = torch.std(metric, dim=-1).mean((-1, -2, -3))
else:
raise TypeError("Unknown latent type!")
# hmm, how do we deal with batch size != 1 here?
# that might be a problem
# for batch size > 1, we pick separate indices per batch
# for now, JUST FOR TESTING, we'll merge all the batches and use the indices that are the most relevant for all batches
metric = metric.mean(dim=0)
metric = metric.flatten()
if self.drop_count is None:
self.drop_count = torch.zeros(
@@ -191,23 +256,26 @@ class RASManager:
cache_indices = torch.cat([indices[:low_bar], indices[-high_bar:]])
self.live_img_indices = indices[low_bar:-high_bar]
self.drop_count[cache_indices] += 1
# TODO: for now we keep all txt tokens
# in the future, we can probably do something like randomly keep a fraction of them
self.live_txt_indices = torch.arange(
0, self.n_txt, dtype=torch.int, device=diff.device
)
if self.n_txt > 0:
self.live_txt_indices = torch.arange(
0, self.n_txt, dtype=torch.int, device=diff.device
)
def live_indices(self):
if self.live_img_indices is None or self.live_txt_indices is None:
return None
return self.live_img_indices
if self.flipped_img_txt:
return torch.cat(
result = torch.cat(
(self.live_img_indices, self.live_txt_indices + self.n_img)
)
else:
return torch.cat(
result = torch.cat(
(self.live_txt_indices, self.live_img_indices + self.n_txt)
)
return result
class DoubleStreamBlockWrapper(DoubleStreamBlock):
@@ -428,3 +496,333 @@ class LastLayerWrapper(LastLayer):
output = super().forward(x, vec)
self.manager.select_indices(output, self.manager.timestep)
return output
class WanAttentionBlockWrapper(WanAttentionBlock):
"""
Same as the WanAttentionBlock, but uses a RASManager and RASCache to do KV caching.
"""
def __init__(self, original: WanAttentionBlock, manager: RASManager, first, last):
nn.Module.__init__(self)
take_attributes_from(
original,
self,
[
"dim",
"ffn_dim",
"num_heads",
"window_size",
"qk_norm",
"cross_attn_norm",
"eps",
"norm1",
"self_attn",
"norm3",
"cross_attn",
"norm2",
"ffn",
"modulation",
],
)
self.manager = manager
self.k_cache: torch.Tensor
self.v_cache: torch.Tensor
self.first = first
self.last = last
def forward(
self, x, e, freqs, context, context_img_len=257, transformer_options={}
):
# RAS: if this is the first block, then we should drop some tokens
idx = self.manager.live_indices()
if self.first:
# we just have img tokens
self.manager.n_txt = 0
self.manager.n_img = x.shape[1]
if idx is not None:
x = x[..., idx, :]
# Modulation handling (copied from original)
if e.ndim < 4:
e = (
comfy.model_management.cast_to(
self.modulation, dtype=x.dtype, device=x.device
)
+ e
).chunk(6, dim=1)
else:
e = (
comfy.model_management.cast_to(
self.modulation, dtype=x.dtype, device=x.device
).unsqueeze(0)
+ e
).unbind(2)
# Self-attention with RAS caching
y = self.self_attn_with_cache(
torch.addcmul(
self.repeat_e(e[0], x), self.norm1(x), 1 + self.repeat_e(e[1], x)
),
freqs,
idx,
transformer_options=transformer_options,
)
x = torch.addcmul(x, y, self.repeat_e(e[2], x))
# Cross-attention & ffn (unchanged from original)
x = x + self.cross_attn(
self.norm3(x),
context,
context_img_len=context_img_len,
transformer_options=transformer_options,
)
y = self.ffn(
torch.addcmul(
self.repeat_e(e[3], x), self.norm2(x), 1 + self.repeat_e(e[4], x)
)
)
x = torch.addcmul(x, y, self.repeat_e(e[5], x))
if self.last:
# put the relevant tokens back into the cached output
if idx is None or self.manager.cached_output is None:
self.manager.cached_output = x.clone()
else:
self.manager.cached_output[..., idx, :] = x
return self.manager.cached_output
return x
def repeat_e(self, e, x):
"""Helper function for modulation broadcasting"""
repeats = 1
if e.size(1) > 1:
repeats = x.size(1) // e.size(1)
if repeats == 1:
return e
if repeats * e.size(1) == x.size(1):
return torch.repeat_interleave(e, repeats, dim=1)
else:
return torch.repeat_interleave(e, repeats + 1, dim=1)[:, : x.size(1)]
def self_attn_with_cache(self, x, freqs, idx, transformer_options={}):
"""Modified self-attention that uses KV caching"""
b, s, n, d = *x.shape[:2], self.num_heads, self.self_attn.head_dim
# just pull out part of the frequencies
if idx is not None:
freqs = freqs[:, idx]
# Compute QKV like original Wan self-attention
q = self.self_attn.norm_q(self.self_attn.q(x)).view(b, s, n, d)
q = apply_rope1(q, freqs).view(b, s, n * d)
k = self.self_attn.norm_k(self.self_attn.k(x)).view(b, s, n, d)
k = apply_rope1(k, freqs).view(b, s, n * d)
v = self.self_attn.v(x).view(b, s, n * d)
# RAS: KV Cache management
if idx is None:
self.k_cache = k
self.v_cache = v
else:
self.k_cache[:, idx, :] = k
self.v_cache[:, idx, :] = v
x = optimized_attention(
q,
self.k_cache,
self.v_cache,
heads=n,
transformer_options=transformer_options,
)
x = self.self_attn.o(x)
return x
class VaceWanAttentionBlockWrapper(VaceWanAttentionBlock):
"""
Same as the VaceWanAttentionBlock, but uses a RASManager and RASCache to do KV caching.
"""
def __init__(
self, original: VaceWanAttentionBlock, manager: RASManager, first=False
):
nn.Module.__init__(self)
# Copy attributes, handling the case where some might not exist
attributes_to_copy = [
"dim",
"ffn_dim",
"num_heads",
"window_size",
"qk_norm",
"cross_attn_norm",
"eps",
"block_id",
"norm1",
"self_attn",
"norm3",
"cross_attn",
"norm2",
"ffn",
"modulation",
]
# Handle optional attributes that might not exist
optional_attributes = ["before_proj", "after_proj"]
take_attributes_from(original, self, attributes_to_copy)
# Copy optional attributes if they exist
for attr in optional_attributes:
if hasattr(original, attr):
setattr(self, attr, getattr(original, attr))
self.manager = manager
self.k_cache: torch.Tensor
self.v_cache: torch.Tensor
self.first = first
def forward(
self, c, x, e, freqs, context, context_img_len=257, transformer_options={}
):
# RAS: if this is the first block, then we should drop some tokens
idx = self.manager.live_indices()
if self.first:
# For VaceWan models, we track the sequence length differently
self.manager.n_txt = (
0 # Wan doesn't have separate text tokens in the same way
)
# TODO: get vace working
self.manager.n_img = c.shape[1]
if idx is not None:
c = c[..., idx, :]
# VaceWan specific logic (copied from original)
if self.block_id == 0 and hasattr(self, "before_proj"):
c = self.before_proj(c) + x
# Call the parent WanAttentionBlock logic but with our modified attention
# Modulation handling (copied from original)
if e.ndim < 4:
e = (
comfy.model_management.cast_to(
self.modulation, dtype=c.dtype, device=c.device
)
+ e
).chunk(6, dim=1)
else:
e = (
comfy.model_management.cast_to(
self.modulation, dtype=c.dtype, device=c.device
).unsqueeze(0)
+ e
).unbind(2)
# Self-attention with RAS caching
y = self.self_attn_with_cache(
torch.addcmul(
self.repeat_e(e[0], c), self.norm1(c), 1 + self.repeat_e(e[1], c)
),
freqs,
idx,
transformer_options=transformer_options,
)
c = torch.addcmul(c, y, self.repeat_e(e[2], c))
# Cross-attention & ffn (unchanged from original)
c = c + self.cross_attn(
self.norm3(c),
context,
context_img_len=context_img_len,
transformer_options=transformer_options,
)
y = self.ffn(
torch.addcmul(
self.repeat_e(e[3], c), self.norm2(c), 1 + self.repeat_e(e[4], c)
)
)
c = torch.addcmul(c, y, self.repeat_e(e[5], c))
# VaceWan specific projection
c_skip = self.after_proj(c)
return c_skip, c
def repeat_e(self, e, x):
"""Helper function for modulation broadcasting"""
repeats = 1
if e.size(1) > 1:
repeats = x.size(1) // e.size(1)
if repeats == 1:
return e
if repeats * e.size(1) == x.size(1):
return torch.repeat_interleave(e, repeats, dim=1)
else:
return torch.repeat_interleave(e, repeats + 1, dim=1)[:, : x.size(1)]
def self_attn_with_cache(self, x, freqs, idx, transformer_options={}):
"""Modified self-attention that uses KV caching"""
# NOTE the following attention function doesn't work
# it's fucked
# needs to be replaced
# just use the one from wan (above)
b, s, n, d = *x.shape[:2], self.num_heads, x.shape[-1] // self.num_heads
# Compute QKV like original Wan self-attention
q = self.self_attn.norm_q(self.self_attn.q(x)).view(b, s, n, d)
k = self.self_attn.norm_k(self.self_attn.k(x)).view(b, s, n, d)
v = self.self_attn.v(x).view(b, s, n, d)
# Apply rope to Q and K
from comfy.ldm.flux.math import apply_rope1
if idx is not None:
freqs = freqs[:, idx]
q = apply_rope1(q, freqs)
k = apply_rope1(k, freqs)
# RAS: KV Cache management
if idx is None:
# Full hydrate - store the complete k,v tensors as cache
self.k_cache = k
self.v_cache = v
else:
self.k_cache[:, idx, :, :] = k
self.v_cache[:, idx, :, :] = v
q_flat = q.view(b, s, n * d)
k_flat = self.k_cache.view(b, self.k_cache.shape[1], n * d)
v_flat = self.v_cache.view(b, self.v_cache.shape[1], n * d)
attn = optimized_attention(
q_flat,
k_flat,
v_flat,
heads=n,
transformer_options=transformer_options,
)
result = self.self_attn.o(attn)
return result
class HeadWrapper(Head):
"""
Same as the Head (Wan final layer), but reports its output to a manager.
"""
def __init__(self, original: Head, manager: RASManager):
nn.Module.__init__(self)
take_attributes_from(
original,
self,
["dim", "out_dim", "patch_size", "eps", "norm", "head", "modulation"],
)
self.manager = manager
def forward(self, x, e) -> Tensor:
output = super().forward(x, e)
self.manager.select_indices(output, self.manager.timestep)
return output