diff --git a/comfy/ldm/lightricks/av_model.py b/comfy/ldm/lightricks/av_model.py index ef9938465..8e360f6a8 100644 --- a/comfy/ldm/lightricks/av_model.py +++ b/comfy/ldm/lightricks/av_model.py @@ -16,7 +16,9 @@ from comfy.ldm.lightricks.model import ( from comfy.ldm.lightricks.symmetric_patchifier import AudioPatchifier from comfy.ldm.lightricks.embeddings_connector import Embeddings1DConnector import comfy.ldm.common_dit +import comfy.model_management import comfy.model_prefetch +import comfy.quant_ops class CompressedTimestep: """Store video timestep embeddings in compressed form using per-frame indexing.""" @@ -271,7 +273,10 @@ class BasicAVTransformerBlock(nn.Module): if run_vx: # video self-attention vshift_msa, vscale_msa = (self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(0, 2))) - norm_vx = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_msa) + vshift_msa + if comfy.model_management.in_training: + norm_vx = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_msa) + vshift_msa + else: + norm_vx = comfy.quant_ops.ck.rms_adaln(vx, vscale_msa, vshift_msa) del vshift_msa, vscale_msa attn1_out = self.attn1(norm_vx, pe=v_pe, mask=self_attention_mask, transformer_options=transformer_options) del norm_vx @@ -305,7 +310,6 @@ class BasicAVTransformerBlock(nn.Module): # video - audio cross attention. if run_a2v or run_v2a: - vx_norm3 = comfy.ldm.common_dit.rms_norm(vx) ax_norm3 = comfy.ldm.common_dit.rms_norm(ax) # audio to video cross attention @@ -315,7 +319,10 @@ class BasicAVTransformerBlock(nn.Module): scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v = self.get_ada_values( self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[:2] - vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v_v) + shift_ca_video_hidden_states_a2v_v + if comfy.model_management.in_training: + vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + scale_ca_video_hidden_states_a2v_v) + shift_ca_video_hidden_states_a2v_v + else: + vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v) ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v del scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v, scale_ca_audio_hidden_states_a2v, shift_ca_audio_hidden_states_a2v @@ -334,7 +341,10 @@ class BasicAVTransformerBlock(nn.Module): self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[2:4] ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a - vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a + if comfy.model_management.in_training: + vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a + else: + vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a) del scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a, scale_ca_audio_hidden_states_v2a, shift_ca_audio_hidden_states_v2a v2a_out = self.video_to_audio_attn(ax_scaled, context=vx_scaled, pe=a_cross_pe, k_pe=v_cross_pe, transformer_options=transformer_options) @@ -344,12 +354,14 @@ class BasicAVTransformerBlock(nn.Module): ax.addcmul_(v2a_out, gate_out_v2a) del gate_out_v2a, v2a_out - del vx_norm3, ax_norm3 # video feedforward if run_vx: vshift_mlp, vscale_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(3, 5)) - vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_mlp) + vshift_mlp + if comfy.model_management.in_training: + vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_mlp) + vshift_mlp + else: + vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, vscale_mlp, vshift_mlp) del vshift_mlp, vscale_mlp ff_out = self.ff(vx_scaled) diff --git a/comfy/ldm/lightricks/model.py b/comfy/ldm/lightricks/model.py index f9de3a38e..f80bffba7 100644 --- a/comfy/ldm/lightricks/model.py +++ b/comfy/ldm/lightricks/model.py @@ -13,6 +13,7 @@ import comfy.patcher_extension import comfy.ldm.modules.attention import comfy.ldm.common_dit import comfy.model_management +import comfy.ops import comfy.quant_ops from .symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords @@ -321,7 +322,11 @@ class FeedForward(nn.Module): ) def forward(self, x): - return self.net(x) + # net = [GELU_approx(proj), Dropout, Linear]; the fused path skips the + # Dropout, so leave it to the stock path whenever it could be active. + if comfy.model_management.in_training: + return self.net(x) + return comfy.ops.linear_input_act(self.net[2], self.net[0].proj(x), "gelu_tanh") def apply_rotary_emb(input_tensor, freqs_cis): rotation_matrix, split_pe = freqs_cis @@ -535,7 +540,12 @@ class BasicTransformerBlock(nn.Module): def forward(self, x, context=None, attention_mask=None, timestep=None, pe=None, transformer_options={}, self_attention_mask=None, prompt_timestep=None): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None, None, :6].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)[:, :, :6, :]).unbind(dim=2) - x += self.attn1(comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa, pe=pe, mask=self_attention_mask, transformer_options=transformer_options) * gate_msa + if comfy.model_management.in_training: + norm_x = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa + else: + norm_x = comfy.quant_ops.ck.rms_adaln(x, scale_msa, shift_msa) + + x += self.attn1(norm_x, pe=pe, mask=self_attention_mask, transformer_options=transformer_options) * gate_msa if self.cross_attention_adaln: shift_q_mca, scale_q_mca, gate_mca = (self.scale_shift_table[None, None, 6:9].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)[:, :, 6:9, :]).unbind(dim=2) @@ -589,7 +599,10 @@ def apply_cross_attention_adaln( prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1) ).unbind(dim=2) - attn_input = comfy.ldm.common_dit.rms_norm(x) * (1 + q_scale) + q_shift + if comfy.model_management.in_training: + attn_input = comfy.ldm.common_dit.rms_norm(x) * (1 + q_scale) + q_shift + else: + attn_input = comfy.quant_ops.ck.rms_adaln(x, q_scale, q_shift) encoder_hidden_states = context * (1 + scale_kv) + shift_kv return attn(attn_input, context=encoder_hidden_states, mask=attention_mask, transformer_options=transformer_options) * q_gate diff --git a/comfy/ldm/wan/model.py b/comfy/ldm/wan/model.py index c042e93c4..dca6efba1 100644 --- a/comfy/ldm/wan/model.py +++ b/comfy/ldm/wan/model.py @@ -11,6 +11,7 @@ from comfy.ldm.flux.layers import EmbedND from comfy.ldm.flux.math import apply_rope1, rope import comfy.ldm.common_dit import comfy.model_management +import comfy.ops import comfy.patcher_extension @@ -174,6 +175,13 @@ def repeat_e(e, x): return torch.repeat_interleave(e, repeats + 1, dim=1)[:, :x.size(1)] +class WanFeedForward(nn.Sequential): + """[Linear, GELU(tanh), Linear], with the GELU folded into the down-projection.""" + + def forward(self, x): + return comfy.ops.linear_input_act(self[2], self[0](x), "gelu_tanh") + + class WanAttentionBlock(nn.Module): def __init__(self, @@ -207,7 +215,7 @@ class WanAttentionBlock(nn.Module): qk_norm, eps, operation_settings=operation_settings) self.norm2 = operation_settings.get("operations").LayerNorm(dim, eps, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) - self.ffn = nn.Sequential( + self.ffn = WanFeedForward( operation_settings.get("operations").Linear(dim, ffn_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")), nn.GELU(approximate='tanh'), operation_settings.get("operations").Linear(ffn_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))) diff --git a/comfy/ldm/wan/uni3c.py b/comfy/ldm/wan/uni3c.py index 827ad2339..f4bf68200 100644 --- a/comfy/ldm/wan/uni3c.py +++ b/comfy/ldm/wan/uni3c.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn from comfy.ldm.flux.layers import EmbedND -from .model import WanSelfAttention +from .model import WanFeedForward, WanSelfAttention class Uni3CLayerNormZero(nn.Module): @@ -41,7 +41,7 @@ class Uni3CAttentionBlock(nn.Module): self.norm1 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations) self.self_attn = WanSelfAttention(dim, num_heads, qk_norm=True, eps=eps, operation_settings=operation_settings) self.norm2 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations) - self.ffn = nn.Sequential( + self.ffn = WanFeedForward( operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'), operations.Linear(ffn_dim, dim, device=device, dtype=dtype))