312 lines
12 KiB
Python
312 lines
12 KiB
Python
from typing import List
|
|
|
|
import comfy.samplers
|
|
import torch
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy.patcher_extension import CallbacksMP
|
|
|
|
from ..nodes_registry import comfy_node
|
|
from ..stg import STGBlockWrapper, STGFlag
|
|
from .parameters import GuiderParameters, Modality
|
|
|
|
|
|
class MultimodalGuider(comfy.samplers.CFGGuider):
|
|
def __init__(
|
|
self,
|
|
model: ModelPatcher,
|
|
parameters: dict[str, GuiderParameters],
|
|
skip_blocks: List[int],
|
|
):
|
|
model = model.clone()
|
|
self.current_step = 0
|
|
self.last_denoised_v = None
|
|
self.last_denoised_a = None
|
|
|
|
model.add_callback_with_key(
|
|
CallbacksMP.ON_PRE_RUN, "mm_guider_on_pre_run", self.reset_current_step
|
|
)
|
|
super().__init__(model)
|
|
|
|
self.stg_flag = STGFlag(
|
|
do_skip=False,
|
|
skip_layers=skip_blocks,
|
|
)
|
|
|
|
self.patch_model(model, self.stg_flag)
|
|
self.parameters = parameters
|
|
|
|
def reset_current_step(self, model_patcher=None):
|
|
self.current_step = 0
|
|
self.last_denoised_v = None
|
|
self.last_denoised_a = None
|
|
|
|
@classmethod
|
|
def patch_model(cls, model: ModelPatcher, stg_flag: STGFlag):
|
|
transformer_blocks = cls.get_transformer_blocks(model)
|
|
|
|
for i, block in enumerate(transformer_blocks):
|
|
model.set_model_patch_replace(
|
|
STGBlockWrapper(block, stg_flag, i), "dit", "double_block", i
|
|
)
|
|
|
|
@staticmethod
|
|
def get_transformer_blocks(model: ModelPatcher):
|
|
diffusion_model = model.get_model_object("diffusion_model")
|
|
key = "diffusion_model.transformer_blocks"
|
|
if diffusion_model.__class__.__name__ == "LTXVTransformer3D":
|
|
key = "diffusion_model.transformer.transformer_blocks"
|
|
return model.get_model_object(key)
|
|
|
|
def set_conds(self, positive, negative):
|
|
self.inner_set_conds({"positive": positive, "negative": negative})
|
|
|
|
def calc_stg_indexes(
|
|
self,
|
|
run_vx: bool,
|
|
run_ax: bool,
|
|
audio_ptb: bool,
|
|
video_ptb: bool,
|
|
):
|
|
stg_indexes = set()
|
|
num_self_attns = 0
|
|
if run_vx:
|
|
num_self_attns += 1
|
|
if run_ax:
|
|
num_self_attns += 1
|
|
|
|
video_attn_idx = 0
|
|
audio_attn_idx = 0 if num_self_attns == 1 else 2
|
|
|
|
if video_ptb:
|
|
stg_indexes.add(video_attn_idx)
|
|
if audio_ptb:
|
|
stg_indexes.add(audio_attn_idx)
|
|
|
|
return list(stg_indexes)
|
|
|
|
def unpack_latents(self, x: torch.Tensor):
|
|
latent_shapes = (
|
|
self.conds.get("positive", {})[0]
|
|
.get("model_conds", {})
|
|
.get("latent_shapes", None)
|
|
.cond
|
|
)
|
|
vx, ax = comfy.utils.unpack_latents(x, latent_shapes)
|
|
return vx, ax
|
|
|
|
def pack_latents(self, vx: torch.Tensor, ax: torch.Tensor):
|
|
x, latent_shapes = comfy.utils.pack_latents([vx, ax])
|
|
return x, latent_shapes
|
|
|
|
def predict_noise(
|
|
self,
|
|
x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
model_options: dict = {},
|
|
seed=None,
|
|
):
|
|
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
|
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
|
current_step = self.current_step
|
|
self.current_step = current_step + 1
|
|
|
|
positive_cond = self.conds.get("positive", None)
|
|
negative_cond = self.conds.get("negative", None)
|
|
|
|
audio_params: GuiderParameters = self.parameters.get(
|
|
Modality.AUDIO.value, GuiderParameters()
|
|
)
|
|
video_params: GuiderParameters = self.parameters.get(
|
|
Modality.VIDEO.value, GuiderParameters()
|
|
)
|
|
|
|
run_vx = not video_params.do_skip(current_step)
|
|
run_ax = not audio_params.do_skip(current_step)
|
|
run_a2v = video_params.do_cross_attn(current_step)
|
|
run_v2a = audio_params.do_cross_attn(current_step)
|
|
|
|
vx, ax = self.unpack_latents(x)
|
|
if not run_vx:
|
|
vx = self.last_denoised_v
|
|
if not run_ax:
|
|
ax = self.last_denoised_a
|
|
x, _ = self.pack_latents(vx, ax)
|
|
|
|
if not run_vx and not run_ax:
|
|
return x
|
|
|
|
try:
|
|
model_options["transformer_options"]["run_vx"] = run_vx
|
|
model_options["transformer_options"]["run_ax"] = run_ax
|
|
model_options["transformer_options"]["a2v_cross_attn"] = run_a2v
|
|
model_options["transformer_options"]["v2a_cross_attn"] = run_v2a
|
|
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
finally:
|
|
del model_options["transformer_options"]["run_vx"]
|
|
del model_options["transformer_options"]["run_ax"]
|
|
del model_options["transformer_options"]["a2v_cross_attn"]
|
|
del model_options["transformer_options"]["v2a_cross_attn"]
|
|
v_noise_pred_pos, a_noise_pred_pos = self.unpack_latents(noise_pred_pos)
|
|
|
|
a_noise_pred_neg, v_noise_pred_neg = 0, 0
|
|
a_noise_pred_perturbed, v_noise_pred_perturbed = 0, 0
|
|
a_noise_pred_modality, v_noise_pred_modality = 0, 0
|
|
|
|
if any(params.do_uncond() for params in [audio_params, video_params]):
|
|
try:
|
|
model_options["transformer_options"]["run_vx"] = run_vx
|
|
model_options["transformer_options"]["run_ax"] = run_ax
|
|
model_options["transformer_options"]["a2v_cross_attn"] = run_a2v
|
|
model_options["transformer_options"]["v2a_cross_attn"] = run_v2a
|
|
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[negative_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
v_noise_pred_neg, a_noise_pred_neg = self.unpack_latents(noise_pred_neg)
|
|
finally:
|
|
del model_options["transformer_options"]["run_vx"]
|
|
del model_options["transformer_options"]["run_ax"]
|
|
del model_options["transformer_options"]["a2v_cross_attn"]
|
|
del model_options["transformer_options"]["v2a_cross_attn"]
|
|
|
|
if any(params.do_perturbed() for params in [audio_params, video_params]):
|
|
try:
|
|
stg_indexes = self.calc_stg_indexes(
|
|
run_vx,
|
|
run_ax and ax.numel() > 0,
|
|
audio_params.perturb_attn,
|
|
video_params.perturb_attn,
|
|
)
|
|
model_options["transformer_options"]["run_vx"] = run_vx
|
|
model_options["transformer_options"]["run_ax"] = run_ax
|
|
model_options["transformer_options"]["a2v_cross_attn"] = run_a2v
|
|
model_options["transformer_options"]["v2a_cross_attn"] = run_v2a
|
|
model_options["transformer_options"]["ptb_index"] = 0
|
|
model_options["transformer_options"]["stg_indexes"] = stg_indexes
|
|
self.stg_flag.do_skip = True
|
|
noise_pred_perturbed = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
v_noise_pred_perturbed, a_noise_pred_perturbed = self.unpack_latents(
|
|
noise_pred_perturbed
|
|
)
|
|
finally:
|
|
self.stg_flag.do_skip = False
|
|
del model_options["transformer_options"]["ptb_index"]
|
|
del model_options["transformer_options"]["run_vx"]
|
|
del model_options["transformer_options"]["run_ax"]
|
|
del model_options["transformer_options"]["a2v_cross_attn"]
|
|
del model_options["transformer_options"]["v2a_cross_attn"]
|
|
del model_options["transformer_options"]["stg_indexes"]
|
|
|
|
if any(params.do_modality() for params in [audio_params, video_params]):
|
|
try:
|
|
model_options["transformer_options"]["run_vx"] = run_vx
|
|
model_options["transformer_options"]["run_ax"] = run_ax
|
|
model_options["transformer_options"]["a2v_cross_attn"] = False
|
|
model_options["transformer_options"]["v2a_cross_attn"] = False
|
|
noise_pred_modality = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
v_noise_pred_modality, a_noise_pred_modality = self.unpack_latents(
|
|
noise_pred_modality
|
|
)
|
|
finally:
|
|
del model_options["transformer_options"]["a2v_cross_attn"]
|
|
del model_options["transformer_options"]["v2a_cross_attn"]
|
|
del model_options["transformer_options"]["run_vx"]
|
|
del model_options["transformer_options"]["run_ax"]
|
|
|
|
if run_vx:
|
|
vx = video_params.calculate(
|
|
v_noise_pred_pos,
|
|
v_noise_pred_neg,
|
|
v_noise_pred_perturbed,
|
|
v_noise_pred_modality,
|
|
)
|
|
else:
|
|
vx = self.last_denoised_v
|
|
|
|
if run_ax:
|
|
ax = audio_params.calculate(
|
|
a_noise_pred_pos,
|
|
a_noise_pred_neg,
|
|
a_noise_pred_perturbed,
|
|
a_noise_pred_modality,
|
|
)
|
|
else:
|
|
ax = self.last_denoised_a
|
|
|
|
x, _ = self.pack_latents(vx, ax)
|
|
|
|
# normally this would be done in cfg_function, but we skipped
|
|
# that for efficiency: we can compute the noise predictions in
|
|
# a single call to calc_cond_batch() (rather than two)
|
|
# so we replicate the hook here
|
|
for fn in model_options.get("sampler_post_cfg_function", []):
|
|
args = {
|
|
"denoised": x,
|
|
"cond": positive_cond,
|
|
"uncond": negative_cond,
|
|
"model": self.inner_model,
|
|
"uncond_denoised": noise_pred_neg,
|
|
"cond_denoised": noise_pred_pos,
|
|
"sigma": timestep,
|
|
"model_options": model_options,
|
|
"input": x,
|
|
# not in the original call in samplers.py:cfg_function, but made available for future hooks
|
|
"perturbed_cond": positive_cond,
|
|
"perturbed_cond_denoised": noise_pred_perturbed,
|
|
}
|
|
x = fn(args)
|
|
|
|
self.last_denoised_v, self.last_denoised_a = self.unpack_latents(x)
|
|
return x
|
|
|
|
|
|
@comfy_node(name="MultimodalGuider")
|
|
class MultimodalGuiderNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"parameters": ("GUIDER_PARAMETERS",),
|
|
"skip_blocks": (
|
|
"STRING",
|
|
{"default": "", "multiline": True},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("GUIDER",)
|
|
|
|
FUNCTION = "get_guider"
|
|
CATEGORY = "lightricks/LTXV"
|
|
|
|
def get_guider(self, model, positive, negative, parameters, skip_blocks):
|
|
skip_blocks = [int(n.strip()) for n in skip_blocks.split(",") if n.strip()]
|
|
guider = MultimodalGuider(model, parameters, skip_blocks)
|
|
guider.set_conds(positive, negative)
|
|
guider.raw_conds = (positive, negative)
|
|
return (guider,)
|