Speedup LTX and Wan (#15138)
This commit is contained in:
Jukka Seppänen
GitHub
parent
bbda83647d
commit
15989f87ca
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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")))
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue