576 lines
28 KiB
Python
576 lines
28 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
NODE_CLASS_MAPPINGS: dict[str, type] = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {}
|
|
|
|
|
|
def register_node(name: str, display_name: str):
|
|
def deco(cls: type) -> type:
|
|
NODE_CLASS_MAPPINGS[name] = cls
|
|
NODE_DISPLAY_NAME_MAPPINGS[name] = display_name
|
|
return cls
|
|
return deco
|
|
|
|
|
|
@dataclass
|
|
class Ltx23MotionCacheConfig:
|
|
enabled: bool
|
|
warmup_steps: int
|
|
refresh_interval: int
|
|
motion_threshold: float
|
|
min_update_ratio: float
|
|
max_update_ratio: float
|
|
cache_cfg_delta: bool
|
|
cache_stg_delta: bool
|
|
log_stats: bool
|
|
|
|
|
|
@dataclass
|
|
class Ltx23SampledVideo:
|
|
video_latent: torch.Tensor
|
|
audio_latent: torch.Tensor | None
|
|
generation_config: Any
|
|
|
|
|
|
def _repo_root_candidates() -> list[Path]:
|
|
here = Path(__file__).resolve().parent
|
|
candidates: list[Path] = []
|
|
env_root = os.environ.get("LTX2_REPO_ROOT", "").strip()
|
|
if env_root:
|
|
candidates.append(Path(env_root).expanduser())
|
|
# Optional self-contained runtime folder if the user later vendors LTX-2 here.
|
|
candidates.extend([here / "vendor" / "LTX-2", here / "runtime" / "LTX-2"])
|
|
return candidates
|
|
|
|
|
|
def _ensure_ltx_imports() -> None:
|
|
if importlib.util.find_spec("ltx_core") and importlib.util.find_spec("ltx_trainer"):
|
|
return
|
|
for root in _repo_root_candidates():
|
|
core = root / "packages" / "ltx-core" / "src"
|
|
trainer = root / "packages" / "ltx-trainer" / "src"
|
|
pipelines = root / "packages" / "ltx-pipelines" / "src"
|
|
if core.exists() and trainer.exists():
|
|
for path in (core, trainer, pipelines):
|
|
if path.exists():
|
|
text = str(path)
|
|
if text not in sys.path:
|
|
sys.path.insert(0, text)
|
|
if importlib.util.find_spec("ltx_core") and importlib.util.find_spec("ltx_trainer"):
|
|
return
|
|
raise ImportError(
|
|
"Cannot import ltx_core/ltx_trainer. Install the LTX-2.3 runtime or set LTX2_REPO_ROOT "
|
|
"to a repo containing packages/ltx-core/src and packages/ltx-trainer/src."
|
|
)
|
|
|
|
|
|
def _make_validation_sampler(transformer_handle: Any, vae_decoder_handle: Any, vae_encoder_handle: Any | None) -> Any:
|
|
_ensure_ltx_imports()
|
|
from ltx_trainer.validation_sampler import ValidationSampler
|
|
|
|
transformer = getattr(transformer_handle, "transformer", transformer_handle)
|
|
vae_decoder = getattr(vae_decoder_handle, "vae_decoder", vae_decoder_handle)
|
|
vae_encoder = None if vae_encoder_handle is None else getattr(vae_encoder_handle, "vae_encoder", vae_encoder_handle)
|
|
return ValidationSampler(
|
|
transformer=transformer,
|
|
vae_decoder=vae_decoder,
|
|
vae_encoder=vae_encoder,
|
|
text_encoder=None,
|
|
audio_decoder=None,
|
|
vocoder=None,
|
|
embeddings_processor=None,
|
|
)
|
|
|
|
|
|
def _device_from_transformer(transformer_handle: Any) -> torch.device:
|
|
return torch.device(getattr(transformer_handle, "device", "cuda"))
|
|
|
|
|
|
def _checkpoint_path_from_decoder(vae_decoder_handle: Any) -> str:
|
|
return str(getattr(vae_decoder_handle, "checkpoint_path", ""))
|
|
|
|
|
|
def _select_motion_cache_update_mask(
|
|
current: torch.Tensor,
|
|
cached: torch.Tensor | None,
|
|
denoise_mask: torch.Tensor,
|
|
motion_cache: Ltx23MotionCacheConfig,
|
|
step_idx: int,
|
|
) -> tuple[torch.Tensor, dict[str, float]]:
|
|
if cached is None or not motion_cache.enabled or step_idx < motion_cache.warmup_steps:
|
|
mask = torch.ones_like(denoise_mask, dtype=torch.bool)
|
|
return mask, {"update_ratio": 1.0, "score_mean": 0.0, "score_max": 0.0}
|
|
if motion_cache.refresh_interval > 1 and step_idx % motion_cache.refresh_interval == 0:
|
|
mask = torch.ones_like(denoise_mask, dtype=torch.bool)
|
|
return mask, {"update_ratio": 1.0, "score_mean": 0.0, "score_max": 0.0}
|
|
|
|
score = (current.float() - cached.float()).abs().mean(dim=-1, keepdim=True)
|
|
active = denoise_mask > 0
|
|
if not bool(active.any()):
|
|
mask = torch.zeros_like(denoise_mask, dtype=torch.bool)
|
|
return mask, {"update_ratio": 0.0, "score_mean": 0.0, "score_max": 0.0}
|
|
|
|
active_scores = score[active]
|
|
score_max = active_scores.max().clamp_min(1e-6)
|
|
score_norm = score / score_max
|
|
mask = (score_norm >= motion_cache.motion_threshold) & active
|
|
|
|
flat_score = score_norm[active].flatten()
|
|
flat_mask = mask[active].flatten()
|
|
total = flat_score.numel()
|
|
if total > 0:
|
|
min_k = max(1, int(round(total * motion_cache.min_update_ratio)))
|
|
max_k = max(min_k, int(round(total * motion_cache.max_update_ratio)))
|
|
current_k = int(flat_mask.sum().item())
|
|
if current_k < min_k:
|
|
idx = torch.topk(flat_score, k=min_k, largest=True).indices
|
|
flat_mask = torch.zeros_like(flat_mask, dtype=torch.bool)
|
|
flat_mask[idx] = True
|
|
elif current_k > max_k:
|
|
idx = torch.topk(flat_score, k=max_k, largest=True).indices
|
|
flat_mask = torch.zeros_like(flat_mask, dtype=torch.bool)
|
|
flat_mask[idx] = True
|
|
rebuilt = torch.zeros_like(mask, dtype=torch.bool)
|
|
rebuilt[active] = flat_mask
|
|
mask = rebuilt
|
|
|
|
update_ratio = float(mask[active].float().mean().item()) if bool(active.any()) else 0.0
|
|
return mask, {
|
|
"update_ratio": update_ratio,
|
|
"score_mean": float(active_scores.float().mean().item()),
|
|
"score_max": float(active_scores.float().max().item()),
|
|
}
|
|
|
|
|
|
def _run_denoising_motioncache(
|
|
sampler: Any,
|
|
motion_cache: Ltx23MotionCacheConfig,
|
|
config: Any,
|
|
video_state: Any,
|
|
audio_state: Any | None,
|
|
video_clean_state: Any,
|
|
audio_clean_state: Any | None,
|
|
v_ctx_pos: torch.Tensor,
|
|
a_ctx_pos: torch.Tensor,
|
|
v_ctx_neg: torch.Tensor | None,
|
|
a_ctx_neg: torch.Tensor | None,
|
|
device: torch.device,
|
|
) -> tuple[Any, Any | None]:
|
|
_ensure_ltx_imports()
|
|
from dataclasses import replace
|
|
|
|
from ltx_core.components.diffusion_steps import EulerDiffusionStep
|
|
from ltx_core.components.guiders import CFGGuider, STGGuider
|
|
from ltx_core.components.schedulers import LTX2Scheduler
|
|
from ltx_core.model.transformer.modality import Modality
|
|
from ltx_core.model.transformer.model import X0Model
|
|
|
|
scheduler = LTX2Scheduler()
|
|
sigmas = scheduler.execute(steps=config.num_inference_steps).to(device).float()
|
|
stepper = EulerDiffusionStep()
|
|
cfg_guider = CFGGuider(config.guidance_scale)
|
|
stg_guider = STGGuider(config.stg_scale)
|
|
stg_perturbation_config = sampler._build_stg_perturbation_config(config) if stg_guider.enabled() else None
|
|
|
|
video = Modality(
|
|
enabled=True,
|
|
latent=video_state.latent,
|
|
sigma=sigmas[0].repeat(video_state.latent.shape[0]),
|
|
timesteps=video_state.denoise_mask,
|
|
positions=video_state.positions,
|
|
context=v_ctx_pos,
|
|
context_mask=None,
|
|
)
|
|
audio = None
|
|
if audio_state is not None:
|
|
audio = Modality(
|
|
enabled=True,
|
|
latent=audio_state.latent,
|
|
sigma=sigmas[0].repeat(audio_state.latent.shape[0]),
|
|
timesteps=audio_state.denoise_mask,
|
|
positions=audio_state.positions,
|
|
context=a_ctx_pos,
|
|
context_mask=None,
|
|
)
|
|
|
|
sampler._transformer.to(device)
|
|
x0_model = X0Model(sampler._transformer)
|
|
cached_video_x0: torch.Tensor | None = None
|
|
cached_cfg_delta_video: torch.Tensor | None = None
|
|
cached_stg_delta_video: torch.Tensor | None = None
|
|
stats: list[dict[str, float]] = []
|
|
|
|
with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
|
|
for step_idx, sigma in enumerate(sigmas[:-1]):
|
|
video = replace(
|
|
video,
|
|
latent=video_state.latent,
|
|
sigma=sigma.repeat(video_state.latent.shape[0]),
|
|
timesteps=sigma * video_state.denoise_mask,
|
|
positions=video_state.positions,
|
|
)
|
|
if audio is not None and audio_state is not None:
|
|
audio = replace(
|
|
audio,
|
|
latent=audio_state.latent,
|
|
sigma=sigma.repeat(audio_state.latent.shape[0]),
|
|
timesteps=sigma * audio_state.denoise_mask,
|
|
positions=audio_state.positions,
|
|
)
|
|
|
|
pos_video, pos_audio = x0_model(video=video, audio=audio, perturbations=None)
|
|
denoised_video, denoised_audio = pos_video, pos_audio
|
|
do_full_guidance = (
|
|
step_idx < motion_cache.warmup_steps
|
|
or motion_cache.refresh_interval <= 1
|
|
or step_idx % motion_cache.refresh_interval == 0
|
|
)
|
|
|
|
if cfg_guider.enabled() and v_ctx_neg is not None:
|
|
if do_full_guidance or not motion_cache.cache_cfg_delta or cached_cfg_delta_video is None:
|
|
video_neg = replace(video, context=v_ctx_neg)
|
|
audio_neg = replace(audio, context=a_ctx_neg) if audio is not None else None
|
|
neg_video, neg_audio = x0_model(video=video_neg, audio=audio_neg, perturbations=None)
|
|
cfg_delta_video = cfg_guider.delta(pos_video, neg_video)
|
|
cached_cfg_delta_video = cfg_delta_video.detach()
|
|
denoised_video = denoised_video + cfg_delta_video
|
|
if audio is not None and denoised_audio is not None:
|
|
denoised_audio = denoised_audio + cfg_guider.delta(pos_audio, neg_audio)
|
|
else:
|
|
denoised_video = denoised_video + cached_cfg_delta_video.to(denoised_video.device, denoised_video.dtype)
|
|
|
|
if stg_guider.enabled() and stg_perturbation_config is not None:
|
|
if do_full_guidance or not motion_cache.cache_stg_delta or cached_stg_delta_video is None:
|
|
perturbed_video, perturbed_audio = x0_model(video=video, audio=audio, perturbations=stg_perturbation_config)
|
|
stg_delta_video = stg_guider.delta(pos_video, perturbed_video)
|
|
cached_stg_delta_video = stg_delta_video.detach()
|
|
denoised_video = denoised_video + stg_delta_video
|
|
if audio is not None and denoised_audio is not None and perturbed_audio is not None:
|
|
denoised_audio = denoised_audio + stg_guider.delta(pos_audio, perturbed_audio)
|
|
else:
|
|
denoised_video = denoised_video + cached_stg_delta_video.to(denoised_video.device, denoised_video.dtype)
|
|
|
|
denoised_video = denoised_video * video_state.denoise_mask + video_clean_state.latent.float() * (1 - video_state.denoise_mask)
|
|
if motion_cache.enabled:
|
|
update_mask, step_stats = _select_motion_cache_update_mask(
|
|
current=denoised_video,
|
|
cached=cached_video_x0,
|
|
denoise_mask=video_state.denoise_mask,
|
|
motion_cache=motion_cache,
|
|
step_idx=step_idx,
|
|
)
|
|
if cached_video_x0 is not None:
|
|
denoised_video = torch.where(update_mask, denoised_video, cached_video_x0.to(denoised_video.device, denoised_video.dtype))
|
|
cached_video_x0 = denoised_video.detach()
|
|
step_stats["step"] = float(step_idx)
|
|
stats.append(step_stats)
|
|
|
|
if audio is not None and audio_state is not None and audio_clean_state is not None:
|
|
# Keep audio path correct, but do not pretend to have spatial MotionCache for audio.
|
|
denoised_audio = denoised_audio * audio_state.denoise_mask + audio_clean_state.latent.float() * (1 - audio_state.denoise_mask)
|
|
|
|
video_state = replace(
|
|
video_state,
|
|
latent=stepper.step(sample=video.latent, denoised_sample=denoised_video, sigmas=sigmas, step_index=step_idx),
|
|
)
|
|
if audio is not None and audio_state is not None:
|
|
audio_state = replace(
|
|
audio_state,
|
|
latent=stepper.step(sample=audio.latent, denoised_sample=denoised_audio, sigmas=sigmas, step_index=step_idx),
|
|
)
|
|
if sampler._sampling_context is not None:
|
|
sampler._sampling_context.advance_step()
|
|
|
|
sampler._motioncache_last_stats = stats
|
|
return video_state, audio_state
|
|
|
|
|
|
def _summarize_motioncache(sampler: Any, motion_cache: Ltx23MotionCacheConfig) -> str:
|
|
stats = getattr(sampler, "_motioncache_last_stats", [])
|
|
if not motion_cache.log_stats:
|
|
return ""
|
|
if not stats:
|
|
return "MotionCache disabled or no denoising stats collected."
|
|
avg_update = sum(item.get("update_ratio", 0.0) for item in stats) / len(stats)
|
|
min_update = min(item.get("update_ratio", 0.0) for item in stats)
|
|
max_update = max(item.get("update_ratio", 0.0) for item in stats)
|
|
return (
|
|
f"MotionCache steps={len(stats)}, avg_update_ratio={avg_update:.3f}, "
|
|
f"min_update_ratio={min_update:.3f}, max_update_ratio={max_update:.3f}, "
|
|
f"refresh_interval={motion_cache.refresh_interval}, threshold={motion_cache.motion_threshold:.3f}"
|
|
)
|
|
|
|
|
|
def _with_motioncache_run(sampler: Any, motion_cache: Ltx23MotionCacheConfig, fn):
|
|
old_run_denoising = sampler._run_denoising
|
|
|
|
def patched_run_denoising(*args: Any, **kwargs: Any) -> tuple[Any, Any | None]:
|
|
return _run_denoising_motioncache(sampler, motion_cache, *args, **kwargs)
|
|
|
|
sampler._run_denoising = patched_run_denoising
|
|
try:
|
|
return fn()
|
|
finally:
|
|
sampler._run_denoising = old_run_denoising
|
|
|
|
|
|
def _sample_latent_like_ltx_node(sampler: Any, config: Any, device: torch.device) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Run LTX generation and return patchified target latent, mirroring modular LTX sampler behavior."""
|
|
_ensure_ltx_imports()
|
|
from ltx_core.components.noisers import GaussianNoiser
|
|
from ltx_core.types import LatentState
|
|
|
|
sampler._validate_config(config)
|
|
v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg = sampler._get_prompt_embeddings(config, device)
|
|
generator = torch.Generator(device=device).manual_seed(config.seed)
|
|
noiser = GaussianNoiser(generator=generator)
|
|
|
|
if config.reference_video is not None and config.mask_video is not None:
|
|
ref_video = sampler._preprocess_conditioning_video(config.reference_video, config, "reference")
|
|
mask_video = sampler._preprocess_conditioning_video(config.mask_video, config, "mask")
|
|
ref_latent, ref_positions = sampler._encode_video(ref_video, config.frame_rate, device)
|
|
mask_latent, mask_positions = sampler._encode_video(mask_video, config.frame_rate, device)
|
|
if ref_latent.shape[1] != mask_latent.shape[1]:
|
|
raise ValueError(f"Reference and mask videos must produce same sequence length: {ref_latent.shape[1]} vs {mask_latent.shape[1]}")
|
|
if getattr(config, "reference_downscale_factor", 1) != 1:
|
|
factor = config.reference_downscale_factor
|
|
ref_positions = ref_positions.clone(); ref_positions[:, 1:, ...] *= factor
|
|
mask_positions = mask_positions.clone(); mask_positions[:, 1:, ...] *= factor
|
|
video_tools = sampler._create_video_latent_tools(config)
|
|
audio_tools = sampler._create_audio_latent_tools(config) if config.generate_audio else None
|
|
target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
|
|
audio_clean_state = audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
|
|
if config.condition_image is not None:
|
|
target_clean_state = sampler._apply_image_conditioning(target_clean_state, config.condition_image, config, device)
|
|
ref_seq_len = ref_latent.shape[1]
|
|
mask_seq_len = mask_latent.shape[1]
|
|
prefix_mask = torch.zeros(1, ref_seq_len + mask_seq_len, 1, device=device, dtype=torch.float32)
|
|
combined_clean_state = LatentState(
|
|
latent=torch.cat([ref_latent, mask_latent, target_clean_state.latent], dim=1),
|
|
denoise_mask=torch.cat([prefix_mask, target_clean_state.denoise_mask], dim=1),
|
|
positions=torch.cat([ref_positions, mask_positions, target_clean_state.positions], dim=2),
|
|
clean_latent=torch.cat([ref_latent, mask_latent, target_clean_state.clean_latent], dim=1),
|
|
)
|
|
video_state = noiser(latent_state=combined_clean_state, noise_scale=1.0)
|
|
audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state is not None else None
|
|
video_state, audio_state = sampler._run_denoising(config, video_state, audio_state, combined_clean_state, audio_clean_state, v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg, device)
|
|
video_latent = video_state.latent[:, ref_seq_len + mask_seq_len:].detach().cpu()
|
|
elif config.reference_video is not None:
|
|
ref_video = sampler._preprocess_reference_video(config)
|
|
ref_latent, ref_positions = sampler._encode_video(ref_video, config.frame_rate, device)
|
|
if getattr(config, "reference_downscale_factor", 1) != 1:
|
|
ref_positions = ref_positions.clone()
|
|
ref_positions[:, 1:, ...] *= config.reference_downscale_factor
|
|
video_tools = sampler._create_video_latent_tools(config)
|
|
audio_tools = sampler._create_audio_latent_tools(config) if config.generate_audio else None
|
|
target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
|
|
audio_clean_state = audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
|
|
if config.condition_image is not None:
|
|
target_clean_state = sampler._apply_image_conditioning(target_clean_state, config.condition_image, config, device)
|
|
ref_seq_len = ref_latent.shape[1]
|
|
ref_mask = torch.zeros(1, ref_seq_len, 1, device=device, dtype=torch.float32)
|
|
combined_clean_state = LatentState(
|
|
latent=torch.cat([ref_latent, target_clean_state.latent], dim=1),
|
|
denoise_mask=torch.cat([ref_mask, target_clean_state.denoise_mask], dim=1),
|
|
positions=torch.cat([ref_positions, target_clean_state.positions], dim=2),
|
|
clean_latent=torch.cat([ref_latent, target_clean_state.clean_latent], dim=1),
|
|
)
|
|
video_state = noiser(latent_state=combined_clean_state, noise_scale=1.0)
|
|
audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state is not None else None
|
|
video_state, audio_state = sampler._run_denoising(config, video_state, audio_state, combined_clean_state, audio_clean_state, v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg, device)
|
|
video_latent = video_state.latent[:, ref_seq_len:].detach().cpu()
|
|
else:
|
|
video_tools = sampler._create_video_latent_tools(config)
|
|
audio_tools = sampler._create_audio_latent_tools(config) if config.generate_audio else None
|
|
video_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
|
|
audio_clean_state = audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
|
|
if config.condition_image is not None:
|
|
video_clean_state = sampler._apply_image_conditioning(video_clean_state, config.condition_image, config, device)
|
|
video_state = noiser(latent_state=video_clean_state, noise_scale=1.0)
|
|
audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state is not None else None
|
|
video_state, audio_state = sampler._run_denoising(config, video_state, audio_state, video_clean_state, audio_clean_state, v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg, device)
|
|
video_latent = video_tools.clear_conditioning(video_state).latent.detach().cpu()
|
|
|
|
audio_latent = None
|
|
if audio_state is not None and audio_tools is not None:
|
|
audio_state = audio_tools.clear_conditioning(audio_state)
|
|
audio_state = audio_tools.unpatchify(audio_state)
|
|
audio_latent = audio_state.latent.detach().cpu()
|
|
return video_latent, audio_latent
|
|
|
|
|
|
def _video_tensor_to_comfy(video: torch.Tensor) -> torch.Tensor:
|
|
# LTX ValidationSampler returns [C,F,H,W]. ComfyUI IMAGE is [F,H,W,C].
|
|
if video.dim() != 4:
|
|
raise ValueError(f"Expected video tensor [C,F,H,W], got shape {tuple(video.shape)}")
|
|
return video.detach().cpu().to(torch.float32).permute(1, 2, 3, 0).contiguous().clamp(0, 1)
|
|
|
|
|
|
@register_node("Ltx23MotionCache", "🅛🅣🅧 MotionCache Config")
|
|
class Ltx23MotionCache:
|
|
CATEGORY = "video/LTX-2.3/MotionCache"
|
|
RETURN_TYPES = ("LTX23_MOTION_CACHE_CONFIG",)
|
|
RETURN_NAMES = ("motion_cache",)
|
|
FUNCTION = "build"
|
|
DESCRIPTION = "Build a standalone MotionCache-style config for LTX-2.3 sampling."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"warmup_steps": ("INT", {"default": 2, "min": 0, "max": 50, "step": 1}),
|
|
"refresh_interval": ("INT", {"default": 3, "min": 1, "max": 20, "step": 1}),
|
|
"motion_threshold": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"min_update_ratio": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"max_update_ratio": ("FLOAT", {"default": 0.85, "min": 0.01, "max": 1.0, "step": 0.01}),
|
|
"cache_cfg_delta": ("BOOLEAN", {"default": True}),
|
|
"cache_stg_delta": ("BOOLEAN", {"default": True}),
|
|
"log_stats": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
def build(self, enabled: bool, warmup_steps: int, refresh_interval: int, motion_threshold: float, min_update_ratio: float, max_update_ratio: float, cache_cfg_delta: bool, cache_stg_delta: bool, log_stats: bool) -> tuple[Ltx23MotionCacheConfig]:
|
|
min_update_ratio = max(0.0, min(1.0, float(min_update_ratio)))
|
|
max_update_ratio = max(min_update_ratio, min(1.0, float(max_update_ratio)))
|
|
return (Ltx23MotionCacheConfig(bool(enabled), max(0, int(warmup_steps)), max(1, int(refresh_interval)), max(0.0, min(1.0, float(motion_threshold))), min_update_ratio, max_update_ratio, bool(cache_cfg_delta), bool(cache_stg_delta), bool(log_stats)),)
|
|
|
|
|
|
@register_node("Ltx23SampleVideoMotionCache", "🅛🅣🅧 Sample Video MotionCache")
|
|
class Ltx23SampleVideoMotionCache:
|
|
CATEGORY = "video/LTX-2.3/MotionCache"
|
|
RETURN_TYPES = ("LTX_SAMPLED_VIDEO", "STRING")
|
|
RETURN_NAMES = ("sampled_video", "motioncache_stats")
|
|
FUNCTION = "sample"
|
|
DESCRIPTION = "Standalone MotionCache-style replacement for an LTX-2.3 modular sampler."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"transformer": ("LTX_TRANSFORMER",),
|
|
"vae_decoder": ("LTX_VIDEO_VAE_DECODER",),
|
|
"config": ("LTX_SAMPLING_CONFIG",),
|
|
"motion_cache": ("LTX23_MOTION_CACHE_CONFIG",),
|
|
},
|
|
"optional": {"vae_encoder": ("LTX_VIDEO_VAE_ENCODER",)},
|
|
}
|
|
|
|
def sample(self, transformer: Any, vae_decoder: Any, config: Any, motion_cache: Ltx23MotionCacheConfig, vae_encoder: Any | None = None) -> tuple[Ltx23SampledVideo, str]:
|
|
generation_config = getattr(config, "generation_config", config)
|
|
if (generation_config.condition_image is not None or generation_config.reference_video is not None or generation_config.mask_video is not None) and vae_encoder is None:
|
|
raise ValueError("vae_encoder is required for I2V / IC / Mask workflows")
|
|
sampler = _make_validation_sampler(transformer, vae_decoder, vae_encoder)
|
|
device = _device_from_transformer(transformer)
|
|
video_latent, audio_latent = _with_motioncache_run(
|
|
sampler,
|
|
motion_cache,
|
|
lambda: _sample_latent_like_ltx_node(sampler, generation_config, device),
|
|
)
|
|
return Ltx23SampledVideo(video_latent=video_latent, audio_latent=audio_latent, generation_config=generation_config), _summarize_motioncache(sampler, motion_cache)
|
|
|
|
|
|
@register_node("Ltx23GenerateVideoMotionCache", "🅛🅣🅧 Generate Video MotionCache")
|
|
class Ltx23GenerateVideoMotionCache:
|
|
CATEGORY = "video/LTX-2.3/MotionCache"
|
|
RETURN_TYPES = ("IMAGE", "AUDIO", "STRING")
|
|
RETURN_NAMES = ("frames", "audio", "motioncache_stats")
|
|
FUNCTION = "generate"
|
|
DESCRIPTION = "Generate ComfyUI frames directly with MotionCache-style LTX-2.3 denoising."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"transformer": ("LTX_TRANSFORMER",),
|
|
"vae_decoder": ("LTX_VIDEO_VAE_DECODER",),
|
|
"config": ("LTX_SAMPLING_CONFIG",),
|
|
"motion_cache": ("LTX23_MOTION_CACHE_CONFIG",),
|
|
},
|
|
"optional": {"vae_encoder": ("LTX_VIDEO_VAE_ENCODER",)},
|
|
}
|
|
|
|
def generate(self, transformer: Any, vae_decoder: Any, config: Any, motion_cache: Ltx23MotionCacheConfig, vae_encoder: Any | None = None) -> tuple[torch.Tensor, dict, str]:
|
|
generation_config = getattr(config, "generation_config", config)
|
|
if (generation_config.condition_image is not None or generation_config.reference_video is not None or generation_config.mask_video is not None) and vae_encoder is None:
|
|
raise ValueError("vae_encoder is required for I2V / IC / Mask workflows")
|
|
sampler = _make_validation_sampler(transformer, vae_decoder, vae_encoder)
|
|
device = _device_from_transformer(transformer)
|
|
video, audio = _with_motioncache_run(sampler, motion_cache, lambda: sampler.generate(generation_config, device))
|
|
audio_dict = {}
|
|
if audio is not None:
|
|
sample_rate = int(getattr(getattr(sampler, "_vocoder", None), "output_sampling_rate", 0) or 0)
|
|
audio_dict = {"waveform": audio.detach().cpu(), "sample_rate": sample_rate}
|
|
return _video_tensor_to_comfy(video), audio_dict, _summarize_motioncache(sampler, motion_cache)
|
|
|
|
|
|
@register_node("Ltx23MotionCacheModelPatch", "🅛🅣🅧 MotionCache Model Patch")
|
|
class Ltx23MotionCacheModelPatch:
|
|
CATEGORY = "video/LTX-2.3/MotionCache"
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "patch"
|
|
DESCRIPTION = "Patch a ComfyUI MODEL with native EasyCache-style MotionCache acceleration; designed to drop into existing LTX-2.3 SamplerCustomAdvanced workflows."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"reuse_threshold": ("FLOAT", {"default": 0.20, "min": 0.0, "max": 3.0, "step": 0.01}),
|
|
"start_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"end_percent": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"subsample_factor": ("INT", {"default": 8, "min": 1, "max": 32, "step": 1}),
|
|
"verbose": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
def patch(
|
|
self,
|
|
model: Any,
|
|
reuse_threshold: float,
|
|
start_percent: float,
|
|
end_percent: float,
|
|
subsample_factor: int,
|
|
verbose: bool,
|
|
) -> tuple[Any]:
|
|
# Reuse ComfyUI's native EasyCache wrappers so this standalone project can
|
|
# accelerate existing MODEL→GUIDER→SamplerCustomAdvanced LTX workflows.
|
|
import comfy.patcher_extension
|
|
from comfy_extras.nodes_easycache import (
|
|
EasyCacheHolder,
|
|
easycache_calc_cond_batch_wrapper,
|
|
easycache_forward_wrapper,
|
|
easycache_sample_wrapper,
|
|
)
|
|
|
|
patched = model.clone()
|
|
patched.model_options.setdefault("transformer_options", {})
|
|
output_channels = None
|
|
try:
|
|
output_channels = patched.model.latent_format.latent_channels
|
|
except Exception:
|
|
output_channels = None
|
|
patched.model_options["transformer_options"]["easycache"] = EasyCacheHolder(
|
|
float(reuse_threshold),
|
|
float(start_percent),
|
|
float(end_percent),
|
|
subsample_factor=max(1, int(subsample_factor)),
|
|
offload_cache_diff=False,
|
|
verbose=bool(verbose),
|
|
output_channels=output_channels,
|
|
)
|
|
patched.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, "ltx23_motioncache", easycache_sample_wrapper)
|
|
patched.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.CALC_COND_BATCH, "ltx23_motioncache", easycache_calc_cond_batch_wrapper)
|
|
patched.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, "ltx23_motioncache", easycache_forward_wrapper)
|
|
return (patched,)
|