258 lines
8.6 KiB
Python
258 lines
8.6 KiB
Python
import math
|
|
|
|
import comfy.ldm.common_dit
|
|
import comfy.ldm.modules.attention
|
|
import torch
|
|
from comfy.ldm.lightricks.model import (
|
|
BasicTransformerBlock,
|
|
LTXVModel,
|
|
apply_rotary_emb,
|
|
)
|
|
from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords
|
|
from torch import nn
|
|
|
|
from ..utils.feta_enhance_utils import get_feta_scores
|
|
|
|
|
|
class LTXModifiedCrossAttention(nn.Module):
|
|
def forward(self, x, context=None, mask=None, pe=None, transformer_options={}):
|
|
context = x if context is None else context
|
|
context_v = x if context is None else context
|
|
|
|
step = transformer_options.get("step", -1)
|
|
total_steps = transformer_options.get("total_steps", 0)
|
|
attn_bank = transformer_options.get("attn_bank", None)
|
|
sample_mode = transformer_options.get("sample_mode", None)
|
|
if attn_bank is not None and self.idx in attn_bank["block_map"]:
|
|
len_conds = len(transformer_options["cond_or_uncond"])
|
|
pred_order = transformer_options["pred_order"]
|
|
if (
|
|
sample_mode == "forward"
|
|
and total_steps - step - 1 < attn_bank["save_steps"]
|
|
):
|
|
step_idx = f"{pred_order}_{total_steps-step-1}"
|
|
attn_bank["block_map"][self.idx][step_idx] = x.cpu()
|
|
elif sample_mode == "reverse" and step < attn_bank["inject_steps"]:
|
|
step_idx = f"{pred_order}_{step}"
|
|
inject_settings = attn_bank.get("inject_settings", {})
|
|
if len(inject_settings) > 0:
|
|
inj = (
|
|
attn_bank["block_map"][self.idx][step_idx]
|
|
.to(x.device)
|
|
.repeat(len_conds, 1, 1)
|
|
)
|
|
if "q" in inject_settings:
|
|
x = inj
|
|
if "k" in inject_settings:
|
|
context = inj
|
|
if "v" in inject_settings:
|
|
context_v = inj
|
|
|
|
q = self.to_q(x)
|
|
k = self.to_k(context)
|
|
v = self.to_v(context_v)
|
|
|
|
q = self.q_norm(q)
|
|
k = self.k_norm(k)
|
|
|
|
if pe is not None:
|
|
q = apply_rotary_emb(q, pe)
|
|
k = apply_rotary_emb(k, pe)
|
|
|
|
feta_score = None
|
|
if (
|
|
transformer_options.get("feta_weight", 0) > 0
|
|
and self.idx in transformer_options["feta_layers"]["layers"]
|
|
):
|
|
feta_score = get_feta_scores(q, k, self.heads, transformer_options)
|
|
|
|
alt_attn_fn = (
|
|
transformer_options.get("patches_replace", {})
|
|
.get("layer", {})
|
|
.get(("self_attn", self.idx), None)
|
|
)
|
|
if alt_attn_fn is not None:
|
|
out = alt_attn_fn(
|
|
q,
|
|
k,
|
|
v,
|
|
self.heads,
|
|
attn_precision=self.attn_precision,
|
|
transformer_options=transformer_options,
|
|
)
|
|
elif mask is None:
|
|
out = comfy.ldm.modules.attention.optimized_attention(
|
|
q, k, v, self.heads, attn_precision=self.attn_precision
|
|
)
|
|
else:
|
|
out = comfy.ldm.modules.attention.optimized_attention_masked(
|
|
q, k, v, self.heads, mask, attn_precision=self.attn_precision
|
|
)
|
|
|
|
if feta_score is not None:
|
|
out *= feta_score
|
|
|
|
return self.to_out(out)
|
|
|
|
|
|
class LTXModifiedBasicTransformerBlock(BasicTransformerBlock):
|
|
def forward(
|
|
self,
|
|
x,
|
|
context=None,
|
|
attention_mask=None,
|
|
timestep=None,
|
|
pe=None,
|
|
transformer_options={},
|
|
):
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
|
self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
|
+ timestep.reshape(
|
|
x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1
|
|
)
|
|
).unbind(dim=2)
|
|
x += (
|
|
self.attn1(
|
|
comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa,
|
|
pe=pe,
|
|
transformer_options=transformer_options,
|
|
)
|
|
* gate_msa
|
|
)
|
|
|
|
x += self.attn2(x, context=context, mask=attention_mask)
|
|
|
|
y = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_mlp) + shift_mlp
|
|
x += self.ff(y) * gate_mlp
|
|
|
|
return x
|
|
|
|
|
|
class LTXVModelModified(LTXVModel):
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
timestep,
|
|
context,
|
|
attention_mask,
|
|
frame_rate=25,
|
|
transformer_options={},
|
|
keyframe_idxs=None,
|
|
**kwargs,
|
|
):
|
|
patches_replace = transformer_options.get("patches_replace", {})
|
|
|
|
orig_shape = list(x.shape)
|
|
|
|
x, latent_coords = self.patchifier.patchify(x)
|
|
pixel_coords = latent_to_pixel_coords(
|
|
latent_coords=latent_coords,
|
|
scale_factors=self.vae_scale_factors,
|
|
causal_fix=self.causal_temporal_positioning,
|
|
)
|
|
|
|
if keyframe_idxs is not None:
|
|
pixel_coords[:, :, -keyframe_idxs.shape[2] :] = keyframe_idxs
|
|
|
|
fractional_coords = pixel_coords.to(torch.float32)
|
|
fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
|
|
|
|
x = self.patchify_proj(x)
|
|
timestep = timestep * 1000.0
|
|
|
|
if attention_mask is not None and not torch.is_floating_point(attention_mask):
|
|
attention_mask = (attention_mask - 1).to(x.dtype).reshape(
|
|
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
|
) * torch.finfo(x.dtype).max
|
|
|
|
pe = self._precompute_freqs_cis(
|
|
fractional_coords, dim=self.inner_dim, out_dtype=x.dtype
|
|
)
|
|
|
|
batch_size = x.shape[0]
|
|
timestep, embedded_timestep = self.adaln_single(
|
|
timestep.flatten(),
|
|
{"resolution": None, "aspect_ratio": None},
|
|
batch_size=batch_size,
|
|
hidden_dtype=x.dtype,
|
|
)
|
|
# Second dimension is 1 or number of tokens (if timestep_per_token)
|
|
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
|
|
embedded_timestep = embedded_timestep.view(
|
|
batch_size, -1, embedded_timestep.shape[-1]
|
|
)
|
|
|
|
# 2. Blocks
|
|
if self.caption_projection is not None:
|
|
batch_size = x.shape[0]
|
|
context = self.caption_projection(context)
|
|
context = context.view(batch_size, -1, x.shape[-1])
|
|
|
|
blocks_replace = patches_replace.get("dit", {})
|
|
for i, block in enumerate(self.transformer_blocks):
|
|
if ("double_block", i) in blocks_replace:
|
|
|
|
def block_wrap(args):
|
|
out = {}
|
|
out["img"] = block(
|
|
args["img"],
|
|
context=args["txt"],
|
|
attention_mask=args["attention_mask"],
|
|
timestep=args["vec"],
|
|
pe=args["pe"],
|
|
)
|
|
return out
|
|
|
|
out = blocks_replace[("double_block", i)](
|
|
{
|
|
"img": x,
|
|
"txt": context,
|
|
"attention_mask": attention_mask,
|
|
"vec": timestep,
|
|
"pe": pe,
|
|
},
|
|
{"original_block": block_wrap},
|
|
)
|
|
x = out["img"]
|
|
else:
|
|
x = block(
|
|
x,
|
|
context=context,
|
|
attention_mask=attention_mask,
|
|
timestep=timestep,
|
|
pe=pe,
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
# 3. Output
|
|
scale_shift_values = (
|
|
self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
|
+ embedded_timestep[:, :, None]
|
|
)
|
|
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
|
|
x = self.norm_out(x)
|
|
# Modulation
|
|
x = x * (1 + scale) + shift
|
|
x = self.proj_out(x)
|
|
|
|
x = self.patchifier.unpatchify(
|
|
latents=x,
|
|
output_height=orig_shape[3],
|
|
output_width=orig_shape[4],
|
|
output_num_frames=orig_shape[2],
|
|
out_channels=orig_shape[1] // math.prod(self.patchifier.patch_size),
|
|
)
|
|
|
|
return x
|
|
|
|
|
|
def inject_model(diffusion_model):
|
|
diffusion_model.__class__ = LTXVModelModified
|
|
for idx, transformer_block in enumerate(diffusion_model.transformer_blocks):
|
|
transformer_block.__class__ = LTXModifiedBasicTransformerBlock
|
|
transformer_block.idx = idx
|
|
transformer_block.attn1.__class__ = LTXModifiedCrossAttention
|
|
transformer_block.attn1.idx = idx
|
|
return diffusion_model
|