File size: 9,736 Bytes
832fee5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | import inspect
import logging
import torch
import comfy.conds
import comfy.model_management
from comfy.ldm.wan.model import AudioInjector_WAN, WanModel_S2V
from comfy.model_base import WAN22_S2V
def _append_context_latents(self, x, kwargs):
context_latents = kwargs.get("context_latents", None)
if context_latents is None:
return x
for lat in context_latents:
cl = self.patch_embedding(lat.float().to(x.device)).to(x.dtype).flatten(2).transpose(1, 2)
x = torch.cat([x, cl], dim=1)
return x
def _patch_wan_model_s2v_forward():
if getattr(WanModel_S2V.forward_orig, "__wan_bernini_s2v_v2_patch__", False):
return
try:
source = inspect.getsource(WanModel_S2V.forward_orig)
except (OSError, TypeError):
source = ""
if "context_latents" in source and getattr(WanModel_S2V.forward_orig, "__wan_bernini_s2v_patch__", False):
WanModel_S2V.forward_orig.__wan_bernini_s2v_v2_patch__ = True
return
original = WanModel_S2V.forward_orig
def forward_orig(
self,
x,
t,
context,
audio_embed=None,
reference_latent=None,
control_video=None,
reference_motion=None,
clip_fea=None,
freqs=None,
transformer_options={},
**kwargs,
):
if audio_embed is not None:
num_embeds = x.shape[-3] * 4
audio_emb_global, audio_emb = self.casual_audio_encoder(audio_embed[:, :, :, :num_embeds])
else:
audio_emb = None
audio_emb_global = None
bs, _, time, height, width = x.shape
x = self.patch_embedding(x.float()).to(x.dtype)
if control_video is not None:
x = x + self.cond_encoder(control_video)
if t.ndim == 1:
t = t.unsqueeze(1).repeat(1, x.shape[2])
grid_sizes = x.shape[2:]
x = x.flatten(2).transpose(1, 2)
seq_len = x.size(1)
cond_mask_weight = comfy.model_management.cast_to(self.trainable_cond_mask.weight, dtype=x.dtype, device=x.device).unsqueeze(1).unsqueeze(1)
x = x + cond_mask_weight[0]
x = _append_context_latents(self, x, kwargs)
if reference_latent is not None:
ref = self.patch_embedding(reference_latent.float()).to(x.dtype)
ref = ref.flatten(2).transpose(1, 2)
freqs_ref = self.rope_encode(reference_latent.shape[-3], reference_latent.shape[-2], reference_latent.shape[-1], t_start=max(30, time + 9), device=x.device, dtype=x.dtype)
ref = ref + cond_mask_weight[1]
x = torch.cat([x, ref], dim=1)
freqs = torch.cat([freqs, freqs_ref], dim=1)
t = torch.cat([t, torch.zeros((t.shape[0], reference_latent.shape[-3]), device=t.device, dtype=t.dtype)], dim=1)
if reference_motion is not None:
motion_encoded, freqs_motion = self.frame_packer(reference_motion, self)
motion_encoded = motion_encoded + cond_mask_weight[2]
x = torch.cat([x, motion_encoded], dim=1)
freqs = torch.cat([freqs, freqs_motion], dim=1)
t = torch.repeat_interleave(t, 2, dim=1)
t = torch.cat([t, torch.zeros((t.shape[0], 3), device=t.device, dtype=t.dtype)], dim=1)
from comfy.ldm.wan.model import sinusoidal_embedding_1d
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
e = e.reshape(t.shape[0], -1, e.shape[-1])
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
context = self.text_embedding(context)
patches_replace = transformer_options.get("patches_replace", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
for i, block in enumerate(self.blocks):
transformer_options["block_index"] = i
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], transformer_options=args["transformer_options"])
return out
out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
x = block(x, e=e0, freqs=freqs, context=context, transformer_options=transformer_options)
if audio_emb is not None:
inject_scale = kwargs.get("audio_inject_scale", 1.0)
if isinstance(inject_scale, torch.Tensor):
inject_scale = inject_scale.reshape(-1)[0].item()
x = self.audio_injector(
x, i, audio_emb, audio_emb_global, seq_len,
scale=inject_scale,
token_mask=kwargs.get("audio_inject_mask", None),
)
x = self.head(x, e)
x = self.unpatchify(x, grid_sizes)
return x
forward_orig.__wan_bernini_s2v_v2_patch__ = True
forward_orig.__wan_bernini_s2v_patch__ = True
forward_orig.__wan_bernini_s2v_original__ = original
WanModel_S2V.forward_orig = forward_orig
def _patch_audio_injector():
if getattr(AudioInjector_WAN.forward, "__wan_bernini_s2v_v2_masked_patch__", False):
return
original_forward = AudioInjector_WAN.forward
def forward(self, x, block_id, audio_emb, audio_emb_global, seq_len, scale=1.0, token_mask=None):
if token_mask is None:
return original_forward(self, x, block_id, audio_emb, audio_emb_global, seq_len, scale=scale)
audio_attn_id = self.injected_block_id.get(block_id, None)
if audio_attn_id is None:
return x
from einops import rearrange
num_frames = audio_emb.shape[1]
input_hidden_states = rearrange(x[:, :seq_len], "b (t n) c -> (b t) n c", t=num_frames)
if self.enable_adain and self.adain_mode == "attn_norm":
audio_emb_global = rearrange(audio_emb_global, "b t n c -> (b t) n c")
adain_hidden_states = self.injector_adain_layers[audio_attn_id](input_hidden_states, temb=audio_emb_global[:, 0])
attn_hidden_states = adain_hidden_states
else:
attn_hidden_states = self.injector_pre_norm_feat[audio_attn_id](input_hidden_states)
if audio_emb.dim() == 3:
attn_audio_emb = rearrange(audio_emb, "b t c -> (b t) 1 c", t=num_frames)
else:
attn_audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c", t=num_frames)
residual_out = self.injector[audio_attn_id](x=attn_hidden_states, context=attn_audio_emb)
residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames)
if token_mask.ndim == 4:
token_mask = token_mask.flatten(1, 2)
if token_mask.shape[1] == residual_out.shape[1]:
residual_out = residual_out * token_mask.to(device=residual_out.device, dtype=residual_out.dtype)
else:
logging.warning(
"ComfyUI-WanBerniniS2V_v2: mask length %s does not match token count %s; using global audio injection",
token_mask.shape[1],
residual_out.shape[1],
)
x[:, :seq_len] = x[:, :seq_len] + residual_out * scale
return x
forward.__wan_bernini_s2v_v2_masked_patch__ = True
forward.__wan_bernini_s2v_masked_patch__ = True
forward.__wan_bernini_s2v_masked_original__ = original_forward
AudioInjector_WAN.forward = forward
def _patch_wan22_s2v_extra_conds():
if getattr(WAN22_S2V.extra_conds, "__wan_bernini_s2v_v2_masked_patch__", False):
return
original_extra_conds = WAN22_S2V.extra_conds
def extra_conds(self, **kwargs):
out = original_extra_conds(self, **kwargs)
audio_inject_mask = kwargs.get("audio_inject_mask", None)
if audio_inject_mask is not None:
out["audio_inject_mask"] = comfy.conds.CONDRegular(audio_inject_mask)
audio_inject_scale = kwargs.get("audio_inject_scale", None)
if audio_inject_scale is not None:
out["audio_inject_scale"] = comfy.conds.CONDRegular(torch.FloatTensor([audio_inject_scale]))
return out
extra_conds.__wan_bernini_s2v_v2_masked_patch__ = True
extra_conds.__wan_bernini_s2v_masked_patch__ = True
extra_conds.__wan_bernini_s2v_masked_original__ = original_extra_conds
WAN22_S2V.extra_conds = extra_conds
original_resize = WAN22_S2V.resize_cond_for_context_window
def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]):
if cond_key == "audio_inject_mask":
mask = cond_value.cond
if mask.ndim == 4 and mask.shape[1] == x_in.shape[2]:
return cond_value._copy_with(window.get_tensor(mask, device, dim=1))
return original_resize(self, cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
resize_cond_for_context_window.__wan_bernini_s2v_v2_masked_patch__ = True
resize_cond_for_context_window.__wan_bernini_s2v_masked_patch__ = True
resize_cond_for_context_window.__wan_bernini_s2v_masked_original__ = original_resize
WAN22_S2V.resize_cond_for_context_window = resize_cond_for_context_window
def apply_model_patches():
_patch_wan_model_s2v_forward()
_patch_audio_injector()
_patch_wan22_s2v_extra_conds()
logging.info("ComfyUI-WanBerniniS2V_v2: applied Bernini S2V model patches") |