feat: Support Wan-Animate2 (CORE-358) (#15362)
This commit is contained in:
Jukka Seppänen
GitHub
parent
88fec4b605
commit
a464ac3358
|
|
@ -147,7 +147,8 @@ class WanI2VCrossAttention(WanSelfAttention):
|
|||
v = self.v(context)
|
||||
k_img = self.norm_k_img(self.k_img(context_img))
|
||||
v_img = self.v_img(context_img)
|
||||
img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options)
|
||||
# Sageattn can cause Nans here, don't allow it as there is no speed difference anyway as img attention is tiny.
|
||||
img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options, low_precision_attention=False)
|
||||
# compute attention
|
||||
x = optimized_attention(q, k, v, heads=self.num_heads, transformer_options=transformer_options)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,387 @@
|
|||
# Wan-Animate-2: https://github.com/Wan-Video/Wan-Animate-2
|
||||
"""Wan2.1-I2V-14B weights, driven by a video instead of a motion extractor.
|
||||
|
||||
A pose branch over the pose video's latents runs in lockstep with the generation
|
||||
branch, feeding it K/V per block. The reference image is one extra latent frame at the
|
||||
front of the generation branch, trimmed off by the caller. Upstream calls the pose video
|
||||
the driving video and its branch forward_ref, not to be confused with the reference image.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.model_management
|
||||
import comfy.quant_ops
|
||||
import comfy.utils
|
||||
from comfy.ldm.flux.math import apply_rope1
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
|
||||
from .model import WanAttentionBlock, WanModel, WanSelfAttention, repeat_e, sinusoidal_embedding_1d
|
||||
|
||||
|
||||
class WanAnimate2SelfAttention(WanSelfAttention):
|
||||
|
||||
def qkv(self, x, freqs):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
q = apply_rope1(self.norm_q(self.q(x)).view(b, s, n, d), freqs)
|
||||
k = apply_rope1(self.norm_k(self.k(x)).view(b, s, n, d), freqs)
|
||||
return q, k, self.v(x).view(b, s, n, d)
|
||||
|
||||
def _attn1_patch(self, x, q, k, transformer_options):
|
||||
for p in transformer_options.get("patches", {}).get("attn1_patch", []):
|
||||
x = p({"x": x, "q": q, "k": k, "transformer_options": transformer_options})
|
||||
return x
|
||||
|
||||
def kv(self, x, freqs):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
return apply_rope1(self.norm_k(self.k(x)).view(b, s, n, d), freqs), self.v(x).view(b, s, n, d)
|
||||
|
||||
def forward_pose(self, x, freqs, transformer_options={}):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
q, k, v = self.qkv(x, freqs)
|
||||
out = optimized_attention(q.reshape(b, s, n * d), k.reshape(b, s, n * d), v.reshape(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options)
|
||||
return self.o(self._attn1_patch(out, q, k, transformer_options)), k, v
|
||||
|
||||
def forward_gen(self, x, freqs, k_pose, v_pose, f_gen, hw, buffers, ref_strength=1.0, transformer_options={}):
|
||||
# frame j attends every gen token plus pose frame j-1 (frame 0 is the reference slot and has none)
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
q, k, v = self.qkv(x, freqs)
|
||||
if ref_strength != 1.0:
|
||||
v[:, :hw] *= ref_strength # frame 0 is the reference image's slot
|
||||
|
||||
if k_pose is None: # pose influence windowed out: plain self-attention, no per-frame loop
|
||||
out = optimized_attention(q.reshape(b, s, n * d), k.reshape(b, s, n * d), v.reshape(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options)
|
||||
return self.o(self._attn1_patch(out, q, k, transformer_options))
|
||||
|
||||
# gen half is the same every frame; only the hw-token pose tail is rewritten
|
||||
kbuf, vbuf, out = buffers
|
||||
kbuf[:, :s] = k
|
||||
vbuf[:, :s] = v
|
||||
|
||||
for j in range(f_gen):
|
||||
q_j = q[:, j * hw:(j + 1) * hw].reshape(b, hw, n * d)
|
||||
if j == 0:
|
||||
kk, vv = k, v
|
||||
else:
|
||||
kbuf[:, s:] = k_pose[:, (j - 1) * hw:j * hw]
|
||||
vbuf[:, s:] = v_pose[:, (j - 1) * hw:j * hw]
|
||||
kk, vv = kbuf, vbuf
|
||||
out[:, j * hw:(j + 1) * hw] = optimized_attention(q_j, kk.reshape(b, kk.shape[1], n * d), vv.reshape(b, kk.shape[1], n * d), heads=self.num_heads, transformer_options=transformer_options)
|
||||
return self.o(self._attn1_patch(out, q, k, transformer_options))
|
||||
|
||||
|
||||
class WanAnimate2Block(WanAttentionBlock):
|
||||
|
||||
def __init__(self, cross_attn_type, dim, ffn_dim, num_heads, window_size=(-1, -1), qk_norm=True, cross_attn_norm=False, eps=1e-6, operation_settings={}):
|
||||
super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
|
||||
self.self_attn = WanAnimate2SelfAttention(dim, num_heads, window_size, qk_norm, eps, operation_settings=operation_settings)
|
||||
|
||||
def _modulation(self, e, x):
|
||||
if e.ndim < 4:
|
||||
return (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
|
||||
return (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e).unbind(2)
|
||||
|
||||
def _cross_attn_ffn(self, x, e, context, context_img_len, transformer_options):
|
||||
x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
|
||||
for p in transformer_options.get("patches", {}).get("attn2_patch", []):
|
||||
x = p({"x": x, "transformer_options": transformer_options})
|
||||
y = self.ffn(torch.addcmul(repeat_e(e[3], x), self.norm2(x), 1 + repeat_e(e[4], x)))
|
||||
return torch.addcmul(x, y, repeat_e(e[5], x))
|
||||
|
||||
def forward_pose(self, x, e, freqs, context, context_img_len=257, transformer_options={}):
|
||||
e = self._modulation(e, x)
|
||||
x = x.contiguous()
|
||||
y, k, v = self.self_attn.forward_pose(torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)), freqs, transformer_options=transformer_options)
|
||||
x = torch.addcmul(x, y, repeat_e(e[2], x))
|
||||
del y
|
||||
return self._cross_attn_ffn(x, e, context, context_img_len, transformer_options), k, v
|
||||
|
||||
def kv_from_input(self, x_pose, e, freqs, transformer_options={}):
|
||||
e = self._modulation(e, x_pose)
|
||||
x_pose = x_pose.contiguous()
|
||||
return self.self_attn.kv(torch.addcmul(repeat_e(e[0], x_pose), self.norm1(x_pose), 1 + repeat_e(e[1], x_pose)), freqs)
|
||||
|
||||
def forward_gen(self, x, e, freqs, context, k_pose, v_pose, f_gen, hw, buffers, ref_strength=1.0, context_img_len=257, transformer_options={}):
|
||||
e = self._modulation(e, x)
|
||||
x = x.contiguous()
|
||||
y = self.self_attn.forward_gen(torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)), freqs, k_pose, v_pose, f_gen, hw, buffers, ref_strength=ref_strength, transformer_options=transformer_options)
|
||||
x = torch.addcmul(x, y, repeat_e(e[2], x))
|
||||
del y
|
||||
return self._cross_attn_ffn(x, e, context, context_img_len, transformer_options)
|
||||
|
||||
|
||||
class PoseBranchCache:
|
||||
"""Pose-branch block inputs, reused across the sampling steps of one execution.
|
||||
|
||||
Caching the block input rather than its K/V halves the memory; reprojecting K/V on read
|
||||
costs ~4% of re-running the block. One slot per distinct pose sequence, so under
|
||||
context windows each window keeps its own; least recently used slots are evicted when
|
||||
the store device runs low on memory. Created and freed by WanAnimate2Cache.
|
||||
"""
|
||||
|
||||
CONVROT_GROUPSIZE = 256
|
||||
|
||||
def __init__(self, store_device=None, dtype="default"):
|
||||
self.store_device = torch.device(store_device) if store_device is not None else torch.device("cpu")
|
||||
self.dtype = dtype
|
||||
self.slots = [] # most recently used last
|
||||
self.slot = None
|
||||
self._pending = {}
|
||||
self._staging = {}
|
||||
|
||||
def select(self, pose_latents):
|
||||
# select runs at a forward boundary: an interrupted forward can leave copies in flight that a different slot's forward would then mistake for its own
|
||||
if self._pending:
|
||||
for t, stream in self._pending.values():
|
||||
if stream is not None:
|
||||
stream.synchronize()
|
||||
self._pending = {}
|
||||
# keyed on batch element 0, so a cond batch size change mid-run stays valid
|
||||
k = pose_latents[:1]
|
||||
for s in self.slots:
|
||||
if s["key"].shape == k.shape and torch.equal(s["key"], k.to(s["key"].device)):
|
||||
self.slots.remove(s)
|
||||
self.slots.append(s)
|
||||
self.slot = s
|
||||
return
|
||||
# cache what fits: a filled slot is the size estimate for the next one, and least recently used slots make room when the store device runs low
|
||||
est = max((self._slot_bytes(s) for s in self.slots), default=0) * 1.5
|
||||
while self.slots and comfy.model_management.get_free_memory(self.store_device) < est:
|
||||
self._free_slot(self.slots.pop(0))
|
||||
self.slot = {"key": k.clone().to(self.store_device), "blocks": {}, "params": {}, "shape": None, "pinned": []}
|
||||
self.slots.append(self.slot)
|
||||
|
||||
def _free_slot(self, s):
|
||||
for t, stream in self._pending.values():
|
||||
if stream is not None:
|
||||
stream.synchronize() # an aborted forward can leave a copy in flight, still reading memory we are about to unpin
|
||||
self._pending = {}
|
||||
for t in s["pinned"]:
|
||||
comfy.model_management.unpin_memory(t)
|
||||
|
||||
def free(self):
|
||||
for s in self.slots:
|
||||
self._free_slot(s)
|
||||
self.slots = []
|
||||
self.slot = None
|
||||
self._staging = {}
|
||||
|
||||
def filled(self, num_blocks):
|
||||
return self.slot is not None and len(self.slot["blocks"]) == num_blocks
|
||||
|
||||
def put(self, i, x_pose):
|
||||
t = x_pose[:1]
|
||||
params = None
|
||||
if self.dtype in ("int8", "int4"):
|
||||
# convrot is what lets low-bit survive the ~125x per-channel outliers here, and over a [tokens, dim] view per-row scale means per-token. The kernels want 2D and a power-of-4 group that divides dim.
|
||||
self.slot["shape"] = t.shape
|
||||
g = self.CONVROT_GROUPSIZE
|
||||
while g > 4 and t.shape[-1] % g:
|
||||
g //= 4
|
||||
if self.dtype == "int4":
|
||||
t, params = comfy.quant_ops.TensorCoreConvRotW4A4Layout.quantize(t.reshape(-1, t.shape[-1]), convrot_groupsize=g)
|
||||
else:
|
||||
t, params = comfy.quant_ops.TensorWiseINT8Layout.quantize(t.reshape(-1, t.shape[-1]), is_weight=True, per_channel=True, convrot=True, convrot_groupsize=g)
|
||||
|
||||
t = t.to(self.store_device, copy=True)
|
||||
if comfy.model_management.pin_memory(t):
|
||||
self.slot["pinned"].append(t)
|
||||
self.slot["blocks"][i] = t
|
||||
# the scales follow the blocks off the GPU: per-window slots would otherwise pile them up in VRAM (~200 MB per window at 480p int4)
|
||||
self.slot["params"][i] = params if params is None else params.to_device(self.store_device)
|
||||
|
||||
def prefetch(self, i, device, dtype):
|
||||
# call before the compute this should overlap, so the stream waits only on work already enqueued
|
||||
if i not in self.slot["blocks"] or i in self._pending:
|
||||
return
|
||||
t = self.slot["blocks"][i]
|
||||
cast_dtype = None if self.slot["params"][i] is not None else dtype # int8 entries move in their stored dtype and widen in take()
|
||||
stream = None
|
||||
r = None
|
||||
if t.device != device:
|
||||
stream = comfy.model_management.get_offload_stream(device)
|
||||
cs = comfy.model_management.current_stream(device)
|
||||
if stream is not None and cs is not None:
|
||||
# the handed-out stream last waited on the main stream a full rotation ago, which does not cover the previous consumer's reads of this slot; wait now so the copy cannot overwrite a slot still being read
|
||||
stream.wait_stream(cs)
|
||||
# two persistent staging buffers per tensor shape instead of a fresh allocation per block (~29 GB of churn per pass at 720p); windows of different lengths get their own pair
|
||||
buf_key = (tuple(t.shape), cast_dtype if cast_dtype is not None else t.dtype)
|
||||
if buf_key not in self._staging:
|
||||
self._staging[buf_key] = [torch.empty(t.shape, dtype=buf_key[1], device=device) for _ in range(2)]
|
||||
r = self._staging[buf_key][i % 2]
|
||||
self._pending[i] = (comfy.model_management.cast_to(t, cast_dtype, device, non_blocking=True, stream=stream, r=r), stream)
|
||||
|
||||
def take(self, i, device, dtype, batch_size):
|
||||
if i not in self._pending:
|
||||
self.prefetch(i, device, dtype)
|
||||
t, stream = self._pending.pop(i)
|
||||
comfy.model_management.sync_stream(device, stream)
|
||||
params = self.slot["params"][i]
|
||||
if params is not None:
|
||||
layout = comfy.quant_ops.TensorCoreConvRotW4A4Layout if self.dtype == "int4" else comfy.quant_ops.TensorWiseINT8Layout
|
||||
t = layout.dequantize(t, params.to_device(t.device)).reshape(self.slot["shape"]).to(dtype)
|
||||
return comfy.utils.repeat_to_batch_size(t, batch_size)
|
||||
|
||||
def _slot_bytes(self, s):
|
||||
return sum(t.numel() * t.element_size() for t in s["blocks"].values())
|
||||
|
||||
def memory_bytes(self):
|
||||
return sum(self._slot_bytes(s) for s in self.slots)
|
||||
|
||||
|
||||
class WanAnimate2Model(WanModel):
|
||||
|
||||
def __init__(self,
|
||||
model_type='animate2',
|
||||
patch_size=(1, 2, 2),
|
||||
text_len=512,
|
||||
in_dim=36,
|
||||
dim=5120,
|
||||
ffn_dim=13824,
|
||||
freq_dim=256,
|
||||
text_dim=4096,
|
||||
out_dim=16,
|
||||
num_heads=40,
|
||||
num_layers=40,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
eps=1e-6,
|
||||
flf_pos_embed_token_number=None,
|
||||
in_dim_ref_conv=None,
|
||||
image_model=None,
|
||||
device=None, dtype=None, operations=None,
|
||||
):
|
||||
# model_type is 'animate2' in unet_config, but the checkpoint is i2v-shaped
|
||||
super().__init__(model_type='i2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim,
|
||||
text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm,
|
||||
cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, in_dim_ref_conv=in_dim_ref_conv,
|
||||
wan_attn_block_class=WanAnimate2Block, image_model=image_model, device=device, dtype=dtype, operations=operations)
|
||||
|
||||
def rope_encode_pose(self, t, h, w, w_patches, device=None, dtype=None):
|
||||
# t_start=1 lines pose frame j up with gen frame j+1, past the reference slot; shift_x parks it in its own strip of rope space.
|
||||
# The caller's rope_options are a user scaling knob and deliberately not forwarded.
|
||||
return super().rope_encode(t, h, w, t_start=1, device=device, dtype=dtype, transformer_options={"rope_options": {"shift_x": float(w_patches)}})
|
||||
|
||||
def _forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, pose_latents=None, clip_fea_pose=None, context_pose=None, **kwargs):
|
||||
bs, c, t, h, w = x.shape
|
||||
x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
|
||||
|
||||
# h/w pre-pad: rope_encode's rounding reproduces the post-pad grid
|
||||
freqs = self.rope_encode(t, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options)
|
||||
freqs_pose = None
|
||||
if pose_latents is not None: # absent when the node's timestep window excludes this step
|
||||
pose_latents = comfy.ldm.common_dit.pad_to_patch_size(pose_latents.to(x.dtype), self.patch_size)
|
||||
w_patches = (w + (self.patch_size[2] // 2)) // self.patch_size[2]
|
||||
freqs_pose = self.rope_encode_pose(pose_latents.shape[2], h, w, w_patches, device=x.device, dtype=x.dtype)
|
||||
|
||||
return self.forward_orig(x, timestep, context, clip_fea=clip_fea, freqs=freqs, freqs_pose=freqs_pose, pose_latents=pose_latents,
|
||||
clip_fea_pose=clip_fea_pose, context_pose=context_pose, transformer_options=transformer_options, **kwargs)[:, :, :t, :h, :w]
|
||||
|
||||
def forward_orig(self, x, t, context, clip_fea=None, freqs=None, freqs_pose=None, pose_latents=None, clip_fea_pose=None, context_pose=None, pose_strength=1.0, reference_strength=1.0, transformer_options={}, **kwargs):
|
||||
x_input = x[:, :, 1:] # video-only: frame 0 is the reference slot, offset past it below
|
||||
x = self.patch_embedding(x.float()).to(x.dtype)
|
||||
grid_sizes = x.shape[2:]
|
||||
transformer_options["grid_sizes"] = grid_sizes
|
||||
f_gen, gh, gw = grid_sizes
|
||||
hw = gh * gw
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
|
||||
# the node windows the pose influence via cond timestep ranges: outside the window the cond carries no pose latents, and the branch, its cache traffic and the per-frame attention loop are all skipped
|
||||
apply_pose = pose_latents is not None
|
||||
if apply_pose and pose_latents.shape[2] != f_gen - 1: # before cache.select, which would otherwise keep an empty slot keyed to the rejected latents
|
||||
raise ValueError("pose branch has {} latent frames, expected {} (generation frames minus the reference-image slot)".format(pose_latents.shape[2], f_gen - 1))
|
||||
|
||||
cache = transformer_options.get("animate2_cache", None) if apply_pose else None
|
||||
if cache is not None:
|
||||
cache.select(pose_latents)
|
||||
cached = cache is not None and cache.filled(len(self.blocks))
|
||||
|
||||
x_pose = None
|
||||
if not cached and apply_pose:
|
||||
# 36ch = [latents(16) | mask(4) | latents(16)]; latents twice, and the mask is all ones since every pose frame is known
|
||||
x_pose = self.patch_embedding(torch.cat([pose_latents, torch.ones_like(pose_latents[:, :4]), pose_latents], dim=1).float()).to(x.dtype)
|
||||
x_pose = x_pose.flatten(2).transpose(1, 2)
|
||||
|
||||
# time embeddings
|
||||
e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x.dtype))
|
||||
e = e.reshape(t.shape[0], -1, e.shape[-1])
|
||||
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
|
||||
|
||||
e0_pose = None
|
||||
if apply_pose:
|
||||
t_pose = torch.ones_like(t.flatten())
|
||||
e_pose = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t_pose).to(dtype=x.dtype))
|
||||
e_pose = e_pose.reshape(t.shape[0], -1, e_pose.shape[-1])
|
||||
e0_pose = self.time_projection(e_pose).unflatten(2, (6, self.dim))
|
||||
|
||||
context_gen = self.text_embedding(context)
|
||||
|
||||
context_img_len = None
|
||||
if clip_fea is not None:
|
||||
if self.img_emb is not None:
|
||||
context_gen = torch.cat([self.img_emb(clip_fea), context_gen], dim=1)
|
||||
context_img_len = clip_fea.shape[-2]
|
||||
|
||||
context_img_len_pose = None
|
||||
if not cached and apply_pose:
|
||||
context_pose = self.text_embedding(context if context_pose is None else context_pose)
|
||||
clip_fea_pose = clip_fea if clip_fea_pose is None else clip_fea_pose
|
||||
if clip_fea_pose is not None:
|
||||
if self.img_emb is not None:
|
||||
context_pose = torch.cat([self.img_emb(clip_fea_pose), context_pose], dim=1)
|
||||
context_img_len_pose = clip_fea_pose.shape[-2]
|
||||
|
||||
patches_replace = transformer_options.get("patches_replace", {})
|
||||
patches = transformer_options.get("patches", {})
|
||||
blocks_replace = patches_replace.get("dit", {})
|
||||
transformer_options["total_blocks"] = len(self.blocks)
|
||||
transformer_options["block_type"] = "double"
|
||||
|
||||
if cache is not None and not cached and apply_pose and "context_window" in transformer_options:
|
||||
# pose-only prepass, to avoid inflating dynamic VRAM calibration when using multiple context windows
|
||||
for i, block in enumerate(self.blocks):
|
||||
transformer_options["block_index"] = i
|
||||
cache.put(i, x_pose)
|
||||
x_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)[0]
|
||||
x_pose = None
|
||||
cached = True
|
||||
|
||||
buffers = None
|
||||
if apply_pose:
|
||||
# allocated once and reused by every block
|
||||
n, d = self.num_heads, self.dim // self.num_heads
|
||||
buffers = (x.new_empty(x.shape[0], x.shape[1] + hw, n, d), x.new_empty(x.shape[0], x.shape[1] + hw, n, d), x.new_empty(x.shape[0], x.shape[1], self.dim))
|
||||
|
||||
for i, block in enumerate(self.blocks):
|
||||
transformer_options["block_index"] = i
|
||||
|
||||
if not apply_pose:
|
||||
k_pose = v_pose = None
|
||||
elif cached:
|
||||
x_pose_in = cache.take(i, x.device, x.dtype, x.shape[0])
|
||||
cache.prefetch(i + 1, x.device, x.dtype) # queue the next block before the gen compute it should overlap
|
||||
k_pose, v_pose = block.kv_from_input(x_pose_in, e0_pose, freqs_pose, transformer_options=transformer_options)
|
||||
del x_pose_in
|
||||
else:
|
||||
if cache is not None:
|
||||
cache.put(i, x_pose)
|
||||
# runs even under a block replace: its state has to reach block i+1
|
||||
x_pose, k_pose, v_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)
|
||||
if v_pose is not None and pose_strength != 1.0:
|
||||
v_pose = v_pose * pose_strength
|
||||
|
||||
if ("double_block", i) in blocks_replace:
|
||||
def block_wrap(args, block=block, k_pose=k_pose, v_pose=v_pose):
|
||||
return {"img": block.forward_gen(args["img"], args["vec"], args["pe"], args["txt"], k_pose, v_pose, f_gen, hw, buffers, ref_strength=reference_strength, context_img_len=context_img_len, transformer_options=args["transformer_options"])}
|
||||
out = blocks_replace[("double_block", i)]({"img": x, "txt": context_gen, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
|
||||
x = out["img"]
|
||||
else:
|
||||
x = block.forward_gen(x, e0, freqs, context_gen, k_pose, v_pose, f_gen, hw, buffers, ref_strength=reference_strength, context_img_len=context_img_len, transformer_options=transformer_options)
|
||||
|
||||
if "double_block" in patches:
|
||||
for p in patches["double_block"]:
|
||||
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": hw, "transformer_options": transformer_options})
|
||||
x = out["img"]
|
||||
|
||||
return self.unpatchify(self.head(x, e), grid_sizes)
|
||||
|
|
@ -46,6 +46,7 @@ import comfy.ldm.cosmos.predict2
|
|||
import comfy.ldm.lumina.model
|
||||
import comfy.ldm.wan.model
|
||||
import comfy.ldm.wan.model_animate
|
||||
import comfy.ldm.wan.model_animate2
|
||||
import comfy.ldm.wan.ar_model
|
||||
import comfy.ldm.wan.model_wandancer
|
||||
import comfy.ldm.hunyuan3d.model
|
||||
|
|
@ -1813,6 +1814,41 @@ class WAN22_Animate(WAN21):
|
|||
return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1)
|
||||
return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
|
||||
|
||||
class WAN_Animate2(WAN21):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model_animate2.WanAnimate2Model)
|
||||
self.image_to_video = True
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
|
||||
pose_video_latent = kwargs.get("pose_video_latent", None)
|
||||
if pose_video_latent is not None:
|
||||
out['pose_latents'] = comfy.conds.CONDRegular(self.process_latent_in(pose_video_latent))
|
||||
|
||||
clip_vision_output_pose = kwargs.get("clip_vision_output_pose", None)
|
||||
if clip_vision_output_pose is not None:
|
||||
out['clip_fea_pose'] = comfy.conds.CONDRegular(clip_vision_output_pose.penultimate_hidden_states)
|
||||
|
||||
cross_attn_pose = kwargs.get("cross_attn_pose", None)
|
||||
if cross_attn_pose is not None:
|
||||
out['context_pose'] = comfy.conds.CONDRegular(cross_attn_pose)
|
||||
|
||||
pose_strength = kwargs.get("pose_strength", 1.0)
|
||||
if pose_strength != 1.0:
|
||||
out['pose_strength'] = comfy.conds.CONDConstant(pose_strength)
|
||||
|
||||
reference_strength = kwargs.get("reference_strength", 1.0)
|
||||
if reference_strength != 1.0:
|
||||
out['reference_strength'] = comfy.conds.CONDConstant(reference_strength)
|
||||
|
||||
return out
|
||||
|
||||
def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]):
|
||||
if cond_key == "pose_latents":
|
||||
return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1)
|
||||
return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
|
||||
|
||||
class WAN22_S2V(WAN21):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.WanModel_S2V)
|
||||
|
|
|
|||
|
|
@ -1450,6 +1450,20 @@ class WAN22_Animate(WAN21_T2V):
|
|||
out = model_base.WAN22_Animate(self, device=device)
|
||||
return out
|
||||
|
||||
class WAN_Animate2(WAN21_T2V):
|
||||
unet_config = {
|
||||
"image_model": "wan2.1",
|
||||
"model_type": "animate2",
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
"shift": 5.0,
|
||||
}
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
out = model_base.WAN_Animate2(self, device=device)
|
||||
return out
|
||||
|
||||
class WAN22_T2V(WAN21_T2V):
|
||||
unet_config = {
|
||||
"image_model": "wan2.1",
|
||||
|
|
@ -2464,6 +2478,7 @@ models = [
|
|||
WAN22_S2V,
|
||||
WAN21_HuMo,
|
||||
WAN22_Animate,
|
||||
WAN_Animate2,
|
||||
WAN21_FlowRVS,
|
||||
WAN21_SCAIL,
|
||||
WAN21_SCAIL2,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ import comfy.model_management
|
|||
import comfy.utils
|
||||
import comfy.latent_formats
|
||||
import comfy.clip_vision
|
||||
import comfy.ldm.wan.model_animate2
|
||||
import comfy.patcher_extension
|
||||
import json
|
||||
import numpy as np
|
||||
from typing import Tuple, TypedDict
|
||||
|
|
@ -1248,6 +1250,167 @@ class WanAnimateToVideo(io.ComfyNode):
|
|||
out_latent["samples"] = latent
|
||||
return io.NodeOutput(positive, negative, out_latent, trim_latent, max(0, ref_motion_latent_length * 4 - 3), video_frame_offset + length)
|
||||
|
||||
class WanAnimate2ToVideo(io.ComfyNode):
|
||||
|
||||
CONTINUE_MOTION_FRAMES = 1
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="WanAnimate2ToVideo",
|
||||
category="model/conditioning/wan/animate",
|
||||
description="Animate a character in a reference image using a video, effectively replicating the facial animation, body motion and hands gesture from the video.",
|
||||
inputs=[
|
||||
io.Conditioning.Input("positive"),
|
||||
io.Conditioning.Input("negative"),
|
||||
io.Vae.Input("vae"),
|
||||
io.Int.Input("width", default=832, min=16, max=nodes.MAX_RESOLUTION, step=16, tooltip="Output video width in pixels."),
|
||||
io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16, tooltip="Output video height in pixels."),
|
||||
io.Int.Input("length", default=81, min=1, max=nodes.MAX_RESOLUTION, step=4, tooltip="Number of frames to generate."),
|
||||
io.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="Number of videos to generate simultaneously."),
|
||||
io.Image.Input("reference_image", optional=True, tooltip="The character to animate."),
|
||||
io.Image.Input("pose_video", optional=True, tooltip="The video whose motion is transferred to the reference character."),
|
||||
io.ClipVisionOutput.Input("clip_vision_output", optional=True, tooltip="CLIP vision of the reference image."),
|
||||
io.Conditioning.Input("positive_pose", optional=True, tooltip="Prompt for the pose-video branch, describing the motion rather than the character. Defaults to positive. Used for both the cond and uncond passes."),
|
||||
io.ClipVisionOutput.Input("clip_vision_output_pose", optional=True, tooltip="CLIP vision of the pose video's first frame. Defaults to clip_vision_output."),
|
||||
io.Image.Input("continue_motion", optional=True, tooltip="Previous motion sequence to continue from for temporal consistency."),
|
||||
io.Int.Input("video_frame_offset", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1, tooltip="Frames to seek into the pose video. Connect to the video_frame_offset output of the previous node when extending."),
|
||||
io.Float.Input("pose_strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Scales the pose video's influence on the motion. 1.0 is the trained behavior; below weakens adherence, above amplifies. 0.0 mutes it but does not fully remove it."),
|
||||
io.Float.Input("pose_start_percent", default=0.0, min=0.0, max=1.0, step=0.01, tooltip="Sampling percent at which the pose influence starts. Outside the window the pose branch is skipped entirely, which also speeds those steps up."),
|
||||
io.Float.Input("pose_end_percent", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Sampling percent at which the pose influence ends. Motion is mostly established early, so e.g. 0.7 can loosen fine detail while keeping the choreography."),
|
||||
io.Float.Input("reference_image_strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Scales how strongly generated frames attend to the reference image's latent frame. Below 1.0 loosens identity/appearance adherence (e.g. to let the prompt restyle), above tightens it against drift."),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(display_name="positive"),
|
||||
io.Conditioning.Output(display_name="negative"),
|
||||
io.Latent.Output(display_name="latent"),
|
||||
io.Int.Output(display_name="trim_latent", tooltip="Number of latent frames that should be trimmed before decoding."),
|
||||
io.Int.Output(display_name="trim_image", tooltip="Number of overlapping image frames when extending a video."),
|
||||
io.Int.Output(display_name="video_frame_offset", tooltip="Frames to seek into the pose video."),
|
||||
],
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, positive, negative, vae, width, height, length, batch_size, video_frame_offset, reference_image=None, pose_video=None, clip_vision_output=None, positive_pose=None, clip_vision_output_pose=None, continue_motion=None, pose_strength=1.0, pose_start_percent=0.0, pose_end_percent=1.0, reference_image_strength=1.0) -> io.NodeOutput:
|
||||
if pose_start_percent > pose_end_percent:
|
||||
raise ValueError("pose_start_percent ({}) must not be greater than pose_end_percent ({}).".format(pose_start_percent, pose_end_percent))
|
||||
latent_length = ((length - 1) // 4) + 1
|
||||
latent_width = width // 8
|
||||
latent_height = height // 8
|
||||
|
||||
if reference_image is None:
|
||||
reference_image = torch.zeros((1, height, width, 3))
|
||||
|
||||
ref_image = comfy.utils.common_upscale(reference_image[:1].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
|
||||
ref_latent = vae.encode(ref_image[:, :, :, :3])
|
||||
trim_latent = ref_latent.shape[2]
|
||||
|
||||
ref_motion_latent_length = 0
|
||||
if continue_motion is None:
|
||||
image = torch.ones((length, height, width, 3)) * 0.5
|
||||
else:
|
||||
continue_motion = continue_motion[-cls.CONTINUE_MOTION_FRAMES:]
|
||||
video_frame_offset = max(0, video_frame_offset - continue_motion.shape[0])
|
||||
continue_motion = comfy.utils.common_upscale(continue_motion[-length:].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
|
||||
# 0.5 is mid-grey, matching upstream's zeros in [-1, 1] pixel space
|
||||
image = torch.ones((length, height, width, continue_motion.shape[-1]), device=continue_motion.device, dtype=continue_motion.dtype) * 0.5
|
||||
image[:continue_motion.shape[0]] = continue_motion
|
||||
ref_motion_latent_length += ((continue_motion.shape[0] - 1) // 4) + 1
|
||||
|
||||
concat_latent_image = torch.cat((ref_latent, vae.encode(image[:, :, :, :3])), dim=2)
|
||||
|
||||
# 1-channel, 0 == known; concat_cond inverts and repeats it to the 4 mask channels
|
||||
mask = torch.ones((1, 1, latent_length + trim_latent, latent_height, latent_width), device=concat_latent_image.device, dtype=concat_latent_image.dtype)
|
||||
mask[:, :, :trim_latent + ref_motion_latent_length] = 0.0
|
||||
|
||||
positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
|
||||
negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
|
||||
|
||||
if clip_vision_output is not None:
|
||||
positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
|
||||
negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})
|
||||
|
||||
# not windowed with the pose values: the reference frame is part of the latent on every step
|
||||
if reference_image_strength != 1.0:
|
||||
positive = node_helpers.conditioning_set_values(positive, {"reference_strength": reference_image_strength})
|
||||
negative = node_helpers.conditioning_set_values(negative, {"reference_strength": reference_image_strength})
|
||||
|
||||
# set on the negative too: upstream runs the pose branch once, outside the CFG loop, so it never sees the negative prompt
|
||||
pose_values = {}
|
||||
if pose_video is not None:
|
||||
if pose_video.shape[0] <= video_frame_offset:
|
||||
raise ValueError("pose_video has {} frames but video_frame_offset is {} -- nothing left to read.".format(pose_video.shape[0], video_frame_offset))
|
||||
pose_video = pose_video[video_frame_offset:]
|
||||
pose_video = comfy.utils.common_upscale(pose_video[:length].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
|
||||
if pose_video.shape[0] < length: # hold the last frame, as upstream pads its clips
|
||||
pose_video = torch.cat((pose_video,) + (pose_video[-1:],) * (length - pose_video.shape[0]), dim=0)
|
||||
pose_values["pose_video_latent"] = vae.encode(pose_video[:, :, :, :3])
|
||||
|
||||
pose_clip = clip_vision_output_pose if clip_vision_output_pose is not None else clip_vision_output
|
||||
if pose_clip is not None:
|
||||
pose_values["clip_vision_output_pose"] = pose_clip
|
||||
|
||||
pose_cond = positive_pose if positive_pose is not None else positive
|
||||
if len(pose_cond) > 0:
|
||||
pose_values["cross_attn_pose"] = pose_cond[0][0]
|
||||
|
||||
if pose_strength != 1.0:
|
||||
pose_values["pose_strength"] = pose_strength
|
||||
|
||||
if pose_start_percent > 0.0 or pose_end_percent < 1.0:
|
||||
# windowed via cond timestep ranges: the pose values ride a cond limited to the window, and complement conds without them cover the rest, where the model runs without the pose branch at all
|
||||
def windowed(cond):
|
||||
parts = node_helpers.conditioning_set_values(cond, {**pose_values, "start_percent": pose_start_percent, "end_percent": pose_end_percent})
|
||||
if pose_start_percent > 0.0:
|
||||
parts = parts + node_helpers.conditioning_set_values(cond, {"start_percent": 0.0, "end_percent": pose_start_percent})
|
||||
if pose_end_percent < 1.0:
|
||||
parts = parts + node_helpers.conditioning_set_values(cond, {"start_percent": pose_end_percent, "end_percent": 1.0})
|
||||
return parts
|
||||
positive = windowed(positive)
|
||||
negative = windowed(negative)
|
||||
else:
|
||||
positive = node_helpers.conditioning_set_values(positive, pose_values)
|
||||
negative = node_helpers.conditioning_set_values(negative, pose_values)
|
||||
|
||||
latent = torch.zeros([batch_size, 16, latent_length + trim_latent, latent_height, latent_width], device=comfy.model_management.intermediate_device())
|
||||
out_latent = {}
|
||||
out_latent["samples"] = latent
|
||||
return io.NodeOutput(positive, negative, out_latent, trim_latent, max(0, ref_motion_latent_length * 4 - 3), video_frame_offset + length)
|
||||
|
||||
|
||||
class WanAnimate2Cache(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="WanAnimate2Cache",
|
||||
category="model/conditioning/wan/animate",
|
||||
description=(
|
||||
"Caches the pose-video's per-block activations so it runs once instead of on every sampling step. Roughly halves generation time "
|
||||
"Tradeoff is ~12.5 GB of system RAM at 480x832/81 frames in bf16 (scales with resolution and length). "
|
||||
"With context windows each window is cached separately, so RAM scales with the window count; use the static_standard schedule, as uniform schedules shift the windows every step and nothing ever recurs to hit the cache."
|
||||
),
|
||||
inputs=[
|
||||
io.Model.Input("model"),
|
||||
io.Combo.Input("device", options=["cpu", "gpu"], default="cpu",
|
||||
tooltip="Where to keep the cache. cpu (RAM) is the safe choice, the cache will not fit in VRAM alongside the model at typical sizes. gpu (VRAM) can be faster if it fits."),
|
||||
io.Combo.Input("dtype", options=["default", "int8", "int4"], default="default",
|
||||
tooltip="Storage precision. default stores the activations in the model's compute dtype. int8 halves the cache, int4 quarters it, convrot is used to retain accuracy."),
|
||||
],
|
||||
outputs=[io.Model.Output()],
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model, device, dtype="default") -> io.NodeOutput:
|
||||
store = comfy.model_management.get_torch_device() if device == "gpu" else torch.device("cpu")
|
||||
cache = comfy.ldm.wan.model_animate2.PoseBranchCache(store_device=store, dtype=dtype)
|
||||
m = model.clone()
|
||||
m.model_options["transformer_options"]["animate2_cache"] = cache
|
||||
m.add_callback(comfy.patcher_extension.CallbacksMP.ON_CLEANUP, lambda patcher: cache.free())
|
||||
return io.NodeOutput(m)
|
||||
|
||||
|
||||
class Wan22ImageToVideoLatent(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
|
|
@ -1475,6 +1638,8 @@ class WanExtension(ComfyExtension):
|
|||
WanSoundImageToVideoExtend,
|
||||
WanHuMoImageToVideo,
|
||||
WanAnimateToVideo,
|
||||
WanAnimate2ToVideo,
|
||||
WanAnimate2Cache,
|
||||
Wan22ImageToVideoLatent,
|
||||
WanInfiniteTalkToVideo,
|
||||
]
|
||||
|
|
|
|||
Loading…
Reference in New Issue