1052 lines
38 KiB
Python
1052 lines
38 KiB
Python
import copy
|
|
|
|
import comfy
|
|
import comfy_extras
|
|
import nodes
|
|
import torch
|
|
from comfy.ldm.lightricks.av_model import LTXAVModel
|
|
from comfy.nested_tensor import NestedTensor
|
|
from comfy_api.latest import io
|
|
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced, SplitSigmas
|
|
from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide, LTXVCropGuides
|
|
|
|
from .guide import blur_internal
|
|
from .latent_norm import LTXVAdainLatent
|
|
from .latents import LTXVSelectLatents, add_latent_guide
|
|
from .nodes_registry import comfy_node
|
|
|
|
|
|
def _get_raw_conds_from_guider(guider):
|
|
if not hasattr(guider, "raw_conds"):
|
|
if "negative" not in guider.original_conds:
|
|
# for BasicGuider
|
|
raise ValueError(
|
|
"Guider does not have negative conds, cannot use it as a guider."
|
|
)
|
|
raw_pos = guider.original_conds["positive"]
|
|
positive = [[raw_pos[0]["cross_attn"], copy.deepcopy(raw_pos[0])]]
|
|
raw_neg = guider.original_conds["negative"]
|
|
negative = [[raw_neg[0]["cross_attn"], copy.deepcopy(raw_neg[0])]]
|
|
guider.raw_conds = (positive, negative)
|
|
return guider.raw_conds
|
|
|
|
|
|
@comfy_node(
|
|
name="LTXVBaseSampler",
|
|
)
|
|
class LTXVBaseSampler:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "The diffusion model to use."}),
|
|
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
|
"width": (
|
|
"INT",
|
|
{
|
|
"default": 768,
|
|
"min": 64,
|
|
"max": nodes.MAX_RESOLUTION,
|
|
"step": 32,
|
|
},
|
|
),
|
|
"height": (
|
|
"INT",
|
|
{
|
|
"default": 512,
|
|
"min": 64,
|
|
"max": nodes.MAX_RESOLUTION,
|
|
"step": 32,
|
|
},
|
|
),
|
|
"num_frames": (
|
|
"INT",
|
|
{"default": 97, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 8},
|
|
),
|
|
"guider": (
|
|
"GUIDER",
|
|
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
|
),
|
|
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
|
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
|
"noise": ("NOISE", {"tooltip": "The noise to use for the sampling."}),
|
|
},
|
|
"optional": {
|
|
"optional_cond_images": (
|
|
"IMAGE",
|
|
{"tooltip": "The images to use for conditioning the sampling."},
|
|
),
|
|
"optional_cond_indices": (
|
|
"STRING",
|
|
{
|
|
"tooltip": "The indices of the images to use for conditioning the sampling."
|
|
},
|
|
),
|
|
"strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.9,
|
|
"min": 0,
|
|
"max": 1,
|
|
"tooltip": "The strength of the conditioning on the images.",
|
|
},
|
|
),
|
|
"crop": (
|
|
["center", "disabled"],
|
|
{
|
|
"default": "disabled",
|
|
"tooltip": "The crop mode to use for the images.",
|
|
},
|
|
),
|
|
"crf": (
|
|
"INT",
|
|
{
|
|
"default": 35,
|
|
"min": 0,
|
|
"max": 100,
|
|
"tooltip": "The CRF value to use for preprocessing the images.",
|
|
},
|
|
),
|
|
"blur": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 10,
|
|
"tooltip": "The blur value to use for preprocessing the images.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("denoised", "positive", "negative")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
vae,
|
|
width,
|
|
height,
|
|
num_frames,
|
|
guider,
|
|
sampler,
|
|
sigmas,
|
|
noise,
|
|
optional_cond_images=None,
|
|
optional_cond_indices=None,
|
|
strength=0.9,
|
|
crop="disabled",
|
|
crf=35,
|
|
blur=0,
|
|
optional_negative_index_latents=None,
|
|
optional_negative_index=-1,
|
|
optional_negative_index_strength=1.0,
|
|
optional_initialization_latents=None,
|
|
guiding_start_step=0,
|
|
guiding_end_step=1000,
|
|
):
|
|
guider = copy.copy(guider)
|
|
guider.original_conds = copy.deepcopy(guider.original_conds)
|
|
positive, negative = _get_raw_conds_from_guider(guider)
|
|
|
|
if optional_cond_images is not None:
|
|
optional_cond_images = (
|
|
comfy.utils.common_upscale(
|
|
optional_cond_images.movedim(-1, 1),
|
|
width,
|
|
height,
|
|
"bilinear",
|
|
crop=crop,
|
|
)
|
|
.movedim(1, -1)
|
|
.clamp(0, 1)
|
|
)
|
|
optional_cond_images = comfy_extras.nodes_lt.LTXVPreprocess.execute(
|
|
optional_cond_images, crf
|
|
)[0]
|
|
for i in range(optional_cond_images.shape[0]):
|
|
optional_cond_images[i] = blur_internal(
|
|
optional_cond_images[i].unsqueeze(0), blur
|
|
)
|
|
|
|
if optional_cond_indices is not None and optional_cond_images is not None:
|
|
optional_cond_indices = optional_cond_indices.split(",")
|
|
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
|
assert len(optional_cond_indices) == len(
|
|
optional_cond_images
|
|
), "Number of optional cond images must match number of optional cond indices"
|
|
|
|
if optional_initialization_latents is None:
|
|
(latents,) = EmptyLTXVLatentVideo().execute(width, height, num_frames, 1)
|
|
else:
|
|
latents = optional_initialization_latents
|
|
|
|
if optional_cond_images is not None and 0 in optional_cond_indices:
|
|
# apply classical i2v conditioning on the first frame
|
|
idx_0 = optional_cond_indices.index(0)
|
|
encode_pixels = optional_cond_images[idx_0 : idx_0 + 1, :, :, :3]
|
|
t = vae.encode(encode_pixels)
|
|
latents["samples"][:, :, : t.shape[2]] = t
|
|
|
|
if "noise_mask" not in latents:
|
|
conditioning_latent_frames_mask = torch.ones(
|
|
(1, 1, latents["samples"].shape[2], 1, 1),
|
|
dtype=torch.float32,
|
|
device=latents["samples"].device,
|
|
)
|
|
conditioning_latent_frames_mask[:, :, : t.shape[2]] = 1.0 - strength
|
|
latents["noise_mask"] = conditioning_latent_frames_mask
|
|
else:
|
|
latents["noise_mask"][:, :, : t.shape[2]] = 1.0 - strength
|
|
conditioning_latent_frames_mask = latents["noise_mask"]
|
|
else:
|
|
conditioning_latent_frames_mask = None
|
|
|
|
high_sigmas, rest_sigmas = SplitSigmas().get_sigmas(sigmas, guiding_start_step)
|
|
middle_sigmas, low_sigmas = SplitSigmas().get_sigmas(
|
|
rest_sigmas, guiding_end_step - guiding_start_step
|
|
)
|
|
|
|
if len(high_sigmas) > 1:
|
|
print("Denoising with no conditioning on sigmas: ", high_sigmas)
|
|
(_, new_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=high_sigmas,
|
|
latent_image=latents,
|
|
)
|
|
|
|
if optional_cond_images is not None:
|
|
# add conditioning on keyframes with index > 0
|
|
for cond_image, cond_idx in zip(
|
|
optional_cond_images, optional_cond_indices
|
|
):
|
|
if cond_idx == 0:
|
|
# 0 is handled by classical i2v conditioning
|
|
continue
|
|
(
|
|
positive,
|
|
negative,
|
|
latents,
|
|
) = LTXVAddGuide.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
vae=vae,
|
|
latent=latents,
|
|
image=cond_image.unsqueeze(0),
|
|
frame_idx=cond_idx,
|
|
strength=strength,
|
|
)
|
|
|
|
if optional_negative_index_latents is not None:
|
|
(
|
|
positive,
|
|
negative,
|
|
latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=latents,
|
|
guiding_latent=optional_negative_index_latents,
|
|
latent_idx=optional_negative_index,
|
|
strength=optional_negative_index_strength,
|
|
)
|
|
|
|
guider.set_conds(positive, negative)
|
|
|
|
# Denoise the latent video
|
|
print("Denoising with conditioning on sigmas: ", middle_sigmas)
|
|
(output_latents, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=middle_sigmas,
|
|
latent_image=latents,
|
|
)
|
|
|
|
# Clean up guides if image conditioning was used
|
|
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=denoised_output_latents,
|
|
)
|
|
|
|
denoised_output_latents["noise_mask"] = conditioning_latent_frames_mask
|
|
|
|
if len(low_sigmas) > 1:
|
|
print(
|
|
"Denoising with no conditioning but with classical i2v noise mask on sigmas: ",
|
|
low_sigmas,
|
|
)
|
|
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=low_sigmas,
|
|
latent_image=denoised_output_latents,
|
|
)
|
|
|
|
return (denoised_output_latents, positive, negative)
|
|
|
|
|
|
@comfy_node(
|
|
name="LTXVExtendSampler",
|
|
)
|
|
class LTXVExtendSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "The diffusion model to use."}),
|
|
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
|
"latents": (
|
|
"LATENT",
|
|
{"tooltip": "The latents of the video to extend."},
|
|
),
|
|
"num_new_frames": (
|
|
"INT",
|
|
{
|
|
"default": 80,
|
|
"min": -1,
|
|
"max": nodes.MAX_RESOLUTION,
|
|
"step": 1,
|
|
"tooltip": "If -1, the number of frames will be based on the number of frames in the optional_guiding_latents.",
|
|
},
|
|
),
|
|
"frame_overlap": (
|
|
"INT",
|
|
{
|
|
"default": 16,
|
|
"min": 16,
|
|
"max": 128,
|
|
"step": 8,
|
|
"tooltip": "The overlap region to use for conditioning the new frames on the end of the provided latents.",
|
|
},
|
|
),
|
|
"guider": (
|
|
"GUIDER",
|
|
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
|
),
|
|
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
|
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
|
"noise": ("NOISE", {"tooltip": "The noise to use for the sampling."}),
|
|
"strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"tooltip": "The strength of the conditioning on the overlapping latents, when using optional_guiding_latents.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"optional_guiding_latents": (
|
|
"LATENT",
|
|
{"tooltip": "Optional latents to guide the sampling."},
|
|
),
|
|
"optional_cond_images": (
|
|
"IMAGE",
|
|
{"tooltip": "The images to use for conditioning the sampling."},
|
|
),
|
|
"optional_cond_indices": ("STRING",),
|
|
"cond_image_strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"tooltip": "The strength of the conditioning on the images.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("denoised_video", "positive", "negative")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
vae,
|
|
latents,
|
|
num_new_frames,
|
|
frame_overlap,
|
|
guider,
|
|
sampler,
|
|
sigmas,
|
|
noise,
|
|
strength=0.5,
|
|
guiding_strength=1.0,
|
|
cond_image_strength=1.0,
|
|
optional_guiding_latents=None,
|
|
optional_cond_images=None,
|
|
optional_cond_indices=None,
|
|
optional_reference_latents=None,
|
|
optional_initialization_latents=None,
|
|
adain_factor=0.0,
|
|
optional_negative_index_latents=None,
|
|
optional_negative_index=-1,
|
|
optional_negative_index_strength=1.0,
|
|
guiding_start_step=0,
|
|
guiding_end_step=1000,
|
|
normalize_per_frame=False,
|
|
):
|
|
guider = copy.copy(guider)
|
|
guider.original_conds = copy.deepcopy(guider.original_conds)
|
|
|
|
if optional_cond_indices is not None and optional_cond_images is not None:
|
|
optional_cond_indices = optional_cond_indices.split(",")
|
|
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
|
assert len(optional_cond_indices) == len(
|
|
optional_cond_images
|
|
), "Number of optional cond images must match number of optional cond indices"
|
|
|
|
positive, negative = _get_raw_conds_from_guider(guider)
|
|
|
|
samples = latents["samples"]
|
|
batch, channels, frames, height, width = samples.shape
|
|
time_scale_factor, width_scale_factor, height_scale_factor = (
|
|
vae.downscale_index_formula
|
|
)
|
|
overlap = frame_overlap // time_scale_factor
|
|
|
|
if num_new_frames == -1 and optional_guiding_latents is not None:
|
|
num_new_frames = (
|
|
optional_guiding_latents["samples"].shape[2] - overlap
|
|
) * time_scale_factor
|
|
|
|
(last_overlap_latents,) = LTXVSelectLatents().select_latents(
|
|
latents, -overlap, -1
|
|
)
|
|
|
|
if optional_initialization_latents is None:
|
|
new_latents = EmptyLTXVLatentVideo.execute(
|
|
width=width * width_scale_factor,
|
|
height=height * height_scale_factor,
|
|
length=overlap * time_scale_factor + num_new_frames,
|
|
batch_size=1,
|
|
)[0]
|
|
else:
|
|
new_latents = optional_initialization_latents
|
|
|
|
last_overlap_latents["samples"] = last_overlap_latents["samples"].to(
|
|
new_latents["samples"].device
|
|
)
|
|
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=new_latents,
|
|
guiding_latent=last_overlap_latents,
|
|
latent_idx=0,
|
|
strength=strength,
|
|
)
|
|
|
|
high_sigmas, rest_sigmas = SplitSigmas().get_sigmas(sigmas, guiding_start_step)
|
|
middle_sigmas, low_sigmas = SplitSigmas().get_sigmas(
|
|
rest_sigmas, guiding_end_step - guiding_start_step
|
|
)
|
|
|
|
if optional_cond_images is not None:
|
|
print("Adding conditioning on keyframes")
|
|
for cond_image, cond_idx in zip(
|
|
optional_cond_images, optional_cond_indices
|
|
):
|
|
if optional_guiding_latents is not None and cond_idx % 8 == 1:
|
|
raise ValueError(
|
|
f"Conditioning image index {cond_idx} (relative to the in the temporal chunk) is a multiple of 8 + 1,"
|
|
"and guiding latents are used. Please provide other conditioning image indices"
|
|
)
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = LTXVAddGuide.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
vae=vae,
|
|
latent=new_latents,
|
|
image=cond_image.unsqueeze(0),
|
|
frame_idx=cond_idx,
|
|
strength=cond_image_strength,
|
|
)
|
|
|
|
if len(high_sigmas) > 1:
|
|
guider.set_conds(positive, negative)
|
|
print("Denoising with overlap conditioning only on sigmas: ", high_sigmas)
|
|
(_, new_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=high_sigmas,
|
|
latent_image=new_latents,
|
|
)
|
|
|
|
if optional_guiding_latents is not None:
|
|
optional_guiding_latents = LTXVSelectLatents().select_latents(
|
|
optional_guiding_latents, overlap, -1
|
|
)[0]
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=new_latents,
|
|
guiding_latent=optional_guiding_latents,
|
|
latent_idx=last_overlap_latents["samples"].shape[2],
|
|
strength=guiding_strength,
|
|
)
|
|
|
|
if optional_negative_index_latents is not None:
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=new_latents,
|
|
guiding_latent=optional_negative_index_latents,
|
|
latent_idx=optional_negative_index,
|
|
strength=optional_negative_index_strength,
|
|
)
|
|
|
|
guider.set_conds(positive, negative)
|
|
|
|
# Denoise the latent video
|
|
print("Denoising with full conditioning on sigmas: ", middle_sigmas)
|
|
(output_latents, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=middle_sigmas,
|
|
latent_image=new_latents,
|
|
)
|
|
|
|
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=denoised_output_latents,
|
|
)
|
|
|
|
if len(low_sigmas) > 1:
|
|
(
|
|
positive,
|
|
negative,
|
|
denoised_output_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=denoised_output_latents,
|
|
guiding_latent=last_overlap_latents,
|
|
latent_idx=0,
|
|
strength=strength,
|
|
)
|
|
|
|
if optional_cond_images is not None:
|
|
print("Adding conditioning on keyframes")
|
|
for cond_image, cond_idx in zip(
|
|
optional_cond_images, optional_cond_indices
|
|
):
|
|
if optional_guiding_latents is not None and cond_idx % 8 == 1:
|
|
raise ValueError(
|
|
f"Conditioning image index {cond_idx} (relative to the in the temporal chunk) is a multiple of 8 + 1,"
|
|
"and guiding latents are used. Please provide other conditioning image indices"
|
|
)
|
|
(
|
|
positive,
|
|
negative,
|
|
denoised_output_latents,
|
|
) = LTXVAddGuide.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
vae=vae,
|
|
latent=denoised_output_latents,
|
|
image=cond_image.unsqueeze(0),
|
|
frame_idx=cond_idx,
|
|
strength=cond_image_strength,
|
|
)
|
|
|
|
guider.set_conds(positive, negative)
|
|
print(
|
|
"Denoising with overlap + keyframes conditioning only on sigmas: ",
|
|
low_sigmas,
|
|
)
|
|
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=low_sigmas,
|
|
latent_image=denoised_output_latents,
|
|
)
|
|
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=denoised_output_latents,
|
|
)
|
|
|
|
if optional_reference_latents is not None:
|
|
denoised_output_latents = LTXVAdainLatent().batch_normalize(
|
|
latents=denoised_output_latents,
|
|
reference=optional_reference_latents,
|
|
factor=adain_factor,
|
|
per_frame=normalize_per_frame,
|
|
)[0]
|
|
|
|
# drop first output latent as it's a reinterpreted 8-frame latent understood as a 1-frame latent
|
|
truncated_denoised_output_latents = LTXVSelectLatents().select_latents(
|
|
denoised_output_latents, 1, -1
|
|
)[0]
|
|
|
|
# Fuse new frames with old ones by calling LinearOverlapLatentTransition
|
|
(latents,) = LinearOverlapLatentTransition().process(
|
|
latents, truncated_denoised_output_latents, overlap - 1, axis=2
|
|
)
|
|
return (latents, positive, negative)
|
|
|
|
|
|
@comfy_node(
|
|
name="LTXVInContextSampler",
|
|
)
|
|
class LTXVInContextSampler:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
|
"guider": (
|
|
"GUIDER",
|
|
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
|
),
|
|
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
|
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
|
"noise": ("NOISE", {"tooltip": "The noise to use for the sampling."}),
|
|
"guiding_latents": (
|
|
"LATENT",
|
|
{
|
|
"tooltip": "The latents to use for guiding the sampling, typically with an IC-LoRA."
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"optional_cond_images": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": "The image to use for conditioning the sampling, if not provided, the sampling will be unconditioned (t2v setup). The image will be resized to the size of the first frame."
|
|
},
|
|
),
|
|
"num_frames": (
|
|
"INT",
|
|
{
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": 1000,
|
|
"step": 1,
|
|
"tooltip": "If -1, the number of frames will be based on the number of frames in the guiding_latents.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("denoised_video", "positive", "negative")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample(
|
|
self,
|
|
vae,
|
|
guider,
|
|
sampler,
|
|
sigmas,
|
|
noise,
|
|
guiding_latents,
|
|
optional_cond_images=None,
|
|
optional_cond_indices=None,
|
|
num_frames=0,
|
|
optional_initialization_latents=None,
|
|
optional_negative_index_latents=None,
|
|
optional_negative_index=-1,
|
|
optional_negative_index_strength=1.0,
|
|
cond_image_strength=1.0,
|
|
guiding_strength=1.0,
|
|
guiding_start_step=0,
|
|
guiding_end_step=1000,
|
|
):
|
|
guider = copy.copy(guider)
|
|
guider.original_conds = copy.deepcopy(guider.original_conds)
|
|
if optional_cond_images is None:
|
|
optional_cond_indices = None
|
|
|
|
if optional_cond_indices is not None and optional_cond_images is not None:
|
|
optional_cond_indices = optional_cond_indices.split(",")
|
|
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
|
assert len(optional_cond_indices) == len(
|
|
optional_cond_images
|
|
), "Number of optional cond images must match number of optional cond indices"
|
|
|
|
positive, negative = _get_raw_conds_from_guider(guider)
|
|
|
|
time_scale_factor, width_scale_factor, height_scale_factor = (
|
|
vae.downscale_index_formula
|
|
)
|
|
|
|
batch, channels, frames, height, width = guiding_latents["samples"].shape
|
|
if num_frames != -1:
|
|
frames = (num_frames - 1) // time_scale_factor + 1
|
|
|
|
if optional_initialization_latents is not None:
|
|
new_latents = optional_initialization_latents
|
|
else:
|
|
new_latents = EmptyLTXVLatentVideo.execute(
|
|
width=width * width_scale_factor,
|
|
height=height * height_scale_factor,
|
|
length=(frames - 1) * time_scale_factor + 1,
|
|
batch_size=1,
|
|
)[0]
|
|
|
|
high_sigmas, rest_sigmas = SplitSigmas().get_sigmas(sigmas, guiding_start_step)
|
|
middle_sigmas, low_sigmas = SplitSigmas().get_sigmas(
|
|
rest_sigmas, guiding_end_step - guiding_start_step
|
|
)
|
|
|
|
if len(high_sigmas) > 1:
|
|
print(
|
|
"Denoising with keyframes only [if available] on sigmas: ",
|
|
high_sigmas,
|
|
)
|
|
(_, new_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=high_sigmas,
|
|
latent_image=new_latents,
|
|
)
|
|
|
|
if optional_cond_indices is not None and 0 in optional_cond_indices:
|
|
guiding_latents = LTXVSelectLatents().select_latents(
|
|
guiding_latents, 1, -1
|
|
)[0]
|
|
skip_one_guiding_latent = True
|
|
else:
|
|
skip_one_guiding_latent = False
|
|
|
|
print("Adding conditioning on guiding latents")
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=new_latents,
|
|
guiding_latent=guiding_latents,
|
|
latent_idx=1 if skip_one_guiding_latent else 0,
|
|
strength=guiding_strength,
|
|
)
|
|
if optional_cond_images is not None:
|
|
print("Adding conditioning on keyframes", optional_cond_indices)
|
|
for cond_image, cond_idx in zip(
|
|
optional_cond_images, optional_cond_indices
|
|
):
|
|
if cond_idx % 8 == 1:
|
|
raise ValueError(
|
|
f"Conditioning image index {cond_idx} is a multiple of 8 + 1 and guiding latents are used. Please provide other cond image indices"
|
|
)
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = LTXVAddGuide.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
vae=vae,
|
|
latent=new_latents,
|
|
image=cond_image.unsqueeze(0),
|
|
frame_idx=cond_idx,
|
|
strength=cond_image_strength,
|
|
)
|
|
if optional_negative_index_latents is not None:
|
|
(
|
|
positive,
|
|
negative,
|
|
new_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=new_latents,
|
|
guiding_latent=optional_negative_index_latents,
|
|
latent_idx=optional_negative_index,
|
|
strength=optional_negative_index_strength,
|
|
)
|
|
|
|
guider.set_conds(positive, negative)
|
|
|
|
# Denoise the latent video
|
|
print("Denoising with full conditioning on sigmas: ", middle_sigmas)
|
|
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=middle_sigmas,
|
|
latent_image=new_latents,
|
|
)
|
|
|
|
# Clean up guides if image conditioning was used
|
|
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=denoised_output_latents,
|
|
)
|
|
|
|
if len(low_sigmas) > 1:
|
|
guider.set_conds(positive, negative)
|
|
print(
|
|
"Denoising with keyframes only [if available] conditioning on sigmas: ",
|
|
low_sigmas,
|
|
)
|
|
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=low_sigmas,
|
|
latent_image=denoised_output_latents,
|
|
)
|
|
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=denoised_output_latents,
|
|
)
|
|
|
|
return (denoised_output_latents, positive, negative)
|
|
|
|
|
|
@comfy_node(description="Linear transition with overlap")
|
|
class LinearOverlapLatentTransition:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"samples1": ("LATENT",),
|
|
"samples2": ("LATENT",),
|
|
"overlap": ("INT", {"default": 1, "min": 1, "max": 256}),
|
|
},
|
|
"optional": {
|
|
"axis": ("INT", {"default": 0}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "process"
|
|
|
|
CATEGORY = "Lightricks/latent"
|
|
|
|
def get_subbatch(self, samples):
|
|
s = samples.copy()
|
|
samples = s["samples"]
|
|
return samples
|
|
|
|
def process(self, samples1, samples2, overlap, axis=0):
|
|
samples1 = self.get_subbatch(samples1)
|
|
samples2 = self.get_subbatch(samples2)
|
|
|
|
# Create transition coefficients
|
|
alpha = torch.linspace(1, 0, overlap + 2)[1:-1].to(samples1.device)
|
|
|
|
# Create shape for broadcasting based on the axis
|
|
shape = [1] * samples1.dim()
|
|
shape[axis] = alpha.size(0)
|
|
alpha = alpha.reshape(shape)
|
|
|
|
# Create slices for the overlap regions
|
|
slice_all = [slice(None)] * samples1.dim()
|
|
slice_overlap1 = slice_all.copy()
|
|
slice_overlap1[axis] = slice(-overlap, None)
|
|
slice_overlap2 = slice_all.copy()
|
|
slice_overlap2[axis] = slice(0, overlap)
|
|
slice_rest1 = slice_all.copy()
|
|
slice_rest1[axis] = slice(None, -overlap)
|
|
slice_rest2 = slice_all.copy()
|
|
slice_rest2[axis] = slice(overlap, None)
|
|
|
|
# Combine samples
|
|
parts = [
|
|
samples1[tuple(slice_rest1)],
|
|
alpha * samples1[tuple(slice_overlap1)]
|
|
+ (1 - alpha) * samples2[tuple(slice_overlap2)],
|
|
samples2[tuple(slice_rest2)],
|
|
]
|
|
|
|
combined_samples = torch.cat(parts, dim=axis)
|
|
combined_batch_index = torch.arange(0, combined_samples.shape[0]).to(
|
|
dtype=torch.float32
|
|
)
|
|
|
|
return (
|
|
{
|
|
"samples": combined_samples,
|
|
"batch_index": combined_batch_index,
|
|
},
|
|
)
|
|
|
|
|
|
@comfy_node(description="LTXV Normalizing Sampler")
|
|
class LTXVNormalizingSampler(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LTXVNormalizingSampler",
|
|
category="utility",
|
|
inputs=[
|
|
io.Noise.Input("noise"),
|
|
io.Guider.Input("guider"),
|
|
io.Sampler.Input("sampler"),
|
|
io.Sigmas.Input("sigmas"),
|
|
io.Latent.Input("latent_image"),
|
|
io.String.Input(
|
|
"video_normalization_factors", default="1,1,1,1,1,1,1,1"
|
|
),
|
|
io.String.Input(
|
|
"audio_normalization_factors", default="1,1,0.25,1,1,0.25,1,1"
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output(display_name="denoised_output"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
noise,
|
|
guider,
|
|
sampler,
|
|
sigmas,
|
|
latent_image,
|
|
video_normalization_factors,
|
|
audio_normalization_factors,
|
|
) -> io.NodeOutput:
|
|
|
|
if (
|
|
guider.model_patcher.model.diffusion_model.__class__.__name__
|
|
!= "LTXAVModel"
|
|
):
|
|
raise ValueError()
|
|
|
|
ltxav: LTXAVModel = guider.model_patcher.model.diffusion_model
|
|
|
|
video_normalization_factors = video_normalization_factors.split(",")
|
|
audio_normalization_factors = audio_normalization_factors.split(",")
|
|
video_normalization_factors = [
|
|
float(factor) for factor in video_normalization_factors
|
|
]
|
|
audio_normalization_factors = [
|
|
float(factor) for factor in audio_normalization_factors
|
|
]
|
|
|
|
# Extend normalization factors to match the length of sigmas
|
|
sigmas_len = len(sigmas) - 1
|
|
if len(video_normalization_factors) < sigmas_len:
|
|
if len(video_normalization_factors) > 0:
|
|
video_normalization_factors.extend(
|
|
[video_normalization_factors[-1]]
|
|
* (sigmas_len - len(video_normalization_factors))
|
|
)
|
|
if len(audio_normalization_factors) < sigmas_len:
|
|
if len(audio_normalization_factors) > 0:
|
|
audio_normalization_factors.extend(
|
|
[audio_normalization_factors[-1]]
|
|
* (sigmas_len - len(audio_normalization_factors))
|
|
)
|
|
|
|
# Calculate indices where both normalization factors are not 1.0
|
|
sampling_split_indices = [
|
|
i + 1
|
|
for i, (v, a) in enumerate(
|
|
zip(video_normalization_factors, audio_normalization_factors)
|
|
)
|
|
if v != 1.0 or a != 1.0
|
|
]
|
|
print("Sampling split indices: %s" % sampling_split_indices, flush=True)
|
|
|
|
# Split sigmas according to sampling_split_indices
|
|
def split_by_indices(arr, indices):
|
|
"""
|
|
Splits arr into chunks according to indices (split points).
|
|
Indices are treated as starting a new chunk at each index in the list.
|
|
"""
|
|
if not indices:
|
|
return [arr]
|
|
split_points = sorted(set(indices))
|
|
chunks = []
|
|
prev = 0
|
|
for idx in split_points:
|
|
if prev < idx:
|
|
chunks.append(arr[prev : idx + 1])
|
|
prev = idx
|
|
if prev < len(arr):
|
|
chunks.append(arr[prev:])
|
|
return chunks
|
|
|
|
sigmas_chunks = split_by_indices(sigmas, sampling_split_indices)
|
|
print("Sigmas chunks: %s" % sigmas_chunks, flush=True)
|
|
|
|
i = 0
|
|
for sigmas_chunk in sigmas_chunks:
|
|
i += len(sigmas_chunk) - 1
|
|
print("Sampling with sigmas %s" % (sigmas_chunk), flush=True)
|
|
(_, latent_image) = SamplerCustomAdvanced().execute(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=sigmas_chunk,
|
|
latent_image=latent_image,
|
|
)
|
|
video_samples, audio_samples = ltxav.separate_audio_and_video_latents(
|
|
latent_image["samples"].tensors,
|
|
None,
|
|
)
|
|
if i - 1 < len(video_normalization_factors) and i - 1 < len(
|
|
audio_normalization_factors
|
|
):
|
|
video_samples = video_samples * video_normalization_factors[i - 1]
|
|
audio_samples = audio_samples * audio_normalization_factors[i - 1]
|
|
latent_image["samples"] = NestedTensor(
|
|
ltxav.recombine_audio_and_video_latents(
|
|
video_samples, audio_samples
|
|
)
|
|
)
|
|
print(
|
|
"After %d steps, the latent image was normalized by %f and %f"
|
|
% (
|
|
i,
|
|
video_normalization_factors[i - 1],
|
|
audio_normalization_factors[i - 1],
|
|
),
|
|
flush=True,
|
|
)
|
|
|
|
return io.NodeOutput(latent_image)
|