From 9b75f73d1d3a7e095012568122a086102ab6bfa8 Mon Sep 17 00:00:00 2001 From: DiffSynth-Studio Bot Date: Mon, 17 Aug 2026 14:17:57 +0800 Subject: [PATCH 01/31] feat: add LTX-2.5 inference pipeline --- .gitignore | 4 + diffsynth/configs/model_configs.py | 77 ++++ .../configs/vram_management_module_maps.py | 18 + diffsynth/models/ltx25_diffvae/NOTICE.md | 9 + diffsynth/models/ltx25_diffvae/__init__.py | 0 .../models/ltx25_diffvae/model/__init__.py | 0 .../model/transformer/__init__.py | 0 .../model/transformer/timestep_embedding.py | 115 +++++ .../ltx25_diffvae/model/video_vae/__init__.py | 0 .../video_vae/diffusion_video_decoder.py | 208 +++++++++ .../ltx25_diffvae/model/video_vae/ops.py | 57 +++ .../model/video_vae/transformer/__init__.py | 3 + .../model/video_vae/transformer/attention.py | 59 +++ .../model/video_vae/transformer/blocks.py | 66 +++ .../transformer/combined/__init__.py | 0 .../video_vae/transformer/combined/attn.py | 114 +++++ .../video_vae/transformer/combined/block.py | 28 ++ .../video_vae/transformer/combined/context.py | 15 + .../video_vae/transformer/combined/mlp.py | 14 + .../video_vae/transformer/det_attn_rope.py | 156 +++++++ .../transformer/fallback_na/__init__.py | 8 + .../transformer/fallback_na/eager.py | 154 +++++++ .../model/video_vae/transformer/layers.py | 66 +++ .../model/video_vae/transformer/qkv.py | 16 + .../model/video_vae/transformer/rope_math.py | 56 +++ .../model/video_vae/transformer/swiglu.py | 38 ++ diffsynth/models/ltx25_duration_head.py | 57 +++ diffsynth/models/ltx25_text_encoder.py | 423 ++++++++++++++++++ diffsynth/models/ltx25_tokenizer.py | 51 +++ diffsynth/models/ltx2_dit.py | 39 +- diffsynth/models/ltx2_text_encoder.py | 4 + diffsynth/pipelines/ltx25_audio_video.py | 120 +++++ .../ltx25_diffusion_video_vae.py | 27 ++ .../ltx25_duration_head.py | 6 + .../ltx25_text_encoder.py | 31 ++ .../state_dict_converters/ltx2_video_vae.py | 8 +- docs/en/Model_Details/LTX-2.5.md | 88 ++++ docs/en/index.rst | 1 + docs/zh/Model_Details/LTX-2.5.md | 82 ++++ docs/zh/index.rst | 1 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 69 +++ .../LTX-2.5-Keyframe-Interpolation.py | 70 +++ .../LTX-2.5-T2AV-DistilledPipeline.py | 54 +++ 43 files changed, 2400 insertions(+), 12 deletions(-) create mode 100644 diffsynth/models/ltx25_diffvae/NOTICE.md create mode 100644 diffsynth/models/ltx25_diffvae/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/transformer/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/ops.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py create mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py create mode 100644 diffsynth/models/ltx25_duration_head.py create mode 100644 diffsynth/models/ltx25_text_encoder.py create mode 100644 diffsynth/models/ltx25_tokenizer.py create mode 100644 diffsynth/pipelines/ltx25_audio_video.py create mode 100644 diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py create mode 100644 diffsynth/utils/state_dict_converters/ltx25_duration_head.py create mode 100644 diffsynth/utils/state_dict_converters/ltx25_text_encoder.py create mode 100644 docs/en/Model_Details/LTX-2.5.md create mode 100644 docs/zh/Model_Details/LTX-2.5.md create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py diff --git a/.gitignore b/.gitignore index 7cce2c12d..dfe725a05 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,10 @@ /models /scripts /diffusers +# Local inference media and isolated upstream/environment artifacts. +# These can contain machine-specific paths, authenticated inspection logs, or large binaries. +/outputs +/packages /.vscode /.opencode *.pkl diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py index aa995365d..d37a21299 100644 --- a/diffsynth/configs/model_configs.py +++ b/diffsynth/configs/model_configs.py @@ -855,6 +855,83 @@ "extra_kwargs": {"separated_audio_video": True, "embedding_dim_gemma": 3840, "num_layers_gemma": 49, "video_attention_heads": 32, "video_attention_head_dim": 128, "audio_attention_heads": 32, "audio_attention_head_dim": 64, "num_connector_layers": 8, "apply_gated_attention": True}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter", }, + { + "model_hash": "7960c5dc4626650824e36f65a8e992e9", + "model_name": "ltx25_dit", + "model_class": "diffsynth.models.ltx2_dit.LTXModel", + "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", + }, + { + "model_hash": "4bc194ac62f5648db68d419916a25688", + "model_name": "ltx25_text_encoder", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + }, + { + "model_hash": "f1c63402b49c39c739f13cdb90714f9e", + "model_name": "ltx25_text_encoder_post_modules", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderPostModulesStateDictConverter", + }, + { + "model_hash": "e19205490f01801d0a7b6d3aba61e26e", + "model_name": "ltx25_video_vae_encoder", + "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", + "extra_kwargs": {"encoder_version": "ltx-2.3"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", + }, + { + "model_hash": "e19205490f01801d0a7b6d3aba61e26e", + "model_name": "ltx25_diffusion_video_vae_decoder", + "model_class": "diffsynth.models.ltx25_diffvae.model.video_vae.diffusion_video_decoder.DiffusionVideoDecoder", + "extra_kwargs": {"stage_channels": [2048, 1024, 512, 512, 256], "stage_depths": [4, 6, 4, 2, 8], "stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]], "stage5_kernel": [11, 11, 11], "timestep_scale_multiplier": 1000.0, "default_num_inference_steps": 1, "model_output_type": "x0"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoderStateDictConverter", + }, + { + "model_hash": "a1d642eecae96baa9c31d4e405564f49", + "model_name": "ltx25_conv_video_vae_encoder", + "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", + "extra_kwargs": {"encoder_version": "ltx-2.3"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", + }, + { + "model_hash": "a1d642eecae96baa9c31d4e405564f49", + "model_name": "ltx25_conv_video_vae_decoder", + "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder", + "extra_kwargs": {"decoder_version": "ltx-2.3"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter", + }, + { + "model_hash": "c2488315f13356abb806f9f217f1e803", + "model_name": "ltx25_audio_vae_decoder", + "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter", + }, + { + "model_hash": "c2488315f13356abb806f9f217f1e803", + "model_name": "ltx25_audio_vocoder", + "model_class": "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter", + }, + { + "model_hash": "c2488315f13356abb806f9f217f1e803", + "model_name": "ltx25_audio_vae_encoder", + "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter", + }, + { + "model_hash": "35840495e440a4f00946450269299bd6", + "model_name": "ltx25_duration_head", + "model_class": "diffsynth.models.ltx25_duration_head.LTX25DurationHead", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_duration_head.LTX25DurationHeadStateDictConverter", + }, + { + "model_hash": "5fbb28ecc6becd9513cee69b2dfb1053", + "model_name": "ltx25_temporal_upsampler", + "model_class": "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler", + "extra_kwargs": {"mid_channels": 512, "spatial_upsample": False, "temporal_upsample": True, "spatial_scale": 1.0, "rational_resampler": True}, + }, ] anima_series = [ { diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index a98ca9ba6..d51389fd3 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -266,6 +266,24 @@ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, + "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextDecoderLayer": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.ltx25_diffvae.model.video_vae.diffusion_video_decoder.DiffusionVideoDecoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffvae.model.video_vae.transformer.swiglu.SwiGLU": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.block.CombinedDiffusionNABlock": "diffsynth.core.vram.layers.AutoWrappedModule", + }, "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": { "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", diff --git a/diffsynth/models/ltx25_diffvae/NOTICE.md b/diffsynth/models/ltx25_diffvae/NOTICE.md new file mode 100644 index 000000000..34a20f287 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/NOTICE.md @@ -0,0 +1,9 @@ +# LTX-2.5 DiffVAE eager port notice + +The Python sources in this directory are a dependency-closed, in-tree port of +selected `ltx_core` DiffVAE decoder sources from `Lightricks/LTX-2` revision +`400fd31054597515f47125691032c04b1c3ee24e`. Original source headers are +retained. The port deliberately excludes NATTEN, Triton, Blackwell DSL, and +runtime imports from the installed `ltx_core` package. It uses the upstream +pure-PyTorch eager tiled-SDPA implementation as its neighborhood-attention +backend. diff --git a/diffsynth/models/ltx25_diffvae/__init__.py b/diffsynth/models/ltx25_diffvae/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffsynth/models/ltx25_diffvae/model/__init__.py b/diffsynth/models/ltx25_diffvae/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffsynth/models/ltx25_diffvae/model/transformer/__init__.py b/diffsynth/models/ltx25_diffvae/model/transformer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py b/diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py new file mode 100644 index 000000000..87d9845ca --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py @@ -0,0 +1,115 @@ +import math + +import torch + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + emb = scale * emb + + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class TimestepEmbedding(torch.nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + out_dim: int | None = None, + post_act_fn: str | None = None, + cond_proj_dim: int | None = None, + sample_proj_bias: bool = True, + ): + super().__init__() + + self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + + if cond_proj_dim is not None: + self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = torch.nn.SiLU() + time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim + + self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + + if post_act_fn is None: + self.post_act = None + + def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Timesteps(torch.nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + return t_emb + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): + def __init__( + self, + embedding_dim: int, + size_emb_dim: int, + ): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward( + self, + timestep: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + return timesteps_emb diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py b/diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py new file mode 100644 index 000000000..1b979e5a5 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +from typing import Literal + +import torch +from torch import nn + +from diffsynth.models.ltx25_diffvae.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings +from diffsynth.models.ltx25_diffvae.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer import ( + AdaLNZero, + ChannelLinear, + CombinedDiffusionNABlock, + LinearPixelShuffleUpsample, + NABlock, +) + + +class DiffusionVideoDecoder(nn.Module): + def __init__( + self, + in_channels: int = 128, + out_channels: int = 3, + patch_size: int = 4, + head_dim: int = 64, + rope_dim_split: tuple[int, int, int] | None = None, + stage_channels: tuple[int, ...] = (1024, 512, 256, 256, 128), + stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8), + stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5), (3, 7, 7)), + upsamples: tuple[tuple[tuple[int, int, int], int], ...] = (((1, 2, 2), 2), ((2, 1, 1), 2), ((2, 2, 2), 1), ((2, 2, 2), 2)), + stage5_kernel: tuple[int, int, int] | None = None, + stage5_channels: int | None = None, + t_emb_dim: int = 384, + default_num_inference_steps: int = 1, + timestep_scale_multiplier: float = 1.0, + model_output_type: Literal["v", "x0"] = "x0", + ) -> None: + super().__init__() + if len(stage_channels) != len(stage_depths) or len(stage_channels) != len(stage_kernels): + raise ValueError("stage_channels, stage_depths, and stage_kernels must have the same length") + if len(upsamples) != len(stage_channels) - 1: + raise ValueError("one fewer upsample than decoder stages is required") + if any(channels % head_dim for channels in stage_channels): + raise ValueError("every stage channel count must be divisible by head_dim") + + self.patch_size = patch_size + self.out_channels = out_channels + self.stage_kernels = stage_kernels + self.upsample_strides = tuple(stride for stride, _ in upsamples) + self.stage5_kernel = tuple(stage5_kernel or stage_kernels[-1]) + self.model_output_type = model_output_type + self.timestep_scale_multiplier = timestep_scale_multiplier + self.register_buffer( + "default_inference_timesteps", + torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps), + persistent=False, + ) + + self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) + self.conv_in = ChannelLinear(in_channels, stage_channels[0], bias=True) + self.det_stages = nn.ModuleList() + self.upsamples = nn.ModuleList() + for channels, depth, kernel, (stride, reduction) in zip( + stage_channels[:-1], stage_depths[:-1], stage_kernels[:-1], upsamples, strict=True + ): + self.det_stages.append( + nn.ModuleList( + [NABlock(channels, kernel, head_dim=head_dim, rope_dim_split=rope_dim_split) for _ in range(depth)] + ) + ) + self.upsamples.append(LinearPixelShuffleUpsample(channels, stride, reduction)) + + context_channels = stage_channels[-1] + diffusion_channels = stage5_channels or context_channels + if diffusion_channels % head_dim: + raise ValueError("stage5_channels must be divisible by head_dim") + self.context_channels = context_channels + self.t_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=t_emb_dim, size_emb_dim=0) + self.conv_in_x_t = ChannelLinear(out_channels * patch_size**2, diffusion_channels, bias=True) + self.shared_adaln = AdaLNZero(dim=diffusion_channels, t_emb_dim=t_emb_dim) + self.diff_blocks = nn.ModuleList( + [ + CombinedDiffusionNABlock( + diffusion_channels, + self.stage5_kernel, + context_channels, + head_dim=head_dim, + rope_dim_split=rope_dim_split, + ) + for _ in range(stage_depths[-1]) + ] + ) + self.norm_out = nn.RMSNorm(diffusion_channels, eps=1e-6) + self.conv_out = ChannelLinear(diffusion_channels, out_channels * patch_size**2, bias=True) + + self.min_latent_shape = self._minimum_latent_shape() + + self._trailing_latent_frames = (stage_kernels[0][0] // 2) * 2 + + def _minimum_latent_shape(self) -> tuple[int, int, int]: + cumulative = [1, 1, 1] + minimum = [1, 1, 1] + for kernel, stride in zip(self.stage_kernels[:-1], self.upsample_strides, strict=True): + for axis in range(3): + minimum[axis] = max(minimum[axis], -(-kernel[axis] // cumulative[axis])) + cumulative[axis] *= stride[axis] + for axis in range(3): + minimum[axis] = max(minimum[axis], -(-self.stage5_kernel[axis] // cumulative[axis])) + return tuple(minimum) + + @staticmethod + def _pad_axis(x: torch.Tensor, axis: int, size: int, *, trailing_only: bool) -> tuple[torch.Tensor, int]: + current = x.shape[axis] + if current >= size: + return x, 0 + missing = size - current + before = 0 if trailing_only else missing // 2 + after = missing - before + pieces = [] + if before: + pieces.append(x.narrow(axis, 0, 1).expand(*x.shape[:axis], before, *x.shape[axis + 1 :])) + pieces.append(x) + if after: + pieces.append(x.narrow(axis, current - 1, 1).expand(*x.shape[:axis], after, *x.shape[axis + 1 :])) + return torch.cat(pieces, dim=axis), before + + def _pad_to_minimum(self, latent: torch.Tensor) -> tuple[torch.Tensor, int, int]: + latent, _ = self._pad_axis(latent, 2, self.min_latent_shape[0], trailing_only=True) + latent, h_before = self._pad_axis(latent, 3, self.min_latent_shape[1], trailing_only=False) + latent, w_before = self._pad_axis(latent, 4, self.min_latent_shape[2], trailing_only=False) + return latent, h_before, w_before + + def _run_det_stage(self, x: torch.Tensor, stage_index: int, drop_leading_frame: bool) -> torch.Tensor: + for block in self.det_stages[stage_index]: + x = block(x) + return self.upsamples[stage_index](x, drop_leading_frame=drop_leading_frame) + + def forward_stages_1_to_3(self, z_noisy: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: + x = self.per_channel_statistics.un_normalize(z_noisy).permute(0, 2, 3, 4, 1) + x = self.conv_in(x) + for stage_index in range(3): + x = self._run_det_stage(x, stage_index, drop_leading_frame) + return x + + def forward_stage_4( + self, + x: torch.Tensor, + drop_leading_frame: bool = True, + pad_trailing: bool = True, + ) -> torch.Tensor: + x = self._run_det_stage(x, 3, drop_leading_frame) + if pad_trailing and self._trailing_latent_frames: + ghost_frames = self._trailing_latent_frames * 8 + keep = min(x.shape[1], max(x.shape[1] - ghost_frames, self.stage5_kernel[0])) + x = x[:, :keep] + return x + + def _context_and_x_for_diff_step(self, context: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor: + pixels = patchify(x_t, patch_size_hw=self.patch_size).permute(0, 2, 3, 4, 1) + return torch.cat([context, self.conv_in_x_t(pixels)], dim=-1) + + def forward_diff_step(self, context_and_x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + x = context_and_x[..., self.context_channels :] + modulation = self.shared_adaln(self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype)) + for block in self.diff_blocks: + x = block(context_and_x, modulation) + context_and_x[..., self.context_channels :].copy_(x) + x = self.conv_out(self.norm_out(x)).permute(0, 4, 1, 2, 3).contiguous() + return unpatchify(x, patch_size_hw=self.patch_size) + + def _euler_step(self, x_t: torch.Tensor, model_out: torch.Tensor, t_now: torch.Tensor, t_next: torch.Tensor) -> torch.Tensor: + dt = (t_now - t_next).view(-1, *([1] * (x_t.ndim - 1))).to(torch.float32) + if self.model_output_type == "x0": + model_out = (x_t.to(torch.float32) - model_out.to(torch.float32)) / t_now.view( + -1, *([1] * (x_t.ndim - 1)) + ).to(torch.float32) + return (x_t.to(torch.float32) - dt * model_out.to(torch.float32)).to(x_t.dtype) + + def forward(self, sample: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + frames = (sample.shape[2] - 1) * 8 + 1 + height, width = sample.shape[3] * 32, sample.shape[4] * 32 + latent, h_before, w_before = self._pad_to_minimum(sample) + trailing = latent[:, :, -1:].expand(-1, -1, self._trailing_latent_frames, -1, -1) + context = self.forward_stages_1_to_3(torch.cat([latent, trailing], dim=2)) + context = self.forward_stage_4(context) + + pixel_shape = (sample.shape[0], self.out_channels, context.shape[1], context.shape[2] * self.patch_size, context.shape[3] * self.patch_size) + x_t = torch.randn(pixel_shape, dtype=sample.dtype, device=sample.device, generator=generator) + timesteps = self.default_inference_timesteps.to(sample.device).expand(sample.shape[0], -1) + for index in range(timesteps.shape[1] - 1): + prediction = self.forward_diff_step(self._context_and_x_for_diff_step(context, x_t), timesteps[:, index]) + x_t = self._euler_step(x_t, prediction, timesteps[:, index], timesteps[:, index + 1]) + prediction = self.forward_diff_step(self._context_and_x_for_diff_step(context, x_t), timesteps[:, -1]) + pixels = prediction if self.model_output_type == "x0" else self._euler_step(x_t, prediction, timesteps[:, -1], torch.zeros_like(timesteps[:, -1])) + return pixels[:, :, :frames, h_before * 32 : h_before * 32 + height, w_before * 32 : w_before * 32 + width].contiguous() + + def decode( + self, + latent: torch.Tensor, + tiled: bool = True, + tile_size_in_pixels: int = 512, + tile_overlap_in_pixels: int = 128, + tile_size_in_frames: int = 128, + tile_overlap_in_frames: int = 24, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + del tiled, tile_size_in_pixels, tile_overlap_in_pixels, tile_size_in_frames, tile_overlap_in_frames + return self.forward(latent, generator=generator) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/ops.py b/diffsynth/models/ltx25_diffvae/model/video_vae/ops.py new file mode 100644 index 000000000..f945dfa79 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/ops.py @@ -0,0 +1,57 @@ +import torch +from einops import rearrange +from torch import nn + + +def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + if patch_size_hw == 1 and patch_size_t == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b c (f p) (h q) (w r) -> b (c p r q) f h w", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + + +def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + if patch_size_hw == 1 and patch_size_t == 1: + return x + + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b (c p r q) f h w -> b c (f p) (h q) (w r)", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + + return x + + +class PerChannelStatistics(nn.Module): + def __init__(self, latent_channels: int = 128): + super().__init__() + self.register_buffer("std-of-means", torch.ones(latent_channels)) + self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py new file mode 100644 index 000000000..1f9d3b95d --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py @@ -0,0 +1,3 @@ +from .blocks import NABlock +from .combined.block import CombinedDiffusionNABlock +from .layers import AdaLNZero, ChannelLinear, LinearPixelShuffleUpsample diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py new file mode 100644 index 000000000..e52d0c304 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py @@ -0,0 +1,59 @@ +import torch +from torch import nn + +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.det_attn_rope import det_qkv_rope +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.fallback_na import EagerSdpaAttention +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.qkv import QKVProjections +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.rope_math import ( + DEFAULT_ABS_ROPE_NUM_TILES, + default_rope_dim_split, + rope_inv_freqs, +) + + +class NeighborhoodAttention3D(nn.Module): + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + head_dim: int = 64, + rope_dim_split: tuple[int, int, int] | None = None, + rope_base: float = 10000.0, + ) -> None: + super().__init__() + if dim % head_dim: + raise ValueError(f"dim={dim} must be divisible by head_dim={head_dim}") + rope_dim_split = rope_dim_split or default_rope_dim_split(head_dim) + if sum(rope_dim_split) != head_dim: + raise ValueError(f"rope_dim_split={rope_dim_split} must sum to head_dim={head_dim}") + + self.dim = dim + self.num_heads = dim // head_dim + self.head_dim = head_dim + self.kernel_size = tuple(kernel_size) + self.scale = head_dim**-0.5 + self.rope_dim_split = rope_dim_split + self.rope_num_tiles = DEFAULT_ABS_ROPE_NUM_TILES + self.rope_compute_dtype = torch.float32 + self.attention_function = EagerSdpaAttention() + self.register_buffer("rope_inv_t", rope_inv_freqs(rope_dim_split[0], rope_base), persistent=False) + self.register_buffer("rope_inv_h", rope_inv_freqs(rope_dim_split[1], rope_base), persistent=False) + self.register_buffer("rope_inv_w", rope_inv_freqs(rope_dim_split[2], rope_base), persistent=False) + self.qkv = QKVProjections(dim) + self.proj = nn.Linear(dim, dim, bias=True) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + + def project_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch, frames, height, width, _ = x.shape + q, k, v = self.qkv(x) + shape = (batch, frames, height, width, self.num_heads, self.head_dim) + return q.view(shape), k.view(shape), v.view(shape) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch, frames, height, width, _ = x.shape + if any(size < kernel for size, kernel in zip((frames, height, width), self.kernel_size, strict=True)): + raise ValueError(f"input {(frames, height, width)} is smaller than neighborhood kernel {self.kernel_size}") + q, k, v = det_qkv_rope(self, x) + output = self.attention_function(self, q.contiguous(), k.contiguous(), v.contiguous()) + return self.proj(output.reshape(batch, frames, height, width, self.dim)) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py new file mode 100644 index 000000000..926f124bc --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import torch +from torch import nn + +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.attention import NeighborhoodAttention3D +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.layers import AdaLNZero +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.swiglu import SwiGLU, plain_mlp + +__all__ = [ + "DiffusionNABlock", + "NABlock", +] + + +class NABlock(nn.Module): + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + head_dim: int = 64, + mlp_ratio: float = 4.0, + rope_dim_split: tuple[int, int, int] | None = None, + ) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(dim, eps=1e-6) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) + self.norm2 = nn.RMSNorm(dim, eps=1e-6) + hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 + self.mlp = SwiGLU(dim, hidden) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x)) + x = plain_mlp(x, self.mlp, self.norm2) + return x + + +class DiffusionNABlock(nn.Module): + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + context_channels: int, + head_dim: int = 64, + mlp_ratio: float = 4.0, + rope_dim_split: tuple[int, int, int] | None = None, + ) -> None: + super().__init__() + self.context_channels = context_channels + self.context_proj = nn.Linear(context_channels, dim, bias=True) + self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim)) + + self.norm1 = nn.RMSNorm(dim, eps=1e-6) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) + self.norm2 = nn.RMSNorm(dim, eps=1e-6) + hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 + self.mlp = SwiGLU(dim, hidden) + self.attn.proj.reset_parameters() + + def _modulation( + self, modulation: tuple[torch.Tensor, ...] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) + ] + return scale_msa, shift_msa, scale_mlp, shift_mlp diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py new file mode 100644 index 000000000..a38ec13af --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import torch +from torch import nn + +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.attention import NeighborhoodAttention3D +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.rope_math import ( + h_positions, + rot_abs_axis_impl, + t_positions, +) + + +_rot_abs_axis = torch.compiler.nested_compile_region(rot_abs_axis_impl) + + +def _apply_nested_abs_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype, +) -> torch.Tensor: + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + xt = _rot_abs_axis(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) + xh = _rot_abs_axis( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = _rot_abs_axis(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _apply_nested_full_volume_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + _apply_nested_abs_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +def _qkv_nested_rope(attn: NeighborhoodAttention3D, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = attn.project_qkv(x) + q = attn.q_norm(q) * attn.scale + k = attn.k_norm(k) + inv_freqs = ( + attn.rope_inv_t.to(device=x.device), + attn.rope_inv_h.to(device=x.device), + attn.rope_inv_w.to(device=x.device), + ) + q = _apply_nested_full_volume_rope( + q, + attn.rope_dim_split, + inv_freqs, + num_tiles=attn.rope_num_tiles, + compute_dtype=attn.rope_compute_dtype, + ) + k = _apply_nested_full_volume_rope( + k, + attn.rope_dim_split, + inv_freqs, + num_tiles=attn.rope_num_tiles, + compute_dtype=attn.rope_compute_dtype, + ) + return q, k, v + + +def full( + x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + y = norm(x) * (1.0 + scale) + shift + batch, t, h, w, _ = y.shape + kt, kh, kw = attn.kernel_size + if t < kt or h < kh or w < kw: + raise ValueError( + f"3D neighborhood attention requires spatial dims >= kernel_size; " + f"got (T,H,W)=({t},{h},{w}) vs kernel={attn.kernel_size}" + ) + + q, k, v = _qkv_nested_rope(attn, y) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = attn.attention_function(attn, q, k, v) + out = out.reshape(batch, t, h, w, attn.dim) + return x + attn.proj(out) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py new file mode 100644 index 000000000..aeb88051d --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import torch + +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.blocks import DiffusionNABlock +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.attn import full as residual_attn +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.context import combined as inject_context +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.mlp import residual_mlp + + +class CombinedDiffusionNABlock(DiffusionNABlock): + def forward_combined( + self, + context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) + x = residual_attn(x, self.attn, self.norm1, scale_msa, shift_msa) + x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp) + return x + + def forward( + self, + context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + return self.forward_combined(context_and_x, modulation) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py new file mode 100644 index 000000000..783ec42b3 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def combined( + context_and_x: torch.Tensor, + w_proj: torch.Tensor, + b_proj: torch.Tensor | None, +) -> torch.Tensor: + context_channels = w_proj.shape[1] + latent_context = context_and_x[..., :context_channels] + x = context_and_x[..., context_channels:] + return x + F.linear(latent_context, w_proj, b_proj) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py new file mode 100644 index 000000000..e3bb98760 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py @@ -0,0 +1,14 @@ +import torch +from torch import nn + +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.layers import modulate + + +def residual_mlp( + x: torch.Tensor, + mlp: nn.Module, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + return x + mlp(modulate(norm(x), scale, shift)) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py new file mode 100644 index 000000000..f8a4c2529 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import torch + +from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.rope_math import ( + DEFAULT_ABS_ROPE_NUM_TILES, + h_positions, + rot_abs_axis_impl, + t_positions, +) + + +def _apply_opaque_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype, +) -> torch.Tensor: + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + xt = rot_abs_axis_impl(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) + xh = rot_abs_axis_impl( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _apply_opaque_tiled_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + _apply_opaque_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +@torch.library.custom_op("diffsynth_ltx25::abs_rope", mutates_args=()) +def _abs_rope_op( + x: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 + return _apply_opaque_tiled_rope( + x, + (d_t, d_h, d_w), + (inv_t, inv_h, inv_w), + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + + +@_abs_rope_op.register_fake +def _abs_rope_fake( + x: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + del inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 + return torch.empty(x.shape, device=x.device, dtype=x.dtype) + + +def _apply_opaque_abs_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"compute_dtype must be float32 or bfloat16, got {compute_dtype}") + d_t, d_h, d_w = rope_split + inv_t, inv_h, inv_w = inv_freqs + return _abs_rope_op( + x, + inv_t, + inv_h, + inv_w, + d_t, + d_h, + d_w, + num_tiles, + compute_dtype == torch.bfloat16, + ) + + +def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = attn.project_qkv(x) + q = attn.q_norm(q) + k = attn.k_norm(k) + q = q * attn.scale + + inv_freqs = ( + attn.rope_inv_t.to(device=x.device), + attn.rope_inv_h.to(device=x.device), + attn.rope_inv_w.to(device=x.device), + ) + num_tiles = getattr(attn, "rope_num_tiles", DEFAULT_ABS_ROPE_NUM_TILES) + compute_dtype = getattr(attn, "rope_compute_dtype", torch.float32) + q = _apply_opaque_abs_rope( + q, + attn.rope_dim_split, + inv_freqs, + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + k = _apply_opaque_abs_rope( + k, + attn.rope_dim_split, + inv_freqs, + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + return q, k, v diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py new file mode 100644 index 000000000..c522eb352 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py @@ -0,0 +1,8 @@ +from .eager import na3d + + +class EagerSdpaAttention: + def __call__(self, attn, q, k, v): + if q.dtype != v.dtype or k.dtype != v.dtype: + q, k = q.to(dtype=v.dtype), k.to(dtype=v.dtype) + return na3d(q, k, v, kernel_size=attn.kernel_size, is_causal=None, scale=1.0) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py new file mode 100644 index 000000000..bd5154c23 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import math + +import torch +from torch.nn import functional + + +NA_SCORE_BUDGET = 2**25 + +NA_KV_STACK_BUDGET = 2**28 + + +def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], list[int]]: + starts: list[int] = [] + ends: list[int] = [] + if causal: + for i in range(length): + starts.append(max(0, i - kernel + 1)) + ends.append(i + 1) + else: + kernel = min(kernel, length) + lo = length - kernel + half = kernel // 2 + for i in range(length): + start = min(max(i - half, 0), lo) + starts.append(start) + ends.append(start + kernel) + return starts, ends + + +def _pick_tiles(dims: tuple[int, int, int], kernels: list[int]) -> list[int]: + tiles = list(dims) + + def cost(ts: list[int]) -> int: + nq = math.prod(ts) + nk = math.prod(min(d, t + k - 1) for t, k, d in zip(ts, kernels, dims, strict=True)) + return nq * nk + + while cost(tiles) > NA_SCORE_BUDGET and max(tiles) > 1: + i = max(range(3), key=lambda a: tiles[a] / kernels[a]) + if tiles[i] <= 1: + break + tiles[i] = max(1, (tiles[i] + 1) // 2) + return tiles + + +def _group_mask( + rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + bools = [] + for starts, ends in rel_bounds: + st = torch.tensor(starts, device=device) + en = torch.tensor(ends, device=device) + kj = torch.arange(int(en.max()), device=device) + bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) + visible = ( + bools[0][:, None, None, :, None, None] + & bools[1][None, :, None, None, :, None] + & bools[2][None, None, :, None, None, :] + ) + nq = visible.shape[0] * visible.shape[1] * visible.shape[2] + nk = visible.shape[3] * visible.shape[4] * visible.shape[5] + mask = torch.zeros((nq, nk), dtype=dtype, device=device) + mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) + return mask.reshape(1, 1, nq, nk) + + +def na3d( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int] | tuple[int, ...], + is_causal: list[bool] | None = None, + scale: float | None = None, +) -> torch.Tensor: + batch, t, h, w, nh, hd = q.shape + dims = (t, h, w) + causal = [False, False, False] if is_causal is None else list(is_causal) + kernels = [k_ if c else min(k_, d) for k_, c, d in zip(kernel_size, causal, dims, strict=True)] + if scale is None: + scale = hd**-0.5 + device = q.device + if scale != 1.0: + q = q * scale + + bounds = [_window_bounds(d, k_, c) for d, k_, c in zip(dims, kernels, causal, strict=True)] + tile_t, tile_h, tile_w = _pick_tiles(dims, [min(k_, d) for k_, d in zip(kernels, dims, strict=True)]) + + groups: dict[ + tuple[ + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + ], + list[tuple[tuple[slice, slice, slice], tuple[slice, slice, slice]]], + ] = {} + for t0 in range(0, t, tile_t): + t1 = min(t0 + tile_t, t) + rt0, rt1 = bounds[0][0][t0], bounds[0][1][t1 - 1] + rel_t = ( + tuple(s - rt0 for s in bounds[0][0][t0:t1]), + tuple(e - rt0 for e in bounds[0][1][t0:t1]), + ) + for h0 in range(0, h, tile_h): + h1 = min(h0 + tile_h, h) + rh0, rh1 = bounds[1][0][h0], bounds[1][1][h1 - 1] + rel_h = ( + tuple(s - rh0 for s in bounds[1][0][h0:h1]), + tuple(e - rh0 for e in bounds[1][1][h0:h1]), + ) + for w0 in range(0, w, tile_w): + w1 = min(w0 + tile_w, w) + rw0, rw1 = bounds[2][0][w0], bounds[2][1][w1 - 1] + rel_w = ( + tuple(s - rw0 for s in bounds[2][0][w0:w1]), + tuple(e - rw0 for e in bounds[2][1][w0:w1]), + ) + groups.setdefault((rel_t, rel_h, rel_w), []).append( + ( + (slice(t0, t1), slice(h0, h1), slice(w0, w1)), + (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), + ) + ) + + out = torch.empty((batch, t, h, w, nh, hd), device=device, dtype=v.dtype) + for rel, tiles in groups.items(): + mask = _group_mask(rel, q.dtype, device) + nq, nk = mask.shape[2], mask.shape[3] + g_max = max(1, NA_KV_STACK_BUDGET // max(1, batch * nh * nk * hd * 2)) if device.type == "cuda" else 1 + qs0, _ = tiles[0] + tq = qs0[0].stop - qs0[0].start + th = qs0[1].stop - qs0[1].start + tw = qs0[2].stop - qs0[2].start + for c0 in range(0, len(tiles), g_max): + chunk = tiles[c0 : c0 + g_max] + g = len(chunk) + q_s = torch.stack([q[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) + k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nq, hd) + k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + o = functional.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) + o = o.view(g, batch, nh, tq, th, tw, hd).permute(0, 1, 3, 4, 5, 2, 6) + for i, (qs, _) in enumerate(chunk): + out[:, qs[0], qs[1], qs[2]] = o[i] + + return out diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py new file mode 100644 index 000000000..f879e247f --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn + + +class ChannelLinear(nn.Linear): + + @property + def in_channels(self) -> int: + return self.in_features + + @property + def out_channels(self) -> int: + return self.out_features + + +class LinearPixelShuffleUpsample(nn.Module): + def __init__( + self, + in_channels: int, + stride: tuple[int, int, int], + out_channels_reduction_factor: int = 1, + ) -> None: + super().__init__() + self.stride = stride + self.proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor + self.out_channels = self.proj_out_channels // math.prod(stride) + self.proj = nn.Linear(in_channels, self.proj_out_channels, bias=True) + + def forward(self, x: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: + x = self.proj(x) + x = rearrange( + x, + "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + if self.stride[0] == 2 and drop_leading_frame: + x = x[:, 1:, :, :, :] + return x + + +class AdaLNZero(nn.Module): + NUM_CHUNKS: int = 7 + + def __init__(self, dim: int, t_emb_dim: int) -> None: + super().__init__() + self.dim = dim + self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True) + nn.init.zeros_(self.proj.weight) + nn.init.zeros_(self.proj.bias) + + def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: + h = self.proj(F.silu(t_emb)) + chunks = h.chunk(self.NUM_CHUNKS, dim=-1) + return tuple(c[:, None, None, None, :] for c in chunks) + + +def modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + return x * (1.0 + scale) + shift diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py new file mode 100644 index 000000000..6096dfad0 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +import torch +from torch import nn + + +class QKVProjections(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + self.to_q = nn.Linear(dim, dim, bias=True) + self.to_k = nn.Linear(dim, dim, bias=True) + self.to_v = nn.Linear(dim, dim, bias=True) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.to_q(x), self.to_k(x), self.to_v(x) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py new file mode 100644 index 000000000..bac4fb16f --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np +import torch + +DEFAULT_ABS_ROPE_NUM_TILES = 4 + + +def t_positions(t: int, device: torch.device) -> torch.Tensor: + return torch.arange(t, dtype=torch.float32, device=device) + + +def h_positions(h: int, device: torch.device) -> torch.Tensor: + return torch.arange(h, dtype=torch.float32, device=device) + + +def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: + assert head_dim % 8 == 0, f"head_dim={head_dim} must be a multiple of 8 for default split" + d_t = (head_dim // 4) // 2 * 2 + d_hw = (head_dim - d_t) // 2 + if d_hw % 2 != 0: + d_t -= 2 + d_hw = (head_dim - d_t) // 2 + assert d_t > 0 + assert d_hw > 0 + return (d_t, d_hw, d_hw) + + +def rope_inv_freqs(dim: int, base: float = 10000.0) -> torch.Tensor: + assert dim % 2 == 0, f"RoPE dim must be even, got {dim}" + exponents = np.arange(0, dim, 2, dtype=np.float64) / dim + inv_freqs = 1.0 / np.power(float(base), exponents) + return torch.from_numpy(inv_freqs).to(torch.float32) + + +def rot_abs_axis_impl( + xc: torch.Tensor, + pos: torch.Tensor, + inv: torch.Tensor, + axis: int, + *, + compute_dtype: torch.dtype, +) -> torch.Tensor: + out_dtype = xc.dtype + pairs = xc.reshape(*xc.shape[:-1], xc.shape[-1] // 2, 2) + xe = pairs[..., 0].to(compute_dtype) + xo = pairs[..., 1].to(compute_dtype) + shape = [1, 1, 1, 1, 1, inv.shape[0]] + shape[axis] = pos.shape[0] + ang = (pos[:, None] * inv[None, :]).reshape(shape) + c = ang.cos().to(compute_dtype) + s = ang.sin().to(compute_dtype) + re = xe * c - xo * s + ro = xe * s + xo * c + out = torch.stack([re, ro], dim=-1).reshape(xc.shape) + return out.to(out_dtype) if out.dtype != out_dtype else out diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py new file mode 100644 index 000000000..d9ad9e0f1 --- /dev/null +++ b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py @@ -0,0 +1,38 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +_DEFAULT_TOKEN_CHUNK = 16_384 + + +def swiglu(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor) -> torch.Tensor: + if x.dtype != w_gate.dtype: + x = x.to(w_gate.dtype) + leading, dim = x.shape[:-1], x.shape[-1] + flat = x.reshape(-1, dim).contiguous() + output = torch.empty_like(flat) + for start in range(0, flat.shape[0], _DEFAULT_TOKEN_CHUNK): + end = min(start + _DEFAULT_TOKEN_CHUNK, flat.shape[0]) + tokens = flat[start:end] + workspace = torch.empty((end - start, w_gate.shape[0]), dtype=x.dtype, device=x.device) + torch.mm(tokens, w_gate.t(), out=workspace) + F.silu(workspace, inplace=True) + workspace.mul_(F.linear(tokens, w_up)) + torch.mm(workspace, w_down.t(), out=output[start:end]) + return output.view(*leading, dim) + + +class SwiGLU(nn.Module): + def __init__(self, dim: int, hidden_dim: int) -> None: + super().__init__() + self.w_up = nn.Linear(dim, hidden_dim, bias=False) + self.w_gate = nn.Linear(dim, hidden_dim, bias=False) + self.w_down = nn.Linear(hidden_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return swiglu(x, self.w_gate.weight, self.w_up.weight, self.w_down.weight) + + +def plain_mlp(x: torch.Tensor, mlp: SwiGLU, norm: nn.RMSNorm) -> torch.Tensor: + return x + mlp(norm(x)) diff --git a/diffsynth/models/ltx25_duration_head.py b/diffsynth/models/ltx25_duration_head.py new file mode 100644 index 000000000..9fdf01ccc --- /dev/null +++ b/diffsynth/models/ltx25_duration_head.py @@ -0,0 +1,57 @@ +import torch +from torch import nn + + +class LTX25AttentionPooler(nn.Module): + def __init__(self, hidden_dim: int = 256, num_queries: int = 1, num_heads: int = 4): + super().__init__() + self.hidden_dim = hidden_dim + self.num_queries = num_queries + self.query_tokens = nn.Parameter(torch.randn(num_queries, hidden_dim) * 0.02) + self.cross_attn = nn.MultiheadAttention( + embed_dim=hidden_dim, + num_heads=num_heads, + batch_first=True, + ) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + queries = self.query_tokens.unsqueeze(0).expand(tokens.shape[0], -1, -1) + pooled, _ = self.cross_attn(queries, tokens, tokens, need_weights=False) + return pooled + + +class LTX25DurationHead(nn.Module): + def __init__( + self, + video_cross_attention_dim: int = 4096, + audio_cross_attention_dim: int = 2048, + pooler_hidden_dim: int = 256, + num_queries: int = 1, + num_pooler_heads: int = 4, + mlp_hidden: int = 256, + ): + super().__init__() + self.pooler_hidden_dim = pooler_hidden_dim + self.video_input_proj = nn.Linear(video_cross_attention_dim, pooler_hidden_dim) + self.video_modality_emb = nn.Parameter(torch.randn(pooler_hidden_dim) * 0.02) + self.audio_input_proj = nn.Linear(audio_cross_attention_dim, pooler_hidden_dim) + self.audio_modality_emb = nn.Parameter(torch.randn(pooler_hidden_dim) * 0.02) + self.attention_pooler = LTX25AttentionPooler(pooler_hidden_dim, num_queries, num_pooler_heads) + self.mlp_hidden = nn.Linear(pooler_hidden_dim * num_queries, mlp_hidden) + self.mlp_out = nn.Linear(mlp_hidden, 1) + + def forward( + self, + video_tokens: torch.Tensor | None = None, + audio_tokens: torch.Tensor | None = None, + ) -> torch.Tensor: + if video_tokens is None and audio_tokens is None: + raise ValueError("LTX25DurationHead.forward requires video_tokens and/or audio_tokens.") + token_groups = [] + if video_tokens is not None: + token_groups.append(self.video_input_proj(video_tokens) + self.video_modality_emb) + if audio_tokens is not None: + token_groups.append(self.audio_input_proj(audio_tokens) + self.audio_modality_emb) + pooled = self.attention_pooler(torch.cat(token_groups, dim=1)) + hidden = torch.nn.functional.gelu(self.mlp_hidden(pooled.reshape(pooled.shape[0], -1)), approximate="tanh") + return self.mlp_out(hidden).squeeze(-1).exp() diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py new file mode 100644 index 000000000..075aae75e --- /dev/null +++ b/diffsynth/models/ltx25_text_encoder.py @@ -0,0 +1,423 @@ +import copy +import math +from typing import NamedTuple + +import torch + +from .ltx2_common import rms_norm +from .ltx2_dit import ( + Attention, + FeedForward, + LTXRopeType, + generate_freq_grid_np, + generate_freq_grid_pytorch, + precompute_freqs_cis, +) + + +LTX25_GEMMA_CONFIG = {'architectures': ['Gemma4UnifiedForConditionalGeneration'], + 'audio_config': {'_name_or_path': '', + 'architectures': None, + 'audio_embed_dim': 640, + 'chunk_size_feed_forward': 0, + 'dtype': 'bfloat16', + 'id2label': {'0': 'LABEL_0', '1': 'LABEL_1'}, + 'initializer_range': 0.02, + 'is_encoder_decoder': False, + 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, + 'model_type': 'gemma4_unified_audio', + 'output_attentions': False, + 'output_hidden_states': False, + 'problem_type': None, + 'return_dict': True, + 'rms_norm_eps': 1e-06}, + 'audio_token_id': 258881, + 'boa_token_id': 256000, + 'boi_token_id': 255999, + 'dtype': 'bfloat16', + 'eoa_token_index': 258883, + 'eoi_token_id': 258882, + 'eos_token_id': [1, 106], + 'gemma_version': 'gemma4-12b-ltx-v1', + 'image_token_id': 258880, + 'initializer_range': 0.02, + 'model_type': 'gemma4_unified', + 'text_config': {'attention_bias': False, + 'attention_dropout': 0.0, + 'attention_k_eq_v': True, + 'bos_token_id': 2, + 'dtype': 'bfloat16', + 'enable_moe_block': False, + 'eos_token_id': 1, + 'final_logit_softcapping': 30.0, + 'global_head_dim': 512, + 'head_dim': 256, + 'hidden_activation': 'gelu_pytorch_tanh', + 'hidden_size': 3840, + 'hidden_size_per_layer_input': 0, + 'initializer_range': 0.02, + 'intermediate_size': 15360, + 'layer_types': ['sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'sliding_attention', + 'full_attention'], + 'max_position_embeddings': 262144, + 'model_type': 'gemma4_unified_text', + 'moe_intermediate_size': None, + 'num_attention_heads': 16, + 'num_experts': None, + 'num_global_key_value_heads': 1, + 'num_hidden_layers': 48, + 'num_key_value_heads': 8, + 'num_kv_shared_layers': 0, + 'pad_token_id': 0, + 'rms_norm_eps': 1e-06, + 'rope_parameters': {'full_attention': {'partial_rotary_factor': 0.25, + 'rope_theta': 1000000.0, + 'rope_type': 'proportional'}, + 'sliding_attention': {'rope_theta': 10000.0, 'rope_type': 'default'}}, + 'sliding_window': 1024, + 'tie_word_embeddings': True, + 'top_k_experts': None, + 'use_bidirectional_attention': 'vision', + 'use_cache': True, + 'use_double_wide_mlp': False, + 'vocab_size': 262144, + 'vocab_size_per_layer_input': 262144}, + 'tie_word_embeddings': True, + 'transformers_version': '5.10.1', + 'video_token_id': 258884, + 'vision_config': {'_name_or_path': '', + 'architectures': None, + 'chunk_size_feed_forward': 0, + 'dtype': 'bfloat16', + 'id2label': {'0': 'LABEL_0', '1': 'LABEL_1'}, + 'initializer_range': 0.02, + 'is_encoder_decoder': False, + 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, + 'mm_embed_dim': 3840, + 'mm_posemb_size': 1120, + 'model_type': 'gemma4_unified_vision', + 'num_soft_tokens': 280, + 'output_attentions': False, + 'output_hidden_states': False, + 'output_proj_dims': 3840, + 'patch_size': 16, + 'pooling_kernel_size': 3, + 'problem_type': None, + 'return_dict': True, + 'rms_norm_eps': 1e-06}} + + +class LTX25TextEncoder(torch.nn.Module): + def __init__(self): + super().__init__() + from transformers import Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration + + self.config = Gemma4UnifiedConfig(**copy.deepcopy(LTX25_GEMMA_CONFIG)) + self.model = Gemma4UnifiedForConditionalGeneration(self.config) + + def forward(self, *args, **kwargs): + return self.model(*args, **kwargs) + + +def norm_and_concat_per_token_rms( + encoded_text: torch.Tensor, + attention_mask: torch.Tensor, +) -> torch.Tensor: + batch_size, sequence_length, embedding_dim, num_layers = encoded_text.shape + variance = torch.mean(encoded_text**2, dim=2, keepdim=True) + normed = encoded_text * torch.rsqrt(variance + 1e-6) + normed = normed.reshape(batch_size, sequence_length, embedding_dim * num_layers) + return torch.where(attention_mask.bool().unsqueeze(-1), normed, torch.zeros_like(normed)) + + +def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor: + return x * math.sqrt(target_dim / source_dim) + + +class LTX25FeatureExtractorV2(torch.nn.Module): + def __init__( + self, + video_aggregate_embed: torch.nn.Linear, + embedding_dim: int, + audio_aggregate_embed: torch.nn.Linear | None = None, + ): + super().__init__() + self.video_aggregate_embed = video_aggregate_embed + self.audio_aggregate_embed = audio_aggregate_embed + self.embedding_dim = embedding_dim + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + padding_side: str = "left", + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del padding_side + encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states + normed = norm_and_concat_per_token_rms(encoded, attention_mask).to(encoded.dtype) + video = self.video_aggregate_embed( + _rescale_norm(normed, self.video_aggregate_embed.out_features, self.embedding_dim) + ) + audio = None + if self.audio_aggregate_embed is not None: + audio = self.audio_aggregate_embed( + _rescale_norm(normed, self.audio_aggregate_embed.out_features, self.embedding_dim) + ) + return video, audio + + +class LTX25Embeddings1DConnector(torch.nn.Module): + def __init__( + self, + attention_head_dim: int, + num_attention_heads: int, + num_layers: int, + positional_embedding_theta: float = 10000.0, + positional_embedding_max_pos: list[int] | None = None, + num_learnable_registers: int | None = 128, + rope_type: LTXRopeType = LTXRopeType.SPLIT, + double_precision_rope: bool = True, + apply_gated_attention: bool = True, + ff_bias: bool = False, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + self.positional_embedding_theta = positional_embedding_theta + self.positional_embedding_max_pos = positional_embedding_max_pos if positional_embedding_max_pos is not None else [1] + self.rope_type = rope_type + self.double_precision_rope = double_precision_rope + self.transformer_1d_blocks = torch.nn.ModuleList( + [ + LTX25BasicTransformerBlock1D( + dim=self.inner_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + rope_type=rope_type, + apply_gated_attention=apply_gated_attention, + ff_bias=ff_bias, + ) + for _ in range(num_layers) + ] + ) + self.num_learnable_registers = num_learnable_registers + if self.num_learnable_registers: + self.learnable_registers = torch.nn.Parameter( + torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0 + ) + + def _replace_padded_with_learnable_registers( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, _ = hidden_states.shape + assert sequence_length % self.num_learnable_registers == 0 + registers = self.learnable_registers.to(hidden_states).repeat(sequence_length // self.num_learnable_registers, 1) + registers = registers.unsqueeze(0).expand(batch_size, -1, -1) + binary_mask = (additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0).to(hidden_states.dtype) + hidden_states = binary_mask * hidden_states + (1 - binary_mask) * registers + return hidden_states, torch.zeros_like(additive_attention_mask) + + def forward( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.num_learnable_registers: + hidden_states, additive_attention_mask = self._replace_padded_with_learnable_registers( + hidden_states, + additive_attention_mask, + ) + indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device) + indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1) + freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch + freqs_cis = precompute_freqs_cis( + indices_grid=indices_grid, + dim=self.inner_dim, + out_dtype=hidden_states.dtype, + theta=self.positional_embedding_theta, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + rope_type=self.rope_type, + freq_grid_generator=freq_grid_generator, + ) + for block in self.transformer_1d_blocks: + hidden_states = block(hidden_states, additive_attention_mask=additive_attention_mask, pe=freqs_cis) + return rms_norm(hidden_states), additive_attention_mask + + +class LTX25BasicTransformerBlock1D(torch.nn.Module): + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + rope_type: LTXRopeType, + apply_gated_attention: bool, + ff_bias: bool, + ): + super().__init__() + self.attn1 = Attention( + query_dim=dim, + heads=heads, + dim_head=dim_head, + rope_type=rope_type, + apply_gated_attention=apply_gated_attention, + ) + self.ff = FeedForward(dim, dim_out=dim, bias=ff_bias) + + def forward( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor | None = None, + pe: torch.Tensor | None = None, + ) -> torch.Tensor: + norm_hidden_states = rms_norm(hidden_states).squeeze(1) + hidden_states = self.attn1(norm_hidden_states, mask=additive_attention_mask, pe=pe) + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + hidden_states = self.ff(rms_norm(hidden_states)) + hidden_states + return hidden_states.squeeze(1) if hidden_states.ndim == 4 else hidden_states + + +class LTX25EmbeddingsProcessorOutput(NamedTuple): + video_encoding: torch.Tensor + audio_encoding: torch.Tensor | None + attention_mask: torch.Tensor + + +def _convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return (attention_mask.to(torch.int64) - 1).to(dtype).reshape( + attention_mask.shape[0], 1, 1, attention_mask.shape[-1] + ) * torch.finfo(dtype).max + + +def _right_pad_order(additive_attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + binary = (additive_attention_mask[:, 0, 0, :] >= 0).to(torch.int32) + sort_indices = torch.argsort(binary, dim=-1, descending=True, stable=True) + reordered = torch.gather(binary, 1, sort_indices) + additive = (reordered.to(additive_attention_mask.dtype) - 1) * torch.finfo(additive_attention_mask.dtype).max + return sort_indices, additive[:, None, None, :] + + +class LTX25TextEncoderPostModules(torch.nn.Module): + def __init__( + self, + embedding_dim: int = 3840, + num_layers: int = 49, + video_attention_heads: int = 32, + video_attention_head_dim: int = 128, + audio_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + num_connector_layers: int = 8, + connector_max_positions: list[int] | None = None, + connector_ff_bias: bool = True, + ): + super().__init__() + self.feature_extractor = LTX25FeatureExtractorV2( + video_aggregate_embed=torch.nn.Linear( + embedding_dim * num_layers, + video_attention_heads * video_attention_head_dim, + bias=True, + ), + embedding_dim=embedding_dim, + audio_aggregate_embed=torch.nn.Linear( + embedding_dim * num_layers, + audio_attention_heads * audio_attention_head_dim, + bias=True, + ), + ) + connector_max_positions = [4096] if connector_max_positions is None else connector_max_positions + self.video_connector = LTX25Embeddings1DConnector( + attention_head_dim=video_attention_head_dim, + num_attention_heads=video_attention_heads, + num_layers=num_connector_layers, + positional_embedding_max_pos=connector_max_positions, + ff_bias=connector_ff_bias, + ) + self.audio_connector = LTX25Embeddings1DConnector( + attention_head_dim=audio_attention_head_dim, + num_attention_heads=audio_attention_heads, + num_layers=num_connector_layers, + positional_embedding_max_pos=connector_max_positions, + ff_bias=connector_ff_bias, + ) + + def create_embeddings( + self, + video_features: torch.Tensor, + audio_features: torch.Tensor | None, + additive_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if audio_features is None: + raise ValueError("LTX-2.5 requires audio features for the audio connector.") + sort_indices, connector_mask = _right_pad_order(additive_attention_mask) + video_features = torch.gather(video_features, 1, sort_indices.unsqueeze(-1).expand_as(video_features)) + video_encoded, video_mask = self.video_connector(video_features, connector_mask) + binary_mask = (video_mask < 0.000001).to(torch.int64).reshape(video_encoded.shape[0], video_encoded.shape[1], 1) + video_encoded = video_encoded * binary_mask + audio_features = torch.gather(audio_features, 1, sort_indices.unsqueeze(-1).expand_as(audio_features)) + audio_encoded, _ = self.audio_connector(audio_features, connector_mask) + return video_encoded, audio_encoded, binary_mask.squeeze(-1) + + def process_hidden_states( + self, + hidden_states: tuple[torch.Tensor, ...], + attention_mask: torch.Tensor, + padding_side: str = "left", + ) -> LTX25EmbeddingsProcessorOutput: + video_features, audio_features = self.feature_extractor(hidden_states, attention_mask, padding_side) + additive_attention_mask = _convert_to_additive_mask(attention_mask, video_features.dtype) + video_encoding, audio_encoding, binary_mask = self.create_embeddings( + video_features, + audio_features, + additive_attention_mask, + ) + return LTX25EmbeddingsProcessorOutput(video_encoding, audio_encoding, binary_mask) diff --git a/diffsynth/models/ltx25_tokenizer.py b/diffsynth/models/ltx25_tokenizer.py new file mode 100644 index 000000000..d6fbbb95a --- /dev/null +++ b/diffsynth/models/ltx25_tokenizer.py @@ -0,0 +1,51 @@ +import json +from pathlib import Path + +import numpy as np +from safetensors import safe_open +from tokenizers import Tokenizer +from transformers import PreTrainedTokenizerFast + + +class LTX25GemmaTokenizer: + def __init__(self, model_path: str | Path, max_length: int = 1024): + model_path = Path(model_path) + with safe_open(model_path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + if "tokenizer_json" not in handle.keys(): + raise ValueError(f"{model_path} does not contain packed tokenizer_json assets.") + tokenizer_bytes = handle.get_tensor("tokenizer_json").detach().cpu().numpy().astype(np.uint8).tobytes() + raw_config = metadata.get("tokenizer_config.json") + if raw_config is None and "hf_asset__tokenizer_config.json" in handle.keys(): + raw_config = handle.get_tensor("hf_asset__tokenizer_config.json").detach().cpu().numpy().astype(np.uint8).tobytes().decode() + config = json.loads(raw_config) if raw_config else {} + ignored = {"tokenizer_class", "auto_map", "model_max_length", "backend", "is_local", "local_files_only", "processor_class", "added_tokens_decoder"} + config = {key: value for key, value in config.items() if key not in ignored} + self.tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer.from_buffer(tokenizer_bytes), + model_max_length=max_length, + **config, + ) + self.tokenizer.model_max_length = max_length + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + self.max_length = max_length + + def tokenize_with_weights(self, text: str) -> dict[str, list[tuple[int, int]]]: + text = text.strip() + bos_id = self.tokenizer.bos_token_id + if bos_id is None: + raise ValueError("Packed Gemma tokenizer has no BOS token id.") + encoded = self.tokenizer(text, padding=False, truncation=True, max_length=self.max_length, return_tensors="pt") + input_ids = encoded.input_ids[0].tolist() + if not input_ids or input_ids[0] != bos_id: + input_ids = [bos_id, *input_ids][: self.max_length] + padded = self.tokenizer.pad( + {"input_ids": [input_ids]}, + padding="max_length", + max_length=self.max_length, + return_tensors="pt", + return_attention_mask=True, + ) + return {"gemma": list(zip(padded.input_ids[0].tolist(), padded.attention_mask[0].tolist(), strict=True))} diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py index 9df0ed3a7..8639113db 100644 --- a/diffsynth/models/ltx2_dit.py +++ b/diffsynth/models/ltx2_dit.py @@ -870,6 +870,7 @@ class TransformerConfig: context_dim: int apply_gated_attention: bool = False cross_attention_adaln: bool = False + ff_bias: bool = True class BasicAVTransformerBlock(torch.nn.Module): @@ -903,7 +904,7 @@ def __init__( norm_eps=norm_eps, apply_gated_attention=video.apply_gated_attention, ) - self.ff = FeedForward(video.dim, dim_out=video.dim) + self.ff = FeedForward(video.dim, dim_out=video.dim, bias=video.ff_bias) video_sst_size = adaln_embedding_coefficient(video.cross_attention_adaln) self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, video.dim)) @@ -926,7 +927,7 @@ def __init__( norm_eps=norm_eps, apply_gated_attention=audio.apply_gated_attention, ) - self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) + self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim, bias=audio.ff_bias) audio_sst_size = adaln_embedding_coefficient(audio.cross_attention_adaln) self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, audio.dim)) @@ -1243,22 +1244,21 @@ def apply_cross_attention_adaln( class GELUApprox(torch.nn.Module): - def __init__(self, dim_in: int, dim_out: int) -> None: + def __init__(self, dim_in: int, dim_out: int, bias: bool = True) -> None: super().__init__() - self.proj = torch.nn.Linear(dim_in, dim_out) + self.proj = torch.nn.Linear(dim_in, dim_out, bias=bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.gelu(self.proj(x), approximate="tanh") class FeedForward(torch.nn.Module): - def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: + def __init__(self, dim: int, dim_out: int, mult: int = 4, bias: bool = True) -> None: super().__init__() inner_dim = int(dim * mult) - project_in = GELUApprox(dim, inner_dim) - - self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out)) + project_in = GELUApprox(dim, inner_dim, bias=bias) + self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out, bias=bias)) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) @@ -1309,6 +1309,10 @@ def __init__( # noqa: PLR0913 double_precision_rope: bool = True, apply_gated_attention: bool = False, cross_attention_adaln: bool = False, + use_prompt_adaln_single: bool = True, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_keyframes_abs_pos_embedding: bool = False, ): super().__init__() self._enable_gradient_checkpointing = False @@ -1319,6 +1323,8 @@ def __init__( # noqa: PLR0913 self.positional_embedding_theta = positional_embedding_theta self.model_type = model_type self.cross_attention_adaln = cross_attention_adaln + self.use_prompt_adaln_single = use_prompt_adaln_single + self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding cross_pe_max_pos = None if model_type.is_video_enabled(): if positional_embedding_max_pos is None: @@ -1362,6 +1368,8 @@ def __init__( # noqa: PLR0913 audio_cross_attention_dim=audio_cross_attention_dim, norm_eps=norm_eps, apply_gated_attention=apply_gated_attention, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, ) @property @@ -1379,7 +1387,12 @@ def _init_video( # Video input components self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient) - self.prompt_adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None + self.prompt_adaln_single = AdaLayerNormSingle( + self.inner_dim, embedding_coefficient=2 + ) if self.cross_attention_adaln and self.use_prompt_adaln_single else None + self.keyframes_abs_pos_embedding = ( + torch.nn.Parameter(torch.zeros(1, self.inner_dim)) if self.use_keyframes_abs_pos_embedding else None + ) # Video caption projection if caption_channels is not None: @@ -1406,7 +1419,9 @@ def _init_audio( self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) self.audio_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=self._adaln_embedding_coefficient) - self.audio_prompt_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None + self.audio_prompt_adaln_single = AdaLayerNormSingle( + self.audio_inner_dim, embedding_coefficient=2 + ) if self.cross_attention_adaln and self.use_prompt_adaln_single else None # Audio caption projection if caption_channels is not None: @@ -1530,6 +1545,8 @@ def _init_transformer_blocks( audio_cross_attention_dim: int, norm_eps: float, apply_gated_attention: bool, + ff_bias: bool, + audio_ff_bias: bool, ) -> None: """Initialize transformer blocks for LTX.""" video_config = ( @@ -1540,6 +1557,7 @@ def _init_transformer_blocks( context_dim=cross_attention_dim, apply_gated_attention=apply_gated_attention, cross_attention_adaln=self.cross_attention_adaln, + ff_bias=ff_bias, ) if self.model_type.is_video_enabled() else None @@ -1552,6 +1570,7 @@ def _init_transformer_blocks( context_dim=audio_cross_attention_dim, apply_gated_attention=apply_gated_attention, cross_attention_adaln=self.cross_attention_adaln, + ff_bias=audio_ff_bias, ) if self.model_type.is_audio_enabled() else None diff --git a/diffsynth/models/ltx2_text_encoder.py b/diffsynth/models/ltx2_text_encoder.py index e4f3b1a3e..3570df568 100644 --- a/diffsynth/models/ltx2_text_encoder.py +++ b/diffsynth/models/ltx2_text_encoder.py @@ -225,6 +225,7 @@ def __init__( dim_head: int, rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, apply_gated_attention: bool = False, + ff_bias: bool = True, ): super().__init__() @@ -239,6 +240,7 @@ def __init__( self.ff = FeedForward( dim, dim_out=dim, + bias=ff_bias, ) def forward( @@ -307,6 +309,7 @@ def __init__( rope_type: LTXRopeType = LTXRopeType.SPLIT, double_precision_rope: bool = True, apply_gated_attention: bool = False, + ff_bias: bool = True, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -326,6 +329,7 @@ def __init__( dim_head=attention_head_dim, rope_type=rope_type, apply_gated_attention=apply_gated_attention, + ff_bias=ff_bias, ) for _ in range(num_layers) ] diff --git a/diffsynth/pipelines/ltx25_audio_video.py b/diffsynth/pipelines/ltx25_audio_video.py new file mode 100644 index 000000000..5d549a950 --- /dev/null +++ b/diffsynth/pipelines/ltx25_audio_video.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Union + +import torch +from tqdm import tqdm + +from ..core import ModelConfig +from ..core.device.npu_compatible_device import get_device_type +from ..models.ltx25_tokenizer import LTX25GemmaTokenizer +from .ltx2_audio_video import ( + LTX2AudioVideoPipeline, + LTX2AudioVideoUnit_PromptEmbedder, +) + + +def _seconds_to_num_frames(seconds: float, frame_rate: float, min_frames: int = 1, max_frames: int = 1024) -> int: + raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames)) + frames = ((raw_frames - 1) // 8) * 8 + 1 + if frames < min_frames: + frames = min(-(-(min_frames - 1) // 8) * 8 + 1, max_frames) + return frames + + +class LTX25AudioVideoUnit_PromptEmbedder(LTX2AudioVideoUnit_PromptEmbedder): + def _preprocess_text(self, pipe, text: str): + token_pairs = pipe.tokenizer.tokenize_with_weights(text)["gemma"] + input_ids = torch.tensor([[token_id for token_id, _ in token_pairs]], device=pipe.device) + attention_mask = torch.tensor([[weight for _, weight in token_pairs]], device=pipe.device) + + outputs = pipe.text_encoder.model.model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + return outputs.hidden_states, attention_mask + + +class LTX25AudioVideoPipeline(LTX2AudioVideoPipeline): + def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): + super().__init__(device=device, torch_dtype=torch_dtype) + self.duration_head = None + self.units[2] = LTX25AudioVideoUnit_PromptEmbedder() + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = get_device_type(), + model_configs: list[ModelConfig] = [], + gemma_path: str | Path | None = None, + vram_limit: float | None = None, + load_duration_head: bool = False, + stage2_lora_config: ModelConfig | None = None, + stage2_lora_strength: float = 1.0, + ) -> "LTX25AudioVideoPipeline": + if gemma_path is None: + raise ValueError("gemma_path is required for the packed LTX-2.5 Gemma4 tokenizer assets.") + pipe = LTX25AudioVideoPipeline(device=device, torch_dtype=torch_dtype) + model_pool = pipe.download_and_load_models(model_configs, vram_limit) + pipe.text_encoder = model_pool.fetch_model("ltx25_text_encoder") + pipe.text_encoder_post_modules = model_pool.fetch_model("ltx25_text_encoder_post_modules") + pipe.dit = model_pool.fetch_model("ltx25_dit") + pipe.video_vae_encoder = model_pool.fetch_model("ltx25_video_vae_encoder") + pipe.video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") + pipe.audio_vae_decoder = model_pool.fetch_model("ltx25_audio_vae_decoder") + pipe.audio_vocoder = model_pool.fetch_model("ltx25_audio_vocoder") + pipe.audio_vae_encoder = model_pool.fetch_model("ltx25_audio_vae_encoder") + + pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") + if load_duration_head: + pipe.duration_head = model_pool.fetch_model("ltx25_duration_head") + if stage2_lora_config is not None: + stage2_lora_config.download_if_necessary() + pipe.stage2_lora_config = stage2_lora_config + pipe.stage2_lora_strength = stage2_lora_strength + pipe.tokenizer = LTX25GemmaTokenizer(gemma_path) + pipe.vram_management_enabled = pipe.check_vram_management_state() + return pipe + + @torch.no_grad() + def predict_num_frames(self, prompt: str, frame_rate: float = 24.0) -> int: + if self.duration_head is None: + raise ValueError("Automatic duration requires from_pretrained(..., load_duration_head=True) and its ModelConfig.") + self.load_models_to_device(("text_encoder", "text_encoder_post_modules", "duration_head")) + embedder = self.units[2] + hidden_states, attention_mask = embedder._preprocess_text(self, prompt) + video_context, audio_context, _ = self.text_encoder_post_modules.process_hidden_states(hidden_states, attention_mask) + seconds = float(self.duration_head(video_context, audio_context).item()) + return _seconds_to_num_frames(seconds, frame_rate) + + @torch.no_grad() + def __call__( + self, + *args, + use_two_stage_pipeline: bool = True, + use_distilled_pipeline: bool = True, + cfg_scale: float = 1.0, + num_inference_steps: int = 8, + progress_bar_cmd=tqdm, + **kwargs, + ): + if use_distilled_pipeline and not use_two_stage_pipeline: + raise ValueError("LTX-2.5 distilled inference requires the two-stage refinement flow.") + if use_distilled_pipeline and cfg_scale != 1.0: + raise ValueError("LTX-2.5 distilled inference requires cfg_scale=1.0.") + if use_two_stage_pipeline and not use_distilled_pipeline and not hasattr(self, "stage2_lora_config"): + raise ValueError("LTX-2.5 Dev two-stage inference requires stage2_lora_config.") + if kwargs.get("num_frames") is None: + prompt = kwargs.get("prompt", args[0] if args else "") + kwargs["num_frames"] = self.predict_num_frames(prompt, kwargs.get("frame_rate", 24.0)) + return super().__call__( + *args, + use_two_stage_pipeline=use_two_stage_pipeline, + use_distilled_pipeline=use_distilled_pipeline, + cfg_scale=cfg_scale, + num_inference_steps=num_inference_steps, + progress_bar_cmd=progress_bar_cmd, + **kwargs, + ) diff --git a/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py b/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py new file mode 100644 index 000000000..385690982 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py @@ -0,0 +1,27 @@ +def LTX25DiffusionVideoDecoderStateDictConverter(state_dict): + converted = {} + for source_name in state_dict: + if source_name.startswith("decoder."): + name = source_name.removeprefix("decoder.") + elif source_name.startswith("per_channel_statistics."): + name = source_name + else: + continue + + if name == "type_emb" or name.startswith("coarse_") or name.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")): + continue + name = name.replace("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1.") + name = name.replace("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2.") + value = state_dict[source_name] + if name.endswith(".attn.qkv.weight") or name.endswith(".attn.qkv.bias"): + if value.shape[0] % 3 != 0: + raise ValueError(f"Fused QKV tensor has invalid leading dimension: {source_name} {tuple(value.shape)}") + leaf = "weight" if name.endswith(".weight") else "bias" + prefix = name[: -len(leaf)] + q, k, v = value.chunk(3, dim=0) + converted[f"{prefix}to_q.{leaf}"] = q + converted[f"{prefix}to_k.{leaf}"] = k + converted[f"{prefix}to_v.{leaf}"] = v + else: + converted[name] = value + return converted diff --git a/diffsynth/utils/state_dict_converters/ltx25_duration_head.py b/diffsynth/utils/state_dict_converters/ltx25_duration_head.py new file mode 100644 index 000000000..f2a52ea49 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/ltx25_duration_head.py @@ -0,0 +1,6 @@ +def LTX25DurationHeadStateDictConverter(state_dict): + return { + name.removeprefix("duration_head."): state_dict[name] + for name in state_dict + if name.startswith("duration_head.") + } diff --git a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py new file mode 100644 index 000000000..d86afe87f --- /dev/null +++ b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py @@ -0,0 +1,31 @@ +def LTX25TextEncoderStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("model."): + new_name = "model.model.language_model." + name.removeprefix("model.") + elif name.startswith("vision_model."): + new_name = "model.model.embed_vision." + name.removeprefix("vision_model.") + elif name.startswith("multi_modal_projector."): + new_name = "model.model.embed_vision.multimodal_embedder." + name.removeprefix("multi_modal_projector.") + elif name.startswith("audio_projector."): + new_name = "model.model.embed_audio." + name.removeprefix("audio_projector.") + else: + continue + state_dict_[new_name] = state_dict[name] + state_dict_["model.lm_head.weight"] = state_dict_["model.model.language_model.embed_tokens.weight"] + return state_dict_ + + +def LTX25TextEncoderPostModulesStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("text_embedding_projection."): + new_name = "feature_extractor." + name.removeprefix("text_embedding_projection.") + elif name.startswith("model.diffusion_model.video_embeddings_connector."): + new_name = "video_connector." + name.removeprefix("model.diffusion_model.video_embeddings_connector.") + elif name.startswith("model.diffusion_model.audio_embeddings_connector."): + new_name = "audio_connector." + name.removeprefix("model.diffusion_model.audio_embeddings_connector.") + else: + continue + state_dict_[new_name] = state_dict[name] + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/ltx2_video_vae.py b/diffsynth/utils/state_dict_converters/ltx2_video_vae.py index 53df15e54..492bc49ce 100644 --- a/diffsynth/utils/state_dict_converters/ltx2_video_vae.py +++ b/diffsynth/utils/state_dict_converters/ltx2_video_vae.py @@ -4,7 +4,9 @@ def LTX2VideoEncoderStateDictConverter(state_dict): if name.startswith("vae.encoder."): new_name = name.replace("vae.encoder.", "") state_dict_[new_name] = state_dict[name] - elif name.startswith("vae.per_channel_statistics."): + elif name.startswith("encoder."): + state_dict_[name.removeprefix("encoder.")] = state_dict[name] + elif name.startswith("vae.per_channel_statistics.") or name.startswith("per_channel_statistics."): new_name = name.replace("vae.per_channel_statistics.", "per_channel_statistics.") if new_name not in ["per_channel_statistics.channel", "per_channel_statistics.mean-of-stds", "per_channel_statistics.mean-of-stds_over_std-of-means"]: state_dict_[new_name] = state_dict[name] @@ -17,7 +19,9 @@ def LTX2VideoDecoderStateDictConverter(state_dict): if name.startswith("vae.decoder."): new_name = name.replace("vae.decoder.", "") state_dict_[new_name] = state_dict[name] - elif name.startswith("vae.per_channel_statistics."): + elif name.startswith("decoder."): + state_dict_[name.removeprefix("decoder.")] = state_dict[name] + elif name.startswith("vae.per_channel_statistics.") or name.startswith("per_channel_statistics."): new_name = name.replace("vae.per_channel_statistics.", "per_channel_statistics.") if new_name not in ["per_channel_statistics.channel", "per_channel_statistics.mean-of-stds", "per_channel_statistics.mean-of-stds_over_std-of-means"]: state_dict_[new_name] = state_dict[name] diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md new file mode 100644 index 000000000..cbc125a08 --- /dev/null +++ b/docs/en/Model_Details/LTX-2.5.md @@ -0,0 +1,88 @@ +# LTX-2.5 + +DiffSynth-Studio provides portable LTX-2.5 joint audio-video inference through +`LTX25AudioVideoPipeline`. It loads the official split BF16 checkpoints locally +and does not require `ltx_core`, NATTEN, Triton, or ltx-kernels at runtime. + +> LTX-2.5 weights are gated. Obtain access from Lightricks and place the +> checkpoint files in a local model directory before running an example. + +## Implemented components + +- LTX-2.5 22B Distilled and Dev DiT checkpoints +- Fine-tuned Gemma4 12B encoder, packed tokenizer assets, and dual AV connectors +- Duration head and automatic causal-grid frame-count prediction +- DiffVAE video encoder and pure-PyTorch eager diffusion decoder +- Audio VAE, 48 kHz BWE vocoder, spatial x2 latent upsampler, and temporal x2 + upsampler registration +- Dev stage-2 distilled-LoRA loading + +The eager DiffVAE decoder uses tiled scaled-dot-product attention as a portable +neighborhood-attention fallback. Its fixed-input output matches the upstream +eager implementation for deterministic decoder stages and an x0 diffusion step. + +## Inference modes + +The public `LTX25AudioVideoPipeline` API follows the existing LTX-2.3 API. + +| Mode | Pipeline parameters | Status | +|---|---|---| +| Distilled two-stage T2AV | `use_distilled_pipeline=True`, `use_two_stage_pipeline=True` | Supported | +| Distilled two-stage I2AV | `input_images`, `input_images_indexes` | Supported | +| Dev one-stage T2AV | `use_distilled_pipeline=False`, `use_two_stage_pipeline=False` | Supported | +| Dev one-stage I2AV | `input_images`, `use_two_stage_pipeline=False` | Supported | +| Dev two-stage T2AV/I2AV | `stage2_lora_config`, `use_two_stage_pipeline=True` | Supported | +| Audio-to-video | `retake_audio`, `audio_sample_rate`, optional `retake_audio_regions` | Supported | +| Video/audio retake | `retake_video`, `retake_video_regions`, `retake_audio_regions` | Supported | +| Keyframe interpolation | multiple `input_images` and `input_images_indexes` | Supported | +| Pixel Spatial Upscaler IC-LoRA | `in_context_videos`, `in_context_downsample_factor=2` | Supported | + +For a two-stage Dev run, supply the released distilled stage-2 LoRA through +`stage2_lora_config`. The two-stage path is required for distilled inference; +Dev also supports a one-stage path without a stage-2 LoRA. + +The Pixel Spatial Upscaler requires the official LTX-2.5 Pixel IC-LoRA. +Load it with `pipe.load_lora(pipe.dit, ModelConfig(path=...))`, pass the +reference video through `in_context_videos`, and set +`clear_lora_before_state_two=True`. Its reference resolution is one quarter of +the final height and width: stage 1 is half resolution and the adapter's +`reference_downscale_factor` is 2. Do not load an LTX-2.3 adapter into an +LTX-2.5 DiT. + +## Geometry and memory requirements + +- `num_frames % 8 == 1` +- One-stage height and width must be divisible by 32. +- Two-stage height and width must be divisible by 64. +- The low-VRAM examples use BF16 compute, FP8 CPU weight offload, and + fine-grained management for the DiT, Gemma4 encoder, text connectors, and + DiffVAE decoder. +- `LTX25_VRAM_LIMIT_GB` controls the GPU budget used to retain prepared model + layers. It defaults to 16; it is not a hard end-to-end VRAM limit. +- A 960×576×121 distilled T2AV run with `LTX25_VRAM_LIMIT_GB=16` measured a + 31.8 GiB PyTorch allocation peak and 44.4 GiB reserved peak. Treat 48 GiB + as a measured lower bound for this resolution, not a validated 48 GiB target, + and leave headroom for the driver. +- `tiled=True` does not yet split the portable DiffVAE decode volume. The + decoder determines the current full-resolution peak; lower-card support + requires a tiled decoder implementation. + +Run the low-VRAM examples with: + +```bash +python examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +python examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py +python examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +``` + +The split checkpoints are expected under `models/Lightricks/LTX-2.5`. The +Pixel example additionally expects the separately gated adapter under +`models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler`. Set a larger +prepared-layer budget with `LTX25_VRAM_LIMIT_GB=24` when GPU memory is +available; it can improve speed but does not lower the decoder peak. Adjust +paths in an example if the local model directory is different. + +## Scope + +This integration is inference-only. Training, LTX-2.5-specific DFR, Dub-It, +and HDR/EXR pipelines are not included in this scope. diff --git a/docs/en/index.rst b/docs/en/index.rst index 9f427801f..cb154d7f3 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -31,6 +31,7 @@ Welcome to DiffSynth-Studio's Documentation Model_Details/Z-Image Model_Details/Anima Model_Details/LTX-2 + Model_Details/LTX-2.5 Model_Details/ERNIE-Image Model_Details/JoyAI-Image Model_Details/ACE-Step diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md new file mode 100644 index 000000000..f697e7880 --- /dev/null +++ b/docs/zh/Model_Details/LTX-2.5.md @@ -0,0 +1,82 @@ +# LTX-2.5 + +DiffSynth-Studio 通过 `LTX25AudioVideoPipeline` 提供可移植的 LTX-2.5 +音视频联合推理。该实现从本地加载官方 BF16 分组件权重，运行时不依赖 +`ltx_core`、NATTEN、Triton 或 ltx-kernels。 + +> LTX-2.5 权重受门控保护。运行示例前，请先在 Lightricks 页面申请访问权限， +> 并将权重放到本地模型目录。 + +## 已实现的组件 + +- LTX-2.5 22B Distilled 和 Dev DiT +- LTX 微调 Gemma4 12B 编码器、内嵌 tokenizer 资产和音视频双 connector +- Duration Head 与因果帧网格的自动帧数预测 +- DiffVAE 视频编码器和纯 PyTorch eager diffusion 解码器 +- 音频 VAE、48 kHz BWE vocoder、空间 x2 latent upsampler 及时域 x2 + upsampler 注册 +- Dev 第二阶段 distilled-LoRA 加载 + +Eager DiffVAE 解码器使用 tiled scaled-dot-product attention 实现可移植的 +邻域注意力后备路径。固定输入下，其确定性解码阶段和一个 x0 diffusion step +与上游 eager 实现数值一致。 + +## 推理模式 + +公开的 `LTX25AudioVideoPipeline` API 与现有 LTX-2.3 API 对齐。 + +| 模式 | Pipeline 参数 | 状态 | +|---|---|---| +| Distilled 两阶段 T2AV | `use_distilled_pipeline=True`，`use_two_stage_pipeline=True` | 支持 | +| Distilled 两阶段 I2AV | `input_images`，`input_images_indexes` | 支持 | +| Dev 单阶段 T2AV | `use_distilled_pipeline=False`，`use_two_stage_pipeline=False` | 支持 | +| Dev 单阶段 I2AV | `input_images`，`use_two_stage_pipeline=False` | 支持 | +| Dev 两阶段 T2AV/I2AV | `stage2_lora_config`，`use_two_stage_pipeline=True` | 支持 | +| A2V | `retake_audio`，`audio_sample_rate`，可选 `retake_audio_regions` | 支持 | +| Video/audio Retake | `retake_video`，`retake_video_regions`，`retake_audio_regions` | 支持 | +| Keyframe interpolation | 多个 `input_images` 和 `input_images_indexes` | 支持 | +| Pixel Spatial Upscaler IC-LoRA | `in_context_videos`，`in_context_downsample_factor=2` | 支持 | + +Dev 两阶段推理需要通过 `stage2_lora_config` 提供发布的第二阶段 distilled-LoRA。 +Distilled 推理必须使用两阶段；Dev 同时支持不加载第二阶段 LoRA 的单阶段路径。 + +Pixel Spatial Upscaler 需要官方的 LTX-2.5 Pixel IC-LoRA。通过 +`pipe.load_lora(pipe.dit, ModelConfig(path=...))` 加载它，将参考视频传给 +`in_context_videos`，并设置 `clear_lora_before_state_two=True`。参考视频的 +宽高应为最终输出的四分之一：第一阶段为半分辨率，adapter 的 +`reference_downscale_factor` 为 2。不要将 LTX-2.3 adapter 加载到 +LTX-2.5 DiT 中。 + +## 几何与显存要求 + +- `num_frames % 8 == 1` +- 单阶段的高度、宽度必须是 32 的倍数。 +- 两阶段的高度、宽度必须是 64 的倍数。 +- 低显存示例使用 BF16 计算、FP8 CPU 权重卸载，并为 DiT、Gemma4 编码器、 + text connector 和 DiffVAE decoder 启用细粒度显存管理。 +- `LTX25_VRAM_LIMIT_GB` 控制保留在 GPU 上的预加载模型层预算，默认值为 16； + 它不是端到端显存硬上限。 +- 在 `LTX25_VRAM_LIMIT_GB=16` 下，960×576×121 distilled T2AV 实测 PyTorch + allocated 峰值为 31.8 GiB、reserved 峰值为 44.4 GiB。48 GiB 仅是该分辨率 + 的实测下界，不代表已在 48 GiB 显存卡上验证，并应为驱动预留余量。 +- 当前 `tiled=True` 不会切分 portable DiffVAE 的完整 decode volume。decoder + 决定了目前的全分辨率峰值；要支持更低显存卡，需要实现 tiled decoder。 + +运行低显存示例： + +```bash +python examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +python examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py +python examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +``` + +分组件权重默认位于 `models/Lightricks/LTX-2.5`。Pixel 示例还需要单独门控的 +adapter，默认位于 +`models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler`。显存充足时可以 +通过 `LTX25_VRAM_LIMIT_GB=24` 增大预加载层预算以改善速度，但它不会降低 decoder +峰值。如本地模型目录不同，请修改示例中的路径。 + +## 范围 + +该接入仅覆盖推理。训练、LTX-2.5 专有 DFR、Dub-It 和 HDR/EXR +pipeline 不在本次范围内。 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index fe6a8694c..ff9af07a0 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -31,6 +31,7 @@ Model_Details/Z-Image Model_Details/Anima Model_Details/LTX-2 + Model_Details/LTX-2.5 Model_Details/ERNIE-Image Model_Details/JoyAI-Image Model_Details/ACE-Step diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py new file mode 100644 index 000000000..3fde37ccb --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -0,0 +1,69 @@ +import os + +import torch + +from diffsynth.core import ModelConfig +from diffsynth.pipelines.ltx25_audio_video import LTX25AudioVideoPipeline +from diffsynth.utils.data import VideoData +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +MODEL_ROOT = "models/Lightricks/LTX-2.5" +PIXEL_LORA = "models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler/ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors" +INPUT_VIDEO = "input.mp4" +GEMMA = f"{MODEL_ROOT}/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" +TRANSFORMER = f"{MODEL_ROOT}/diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors" +VRAM_LIMIT_GB = float(os.environ.get("LTX25_VRAM_LIMIT_GB", "16")) + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e5m2, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} + +pipe = LTX25AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(path=GEMMA, **vram_config), + ModelConfig(path=[GEMMA, TRANSFORMER], **vram_config), + ModelConfig(path=TRANSFORMER, **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig( + path=f"{MODEL_ROOT}/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + **vram_config, + ), + ], + gemma_path=GEMMA, + vram_limit=VRAM_LIMIT_GB, +) +pipe.load_lora(pipe.dit, PIXEL_LORA) + +height, width, num_frames = 576, 960, 121 +reference_video = VideoData(INPUT_VIDEO, height=height // 4, width=width // 4).raw_data() +video, audio = pipe( + prompt="A colorful sailboat crosses a calm lake at sunrise. Gentle water sounds and distant birds.", + seed=42, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + in_context_videos=[reference_video], + in_context_downsample_factor=2, + tiled=True, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + clear_lora_before_state_two=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_pixel_spatial_upscale.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py new file mode 100644 index 000000000..4d306e77d --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py @@ -0,0 +1,70 @@ +import os + +import torch +from PIL import Image + +from diffsynth.core import ModelConfig +from diffsynth.pipelines.ltx25_audio_video import LTX25AudioVideoPipeline +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +MODEL_ROOT = "models/Lightricks/LTX-2.5" +START_IMAGE = "start.png" +MIDDLE_IMAGE = "middle.png" +END_IMAGE = "end.png" +GEMMA = f"{MODEL_ROOT}/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" +TRANSFORMER = f"{MODEL_ROOT}/diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors" +STAGE2_LORA = f"{MODEL_ROOT}/loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors" +VRAM_LIMIT_GB = float(os.environ.get("LTX25_VRAM_LIMIT_GB", "16")) + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e5m2, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} + +pipe = LTX25AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(path=GEMMA, **vram_config), + ModelConfig(path=[GEMMA, TRANSFORMER], **vram_config), + ModelConfig(path=TRANSFORMER, **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig( + path=f"{MODEL_ROOT}/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + **vram_config, + ), + ], + gemma_path=GEMMA, + stage2_lora_config=ModelConfig(path=STAGE2_LORA), + vram_limit=VRAM_LIMIT_GB, +) + +video, audio = pipe( + prompt="A colorful sailboat crosses a calm lake at sunrise. Gentle water sounds and distant birds.", + seed=42, + height=576, + width=960, + num_frames=121, + frame_rate=24, + input_images=[Image.open(path).convert("RGB") for path in (START_IMAGE, MIDDLE_IMAGE, END_IMAGE)], + input_images_indexes=[0, 60, 120], + tiled=True, + use_distilled_pipeline=False, + use_two_stage_pipeline=True, + cfg_scale=3.0, + num_inference_steps=8, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_keyframe_interpolation.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py new file mode 100644 index 000000000..c9e7e1157 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -0,0 +1,54 @@ +import os + +import torch + +from diffsynth.core import ModelConfig +from diffsynth.pipelines.ltx25_audio_video import LTX25AudioVideoPipeline +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +MODEL_ROOT = "models/Lightricks/LTX-2.5" +GEMMA = f"{MODEL_ROOT}/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" +TRANSFORMER = f"{MODEL_ROOT}/diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors" +VRAM_LIMIT_GB = float(os.environ.get("LTX25_VRAM_LIMIT_GB", "16")) + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e5m2, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} + +pipe = LTX25AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(path=GEMMA, **vram_config), + ModelConfig(path=[GEMMA, TRANSFORMER], **vram_config), + ModelConfig(path=TRANSFORMER, **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(path=f"{MODEL_ROOT}/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + gemma_path=GEMMA, + vram_limit=VRAM_LIMIT_GB, +) + +video, audio = pipe( + prompt="A gentle ocean wave rolls toward a quiet sunrise beach. Natural ambient surf audio.", + seed=42, + height=576, + width=960, + num_frames=121, + tiled=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) From 1ba11f21c912bb384579e9e60a1a8fb1f2c56490 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Mon, 7 Sep 2026 22:50:32 +0800 Subject: [PATCH 02/31] Support LTX-2.5: unified pipeline, flat DiffVAE, INT8 variants and inference examples - Merge LTX-2.5 into LTX2AudioVideoPipeline (drop separate ltx25_audio_video.py) - Register ltx25 DiT / Gemma4 text encoder / feature extractor + connectors / video VAE (DiffVAE + ConvVAE) / audio VAE / vocoder / duration head, BF16 and INT8 ConvRot variants - Flatten DiffVAE into diffsynth/models/ltx25_diffusion_video_vae.py with a single decode interface and auto tiling - Add keyframes abs pos embedding and tokenwise AV cross-attention scale/shift to ltx2_dit - Add stage1 distilled ancestral schedule unit, audio-only (T2A) and A2V freeze support, auto_duration in one pipe call - VRAM management maps for new modules; low-VRAM and standard inference examples for T2AV/I2AV/A2V/T2A/Retake/IC-LoRA/INT8 --- diffsynth/configs/model_configs.py | 77 +- .../configs/vram_management_module_maps.py | 42 +- diffsynth/diffusion/base_pipeline.py | 15 +- diffsynth/diffusion/flow_match.py | 64 +- diffsynth/models/ltx25_diffusion_video_vae.py | 6369 +++++++++++++++++ diffsynth/models/ltx25_diffvae/NOTICE.md | 9 - diffsynth/models/ltx25_diffvae/__init__.py | 0 .../models/ltx25_diffvae/model/__init__.py | 0 .../model/transformer/__init__.py | 0 .../model/transformer/timestep_embedding.py | 115 - .../ltx25_diffvae/model/video_vae/__init__.py | 0 .../video_vae/diffusion_video_decoder.py | 208 - .../ltx25_diffvae/model/video_vae/ops.py | 57 - .../model/video_vae/transformer/__init__.py | 3 - .../model/video_vae/transformer/attention.py | 59 - .../model/video_vae/transformer/blocks.py | 66 - .../transformer/combined/__init__.py | 0 .../video_vae/transformer/combined/attn.py | 114 - .../video_vae/transformer/combined/block.py | 28 - .../video_vae/transformer/combined/context.py | 15 - .../video_vae/transformer/combined/mlp.py | 14 - .../video_vae/transformer/det_attn_rope.py | 156 - .../transformer/fallback_na/__init__.py | 8 - .../transformer/fallback_na/eager.py | 154 - .../model/video_vae/transformer/layers.py | 66 - .../model/video_vae/transformer/qkv.py | 16 - .../model/video_vae/transformer/rope_math.py | 56 - .../model/video_vae/transformer/swiglu.py | 38 - diffsynth/models/ltx25_text_encoder.py | 120 +- diffsynth/models/ltx2_common.py | 3 + diffsynth/models/ltx2_dit.py | 137 +- diffsynth/models/ltx2_video_vae.py | 10 + diffsynth/pipelines/ltx25_audio_video.py | 120 - diffsynth/pipelines/ltx2_audio_video.py | 596 +- .../ltx25_diffusion_video_vae.py | 20 +- .../ltx25_text_encoder.py | 12 +- .../model_inference/LTX-2.5-A2V-TwoStage.py | 56 + .../LTX-2.5-I2AV-DistilledPipeline.py | 87 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 65 + examples/ltx2/model_inference/LTX-2.5-T2A.py | 38 + .../LTX-2.5-T2AV-DistilledPipeline.py | 55 + .../LTX-2.5-T2AV-INT8-ConvRot.py | 50 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 65 + .../LTX-2.5-A2V-TwoStage.py | 57 + .../LTX-2.5-I2AV-DistilledPipeline.py | 88 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 59 +- .../model_inference_low_vram/LTX-2.5-T2A.py | 39 + .../LTX-2.5-T2AV-DistilledPipeline.py | 54 +- .../LTX-2.5-T2AV-INT8-ConvRot.py | 51 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 66 + 50 files changed, 8055 insertions(+), 1542 deletions(-) create mode 100644 diffsynth/models/ltx25_diffusion_video_vae.py delete mode 100644 diffsynth/models/ltx25_diffvae/NOTICE.md delete mode 100644 diffsynth/models/ltx25_diffvae/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/transformer/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/ops.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py delete mode 100644 diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py delete mode 100644 diffsynth/pipelines/ltx25_audio_video.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-T2A.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py index dd6bf0974..02ee04f35 100644 --- a/diffsynth/configs/model_configs.py +++ b/diffsynth/configs/model_configs.py @@ -896,82 +896,127 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-*-transformer-bf16.safetensors") "model_hash": "7960c5dc4626650824e36f65a8e992e9", "model_name": "ltx25_dit", "model_class": "diffsynth.models.ltx2_dit.LTXModel", - "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True}, + "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True, "use_tokenwise_av_ca_scale_shift": True}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", }, { - "model_hash": "4bc194ac62f5648db68d419916a25688", + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-*-transformer-comfy-int8-convrot.safetensors") + "model_hash": "57343d320cac0bbba58a488b8ebe7187", + "model_name": "ltx25_dit", + "model_class": "diffsynth.models.ltx2_dit.LTXModel", + "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True, "use_tokenwise_av_ca_scale_shift": True}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["timestep_embedder.linear_1", "timestep_embedder.linear_2", "adaln_single.linear", "audio_adaln_single.linear", "prompt_adaln_single.linear", "audio_prompt_adaln_single.linear", "av_ca_a2v_gate_adaln_single.linear", "av_ca_audio_scale_shift_adaln_single.linear", "av_ca_v2a_gate_adaln_single.linear", "av_ca_video_scale_shift_adaln_single.linear", "patchify_proj", "audio_patchify_proj", "proj_out", "audio_proj_out", "to_gate_logits"]}, + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") + "model_hash": "055700dc619165899bebb5162f699cd2", + "model_name": "ltx25_text_encoder", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") + "model_hash": "4743ded7a5725b6589bccdb62512723b", "model_name": "ltx25_text_encoder", "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["lm_head", "embedding_projection", "patch_dense"]}, }, { - "model_hash": "f1c63402b49c39c739f13cdb90714f9e", - "model_name": "ltx25_text_encoder_post_modules", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderPostModulesStateDictConverter", + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") + "model_hash": "055700dc619165899bebb5162f699cd2", + "model_name": "ltx25_feature_extractor", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25FeatureExtractorV2", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25FeatureExtractorStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors") + "model_hash": "7960c5dc4626650824e36f65a8e992e9", + "model_name": "ltx25_embeddings_connectors", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") + "model_hash": "4743ded7a5725b6589bccdb62512723b", + "model_name": "ltx25_feature_extractor", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25FeatureExtractorV2", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25FeatureExtractorStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_aggregate_embed", "audio_aggregate_embed"]}, + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors") + "model_hash": "57343d320cac0bbba58a488b8ebe7187", + "model_name": "ltx25_embeddings_connectors", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["to_gate_logits"]}, + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors") "model_hash": "e19205490f01801d0a7b6d3aba61e26e", "model_name": "ltx25_video_vae_encoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", - "extra_kwargs": {"encoder_version": "ltx-2.3"}, + "extra_kwargs": {"encoder_version": "ltx-2.3", "latent_log_var": "constant"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors") "model_hash": "e19205490f01801d0a7b6d3aba61e26e", "model_name": "ltx25_diffusion_video_vae_decoder", - "model_class": "diffsynth.models.ltx25_diffvae.model.video_vae.diffusion_video_decoder.DiffusionVideoDecoder", + "model_class": "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder", "extra_kwargs": {"stage_channels": [2048, 1024, 512, 512, 256], "stage_depths": [4, 6, 4, 2, 8], "stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]], "stage5_kernel": [11, 11, 11], "timestep_scale_multiplier": 1000.0, "default_num_inference_steps": 1, "model_output_type": "x0"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors") "model_hash": "a1d642eecae96baa9c31d4e405564f49", "model_name": "ltx25_conv_video_vae_encoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", - "extra_kwargs": {"encoder_version": "ltx-2.3"}, + "extra_kwargs": {"encoder_version": "ltx-2.3", "latent_log_var": "uniform"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors") "model_hash": "a1d642eecae96baa9c31d4e405564f49", "model_name": "ltx25_conv_video_vae_decoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder", - "extra_kwargs": {"decoder_version": "ltx-2.3"}, + "extra_kwargs": {"decoder_version": "ltx-2.3", "decoder_spatial_padding_mode": "zeros"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") "model_hash": "c2488315f13356abb806f9f217f1e803", "model_name": "ltx25_audio_vae_decoder", "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") "model_hash": "c2488315f13356abb806f9f217f1e803", "model_name": "ltx25_audio_vocoder", "model_class": "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") "model_hash": "c2488315f13356abb806f9f217f1e803", "model_name": "ltx25_audio_vae_encoder", "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter", }, { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors") "model_hash": "35840495e440a4f00946450269299bd6", "model_name": "ltx25_duration_head", "model_class": "diffsynth.models.ltx25_duration_head.LTX25DurationHead", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_duration_head.LTX25DurationHeadStateDictConverter", }, - { - "model_hash": "5fbb28ecc6becd9513cee69b2dfb1053", - "model_name": "ltx25_temporal_upsampler", - "model_class": "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler", - "extra_kwargs": {"mid_channels": 512, "spatial_upsample": False, "temporal_upsample": True, "spatial_scale": 1.0, "rational_resampler": True}, - }, ] anima_series = [ { diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index 2fdd7afca..b05a306d6 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -271,6 +271,7 @@ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", }, "diffsynth.models.ltx2_dit.LTXModel": { + "diffsynth.models.ltx2_dit.BasicAVTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, @@ -282,15 +283,23 @@ "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules": { + "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors": { + "diffsynth.models.ltx25_text_encoder.LTX25Embeddings1DConnector": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_text_encoder.LTX25BasicTransformerBlock1D": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.ltx25_diffvae.model.video_vae.diffusion_video_decoder.DiffusionVideoDecoder": { + "diffsynth.models.ltx25_text_encoder.LTX25FeatureExtractorV2": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + }, + "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder": { + "diffsynth.models.ltx25_diffusion_video_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffusion_video_vae.NABlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffusion_video_vae.CombinedDiffusionNABlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffusion_video_vae.SwiGLU": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffusion_video_vae.ChannelLinear": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_diffvae.model.video_vae.transformer.swiglu.SwiGLU": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.block.CombinedDiffusionNABlock": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": { "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", @@ -298,15 +307,40 @@ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": { + "diffsynth.models.ltx2_video_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": { + "diffsynth.models.ltx2_video_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", }, + "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder": { + "diffsynth.models.ltx2_audio_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + }, "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": { + "diffsynth.models.ltx2_audio_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_audio_vae.LTX2Vocoder": { + "diffsynth.models.ltx2_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.SnakeBeta": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.Activation1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.MelSTFT": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.DownSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE": { + "diffsynth.models.ltx2_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.SnakeBeta": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.Activation1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.MelSTFT": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.DownSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule", }, diff --git a/diffsynth/diffusion/base_pipeline.py b/diffsynth/diffusion/base_pipeline.py index 5af48a110..1a2ee13e5 100644 --- a/diffsynth/diffusion/base_pipeline.py +++ b/diffsynth/diffusion/base_pipeline.py @@ -180,9 +180,10 @@ def load_models_to_device(self, model_names): module.onload() - def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None): + def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None, generator=None): # Initialize Gaussian noise - generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) + if generator is None: + generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype) noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) return noise @@ -223,7 +224,12 @@ def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, if inpaint_mask is not None: noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents) noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask) - latents_next = scheduler.step(noise_pred, timestep, latents) + scheduler_kwargs = {} + if "ancestral_noise_shape" in kwargs: + scheduler_kwargs["ancestral_noise_shape"] = kwargs["ancestral_noise_shape"] + if "ancestral_noise_transform" in kwargs: + scheduler_kwargs["ancestral_noise_transform"] = kwargs["ancestral_noise_transform"] + latents_next = scheduler.step(noise_pred, timestep, latents, **scheduler_kwargs) return latents_next @@ -346,8 +352,9 @@ def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, i if isinstance(noise_pred_posi, tuple): # Separately handling different output types of latents, eg. video and audio latents. + # Disabled modalities return None and stay None under CFG. noise_pred = tuple( - n_nega + cfg_scale * (n_posi - n_nega) + None if n_posi is None or n_nega is None else n_nega + cfg_scale * (n_posi - n_nega) for n_posi, n_nega in zip(noise_pred_posi, noise_pred_nega) ) else: diff --git a/diffsynth/diffusion/flow_match.py b/diffsynth/diffusion/flow_match.py index 83d3497d4..ea0db6311 100644 --- a/diffsynth/diffusion/flow_match.py +++ b/diffsynth/diffusion/flow_match.py @@ -26,6 +26,32 @@ def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "Z "SenseNova-U1": FlowMatchScheduler.set_timesteps_sensenova_u1, }.get(template, FlowMatchScheduler.set_timesteps_flux) self.num_train_timesteps = 1000 + self.step_mode = "euler" + self.ancestral_eta = 1.0 + self.ancestral_s_noise = 1.0 + self.ancestral_generator = None + self.roundtrip_denoised = False + + def set_step_mode( + self, + mode="euler", + eta=1.0, + s_noise=1.0, + noise_seed=None, + device="cpu", + roundtrip_denoised=False, + ): + if mode not in ("euler", "euler_ancestral"): + raise ValueError(f"Unsupported flow-matching step mode: {mode}") + self.step_mode = mode + self.ancestral_eta = eta + self.ancestral_s_noise = s_noise + self.ancestral_generator = None + self.roundtrip_denoised = roundtrip_denoised + if mode == "euler_ancestral": + if noise_seed is None: + raise ValueError("noise_seed is required for ancestral Euler sampling.") + self.ancestral_generator = torch.Generator(device=device).manual_seed(noise_seed) @staticmethod def set_timesteps_flux(num_inference_steps=100, denoising_strength=1.0, shift=None): @@ -368,6 +394,7 @@ def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, trainin denoising_strength=denoising_strength, **kwargs, ) + self.set_step_mode("euler") if training: self.set_training_weight() self.training = True @@ -380,11 +407,42 @@ def step(self, model_output, timestep, sample, to_final=False, **kwargs): timestep_id = torch.argmin((self.timesteps - timestep).abs()) sigma = self.sigmas[timestep_id] if to_final or timestep_id + 1 >= len(self.timesteps): - sigma_ = 0 + sigma_ = torch.zeros_like(sigma) else: sigma_ = self.sigmas[timestep_id + 1] - prev_sample = sample + model_output * (sigma_ - sigma) - return prev_sample + denoised = sample.float() - model_output.float() * sigma.float() + if self.roundtrip_denoised: + denoised = denoised.to(sample.dtype).float() + if self.step_mode == "euler": + if not self.roundtrip_denoised: + return sample + model_output * (sigma_ - sigma) + velocity = ((sample.float() - denoised) / sigma.float()).to(sample.dtype) + return (sample.float() + velocity.float() * (sigma_ - sigma).float()).to(sample.dtype) + + if sigma_ == 0: + return denoised.to(sample.dtype) + downstep_ratio = 1.0 + (sigma_ / sigma - 1.0) * self.ancestral_eta + sigma_down = sigma_ * downstep_ratio + sigma_down_ratio = sigma_down / sigma + prev_sample = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised + alpha_next = 1.0 - sigma_ + alpha_down = 1.0 - sigma_down + renoise_coeff = ( + sigma_ ** 2 - sigma_down ** 2 * alpha_next ** 2 / alpha_down ** 2 + ).clamp(min=0).sqrt() + noise_shape = kwargs.get("ancestral_noise_shape", sample.shape) + noise = torch.randn( + noise_shape, + generator=self.ancestral_generator, + dtype=sample.dtype, + device=sample.device, + ) + noise_transform = kwargs.get("ancestral_noise_transform") + if noise_transform is not None: + noise = noise_transform(noise) + prev_sample = alpha_next / alpha_down * prev_sample + prev_sample = prev_sample + noise.float() * self.ancestral_s_noise * renoise_coeff + return prev_sample.to(sample.dtype) def return_to_timestep(self, timestep, sample, sample_stablized): if isinstance(timestep, torch.Tensor): diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py new file mode 100644 index 000000000..245608a90 --- /dev/null +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -0,0 +1,6369 @@ +"""LTX-2.5 diffusion video decoder with eager attention and internal tiling. + +The implementation is intentionally self-contained so DiffSynth does not depend on +``ltx-core`` or a nested source package at runtime. It preserves checkpoint names +and provides full, tiled, and keyframe-aware decode through one public interface. +""" + +from __future__ import annotations + +import dataclasses +import itertools +import logging +import math +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, replace +from enum import Enum +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Final, List, Literal, NamedTuple, Protocol, Tuple, TYPE_CHECKING + +import numpy as np +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn +from torch.nn import functional + + +class Disposable: + """Compatibility marker used by the target implementation.""" + + +class VideoDecoder: + """Structural marker for video decoders.""" + + +def _clip_generators(count, generator): + if isinstance(generator, Sequence): + if len(generator) != count: + raise ValueError(f"decode_single_frames got {count} latents and {len(generator)} generators") + return generator + return [generator] * count + + +def iter_decoded_single_frames(decoder, latents, generator=None): + generators = _clip_generators(len(latents), generator) + for index, (latent, item_generator) in enumerate(zip(latents, generators, strict=True)): + if latent.ndim != 5 or latent.shape[2] != 1: + raise ValueError( + f"decode_single_frames expects (B, C, 1, H, W) latents, got {tuple(latent.shape)} at index {index}" + ) + chunks = list(decoder.decode_video(latent, tiling_config=None, generator=item_generator)) + if not chunks: + raise RuntimeError(f"Decoder returned no pixels for single-frame latent {index}") + yield torch.cat(chunks, dim=0) + + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + Args + timesteps (torch.Tensor): + a 1-D Tensor of N indices, one per batch element. These may be fractional. + embedding_dim (int): + the dimension of the output. + flip_sin_to_cos (bool): + Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) + downscale_freq_shift (float): + Controls the delta between frequencies between dimensions + scale (float): + Scaling factor applied to the embeddings. + max_period (int): + Controls the maximum frequency of the embeddings + Returns + torch.Tensor: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class TimestepEmbedding(torch.nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + out_dim: int | None = None, + post_act_fn: str | None = None, + cond_proj_dim: int | None = None, + sample_proj_bias: bool = True, + ): + super().__init__() + + self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + + if cond_proj_dim is not None: + self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = torch.nn.SiLU() + time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim + + self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + + if post_act_fn is None: + self.post_act = None + + def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Timesteps(torch.nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + return t_emb + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): + """ + For PixArt-Alpha. + Reference: + https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 + """ + + def __init__( + self, + embedding_dim: int, + size_emb_dim: int, + ): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward( + self, + timestep: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + return timesteps_emb + + +class VideoPixelShape(NamedTuple): + """ + Shape of the tensor representing the video pixel array. Assumes BGR channel format. + """ + + batch: int + frames: int + height: int + width: int + fps: float + + +class SpatioTemporalScaleFactors(NamedTuple): + """ + Describes the spatiotemporal downscaling between decoded video space and + the corresponding VAE latent grid. + Field order matches the (frame/time, height, width) axis layout used by + latent tensors and meshgrid coordinates elsewhere in the codebase. + """ + + time: int + height: int + width: int + + @classmethod + def default(cls) -> "SpatioTemporalScaleFactors": + return cls(time=8, height=32, width=32) + + @classmethod + def from_blocks(cls, blocks: list, patch_size: int) -> "SpatioTemporalScaleFactors": + """Derive the scale factors from a VAE encoder/decoder block list. + Each ``compress_*`` block halves (encoder) or doubles (decoder) its target + axes by a stride of 2, independent of any channel ``multiplier``. The initial + patchify contributes an extra ``patch_size`` of spatial compression. Deriving + the factors from the blocks keeps a single source of truth that stays correct + across VAE variants (e.g. the 32x32x8 default and the 16x16x4 variant) instead + of relying on a hardcoded constant. + """ + spatial_steps = 0 + temporal_steps = 0 + for block_name, _ in blocks: + if block_name.startswith(("compress_space", "compress_all")): + spatial_steps += 1 + if block_name.startswith(("compress_time", "compress_all")): + temporal_steps += 1 + spatial = patch_size * (2**spatial_steps) + return cls(time=2**temporal_steps, height=spatial, width=spatial) + + @classmethod + def from_model_config(cls, model_config: dict) -> "SpatioTemporalScaleFactors": + """Derive the video scale factors from a checkpoint's model config dict. + Reads the embedded VAE block list (see ``from_blocks``). Falls back to the + default when the config carries no VAE block list -- either no ``vae`` section + or a ``vae`` section without encoder/decoder blocks (e.g. audio-only + checkpoints), where video tools are never built. + """ + vae_config = model_config.get("vae", {}) + blocks = vae_config.get("encoder_blocks") or vae_config.get("decoder_blocks") + if not blocks: + return cls.default() + return cls.from_blocks(blocks, vae_config.get("patch_size", 4)) + + +VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default() + + +class VideoLatentShape(NamedTuple): + """ + Shape of the tensor representing video in VAE latent space. + The latent representation is a 5D tensor with dimensions ordered as + (batch, channels, frames, height, width). Spatial and temporal dimensions + are downscaled relative to pixel space according to the VAE's scale factors. + """ + + batch: int + channels: int + frames: int + height: int + width: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.height, self.width]) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": + return VideoLatentShape( + batch=shape[0], + channels=shape[1], + frames=shape[2], + height=shape[3], + width=shape[4], + ) + + def token_count(self) -> int: + """Number of tokens after patchification with the default patch size of 1.""" + return self.frames * self.height * self.width + + def mask_shape(self) -> "VideoLatentShape": + return self._replace(channels=1) + + @staticmethod + def from_pixel_shape( + shape: VideoPixelShape, + latent_channels: int = 128, + scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS, + ) -> "VideoLatentShape": + frames = (shape.frames - 1) // scale_factors.time + 1 + height = shape.height // scale_factors.height + width = shape.width // scale_factors.width + + return VideoLatentShape( + batch=shape.batch, + channels=latent_channels, + frames=frames, + height=height, + width=width, + ) + + def upscale(self, scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS) -> "VideoLatentShape": + return self._replace( + channels=3, + frames=(self.frames - 1) * scale_factors.time + 1, + height=self.height * scale_factors.height, + width=self.width * scale_factors.width, + ) + + +class AudioLatentShape(NamedTuple): + """ + Shape of audio in VAE latent space: (batch, channels, frames, mel_bins). + mel_bins is the number of frequency bins from the mel-spectrogram encoding. + """ + + batch: int + channels: int + frames: int + mel_bins: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.mel_bins]) + + def token_count(self) -> int: + """Number of tokens after patchification.""" + return self.frames + + def mask_shape(self) -> "AudioLatentShape": + return self._replace(channels=1, mel_bins=1) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "AudioLatentShape": + return AudioLatentShape( + batch=shape[0], + channels=shape[1], + frames=shape[2], + mel_bins=shape[3], + ) + + @staticmethod + def from_duration( + batch: int, + duration: float, + channels: int = 8, + mel_bins: int = 16, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + ) -> "AudioLatentShape": + latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor) + + return AudioLatentShape( + batch=batch, + channels=channels, + frames=round(duration * latents_per_second), + mel_bins=mel_bins, + ) + + @staticmethod + def from_video_pixel_shape( + shape: VideoPixelShape, + channels: int = 8, + mel_bins: int = 16, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + ) -> "AudioLatentShape": + return AudioLatentShape.from_duration( + batch=shape.batch, + duration=float(shape.frames) / float(shape.fps), + channels=channels, + mel_bins=mel_bins, + sample_rate=sample_rate, + hop_length=hop_length, + audio_latent_downsample_factor=audio_latent_downsample_factor, + ) + + +@dataclass(frozen=True) +class Audio: + """ + Container for decoded audio samples and metadata. + Attributes: + waveform: Audio waveform tensor. + sampling_rate: Sampling rate (Hz) of the waveform. + """ + + waveform: torch.Tensor + sampling_rate: int + + def to(self, **kwargs: object) -> "Audio": + return replace(self, waveform=self.waveform.to(**kwargs)) + + +@dataclass(frozen=True) +class GeneratedKeyframeLayout: + """Where a state's generated-keyframe slot tokens live, and what they represent. + Recorded by :class:`~ltx_core.conditioning.types.keyframe_slots.VideoGeneratedKeyframeSlots` + when it appends the slots, so they can later be located and extracted *exactly* rather + than by assuming they are the trailing tokens. Conditioning items are applied in list + order and each appends to the end, so a state built with slots plus any other appending + conditioning item has no fixed trailing layout. + Attributes: + pixel_frame_indices: Target pixel-frame index of each slot, in token order. + tokens_per_keyframe: Number of tokens one slot occupies (one latent frame's worth). + first_token: Index of the first slot token in the token sequence. + """ + + pixel_frame_indices: tuple[int, ...] + tokens_per_keyframe: int + first_token: int + + @property + def num_keyframes(self) -> int: + return len(self.pixel_frame_indices) + + @property + def num_tokens(self) -> int: + return self.num_keyframes * self.tokens_per_keyframe + + @property + def token_slice(self) -> slice: + return slice(self.first_token, self.first_token + self.num_tokens) + + +@dataclass(frozen=True) +class LatentState: + """ + State of latents during the diffusion denoising process. + Attributes: + latent: The current noisy latent tensor being denoised. + denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising). + positions: Positional indices for each latent element, used for positional embeddings. + clean_latent: Initial state of the latent before denoising, may include conditioning latents. + attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention, + 0 = no attention. None means full attention everywhere. Built incrementally by conditioning items. + keyframes_mask: Optional per-token marker of shape (B, T, 1) -- same layout as + ``denoise_mask`` -- non-zero on tokens whose latent encodes a *single standalone pixel + frame* rather than the usual multi-frame span. That set is the target's first latent + frame (the video encoder is causal, so its first temporal latent frame covers 1 pixel + frame while the rest cover 8) plus any generated keyframe slots. Selects the tokens + that receive the model's learned keyframe absolute-position embedding; ignored + entirely by models built without ``use_keyframes_abs_pos_embedding``. + generated_keyframe_layout: Set when generated keyframe slots were appended; locates them. + generated_keyframes: Populated by ``clear_conditioning`` when a layout is present: the + denoised slot content as an unpatchified ``(B, C, K, H, W)`` latent, one latent frame + per keyframe. Each frame must be decoded as a standalone one-frame clip, never as a + K-frame video -- a causal decode would blend slots that were never adjacent. + frozen: When True, this stream is held fixed: token denoising is disabled (``denoise_mask`` + should be all zeros; pipeline builders enforce that) and the scalar noise level used for + prompt / cross-modality AdaLN gates is forced to 0 when the state is converted for the + transformer. + """ + + latent: torch.Tensor + denoise_mask: torch.Tensor + positions: torch.Tensor + clean_latent: torch.Tensor + attention_mask: torch.Tensor | None = None + keyframes_mask: torch.Tensor | None = None + generated_keyframe_layout: GeneratedKeyframeLayout | None = None + generated_keyframes: torch.Tensor | None = None + frozen: bool = False + + def clone(self) -> "LatentState": + return LatentState( + latent=self.latent.clone(), + denoise_mask=self.denoise_mask.clone(), + positions=self.positions.clone(), + clean_latent=self.clean_latent.clone(), + attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None, + keyframes_mask=self.keyframes_mask.clone() if self.keyframes_mask is not None else None, + generated_keyframe_layout=self.generated_keyframe_layout, + generated_keyframes=(self.generated_keyframes.clone() if self.generated_keyframes is not None else None), + frozen=self.frozen, + ) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor: + """Root-mean-square (RMS) normalize `x` over its last dimension. + Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized + shape and forwards `weight` and `eps`. + """ + return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) + + +def check_config_value(config: dict, key: str, expected: Any) -> None: # noqa: ANN401 + actual = config.get(key) + if actual != expected: + raise ValueError(f"Config value {key} is {actual}, expected {expected}") + + +def to_velocity( + sample: torch.Tensor, + sigma: float | torch.Tensor, + denoised_sample: torch.Tensor, + calc_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Convert the sample and its denoised version to velocity. + Returns: + Velocity + """ + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(calc_dtype).item() + if sigma == 0: + raise ValueError("Sigma can't be 0.0") + return ((sample.to(calc_dtype) - denoised_sample.to(calc_dtype)) / sigma).to(sample.dtype) + + +def to_denoised( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: float | torch.Tensor, + calc_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Convert the sample and its denoising velocity to denoised sample. + Returns: + Denoised sample + """ + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(calc_dtype) + return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) + + +def find_matching_file(root_path: str, pattern: str) -> Path: + """ + Recursively search for files matching a glob pattern and return the first match. + """ + matches = list(Path(root_path).rglob(pattern)) + if not matches: + raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}") + return matches[0] + + +def compute_trapezoidal_mask_1d( + length: int, + ramp_left: int, + ramp_right: int, + left_starts_from_0: bool = False, +) -> torch.Tensor: + """ + Generate a 1D trapezoidal blending mask with linear ramps. + Args: + length: Output length of the mask. + ramp_left: Fade-in length on the left. + ramp_right: Fade-out length on the right. + left_starts_from_0: Whether the ramp starts from 0 or first non-zero value. + Useful for temporal tiles where the first tile is causal. + Returns: + A 1D tensor of shape `(length,)` with values in [0, 1]. + """ + if length <= 0: + raise ValueError("Mask length must be positive.") + + ramp_left = max(0, min(ramp_left, length)) + ramp_right = max(0, min(ramp_right, length)) + + mask = torch.ones(length) + + if ramp_left > 0: + interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2 + fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1] + if not left_starts_from_0: + fade_in = fade_in[1:] + mask[:ramp_left] *= fade_in + + if ramp_right > 0: + fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1] + mask[-ramp_right:] *= fade_out + + return mask.clamp_(0, 1) + + +def compute_rectangular_mask_1d( + length: int, + left_ramp: int, + right_ramp: int, +) -> torch.Tensor: + """ + Generate a 1D rectangular (pulse) mask. + Args: + length: Output length of the mask. + left_ramp: Number of elements at the start of the mask to set to 0. + right_ramp: Number of elements at the end of the mask to set to 0. + Returns: + A 1D tensor of shape `(length,)` with values 0 or 1. + """ + if length <= 0: + raise ValueError("Mask length must be positive.") + + mask = torch.ones(length) + if left_ramp > 0: + mask[:left_ramp] = 0 + if right_ramp > 0: + mask[-right_ramp:] = 0 + return mask + + +@dataclass(frozen=True) +class DimensionInterval: + start: int + end: int + left_ramp: int + right_ramp: int + + +@dataclass(frozen=True) +class DimensionIntervals: + """Intervals which a single dimension of the latent space is split into. + Each interval is defined by its start, end, left ramp, and right ramp. + The start and end are the indices of the first and last element (exclusive) in the interval. + Ramps are regions of the interval where the value of the mask tensor is + interpolated between 0 and 1 for blending with neighboring intervals. + The left ramp and right ramp values are the lengths of the left and right ramps. + """ + + intervals: list[DimensionInterval] + + +@dataclass(frozen=True) +class LatentIntervals: + """Intervals which the latent tensor of given shape is split into. + Each dimension of the latent space is split into intervals based on the length along said dimension. + """ + + original_shape: torch.Size + dimension_intervals: tuple[DimensionIntervals, ...] + + +SplitOperation = Callable[[int], DimensionIntervals] + + +MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor]]] + + +def default_split_operation(length: int) -> DimensionIntervals: + return DimensionIntervals(intervals=[DimensionInterval(start=0, end=length, left_ramp=0, right_ramp=0)]) + + +DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation + + +def untiled_mask_1d() -> torch.Tensor: + """Length-1 ones that broadcast over an untiled axis (historical ``None`` mask).""" + return torch.ones(1) + + +def default_mapping_operation( + _intervals: DimensionIntervals, +) -> tuple[list[slice], list[torch.Tensor]]: + return [slice(0, None)], [untiled_mask_1d()] + + +DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation + + +def _grow_last_tile_to_min(intervals: list[DimensionInterval], min_tile_size: int) -> list[DimensionInterval]: + """Grow a short last tile left to ``min_tile_size``; widen penultimate ``right_ramp``.""" + if len(intervals) <= 1: + return list(intervals) + last = intervals[-1] + if last.end - last.start >= min_tile_size: + return list(intervals) + new_start = last.end - min_tile_size + prev = intervals[-2] + new_overlap = prev.end - new_start + return [ + *intervals[:-2], + replace(prev, right_ramp=new_overlap), + replace(last, start=new_start, left_ramp=new_overlap), + ] + + +def _validate_tile_intervals(intervals: list[DimensionInterval], *, dim_size: int, min_tile_size: int) -> None: + """Validate coverage, ramp/overlap consistency, and ``min_tile_size``.""" + if not intervals or intervals[0].start != 0 or intervals[-1].end != dim_size: + raise ValueError(f"tiles must cover [0, {dim_size})") + for i, iv in enumerate(intervals): + length = iv.end - iv.start + if length < min_tile_size: + raise ValueError(f"tile {i} length {length} is below min_tile_size={min_tile_size}") + if iv.left_ramp < 0 or iv.right_ramp < 0 or iv.left_ramp > length or iv.right_ramp > length: + raise ValueError(f"tile {i} has invalid ramps: left={iv.left_ramp}, right={iv.right_ramp}, length={length}") + if i == 0: + continue + overlap = intervals[i - 1].end - iv.start + if overlap < 0 or intervals[i - 1].right_ramp != overlap or iv.left_ramp != overlap: + raise ValueError(f"tiles {i - 1}/{i}: ramp/overlap mismatch (overlap={overlap})") + + +def split_by_size(size: int, overlap: int, min_tile_size: int | None = None) -> SplitOperation: + """Split a dimension into overlapping tiles of a given size. + Tiles are sized ``size`` with ``overlap`` shared elements between + consecutive tiles. The last tile may be shorter if the dimension + doesn't divide evenly. If ``min_tile_size`` is set and the last tile is + shorter, it is grown leftward (penultimate ``right_ramp`` widens); the + result is validated and invalid layouts raise ``ValueError``. + Args: + size: Target tile size (in axis units). + overlap: Overlap between consecutive tiles. + min_tile_size: Optional minimum tile length. ``None`` keeps legacy + short-last-tile behavior. + Returns: + A split operation that divides a dimension into tiles. + """ + if size <= 0: + raise ValueError(f"size must be > 0, got {size}") + if overlap < 0 or overlap >= size: + raise ValueError(f"overlap must satisfy 0 <= overlap < size, got overlap={overlap}, size={size}") + if min_tile_size is not None and min_tile_size < 1: + raise ValueError(f"min_tile_size must be >= 1, got {min_tile_size}") + + def split(dimension_size: int) -> DimensionIntervals: + if min_tile_size is not None and dimension_size < min_tile_size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap) + intervals = [ + DimensionInterval(start=0, end=size, left_ramp=0, right_ramp=overlap), + *( + DimensionInterval( + start=i * (size - overlap), + end=i * (size - overlap) + size, + left_ramp=overlap, + right_ramp=overlap, + ) + for i in range(1, amount - 1) + ), + DimensionInterval( + start=(amount - 1) * (size - overlap), end=dimension_size, left_ramp=overlap, right_ramp=0 + ), + ] + if min_tile_size is not None: + intervals = _grow_last_tile_to_min(intervals, min_tile_size) + _validate_tile_intervals(intervals, dim_size=dimension_size, min_tile_size=min_tile_size) + return DimensionIntervals(intervals=intervals) + + return split + + +def split_temporal_causal(size: int, overlap: int, min_tile_size: int | None = None) -> SplitOperation: + """Split a temporal axis into overlapping tiles with causal handling. + Each tile after the first is shifted back by 1 and its left ramp is + increased by 1, ensuring causal continuity through the blend ramps. + Args: + size: Tile size in axis units. + overlap: Overlap between tiles in the same units. + min_tile_size: Optional floor forwarded to :func:`split_by_size`. + Returns: + Split operation that divides temporal dimension with causal handling. + """ + non_causal_split = split_by_size(size, overlap, min_tile_size=min_tile_size) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + dim_intervals = non_causal_split(dimension_size) + if len(dim_intervals.intervals) <= 1: + return dim_intervals + modified_intervals = [dim_intervals.intervals[0]] + [ + replace(interval, start=interval.start - 1, left_ramp=interval.left_ramp + 1) + for interval in dim_intervals.intervals[1:] + ] + return DimensionIntervals(intervals=modified_intervals) + + return split + + +def split_temporal(tile_size_frames: int, overlap_frames: int) -> SplitOperation: + """Split a temporal axis in video frame space into overlapping tiles. + Args: + tile_size_frames: Tile length in frames. + overlap_frames: Overlap between consecutive tiles in frames. + Returns: + Split operation that takes frame count and returns DimensionIntervals in frame indices. + """ + non_causal_split = split_by_size(tile_size_frames, overlap_frames) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= tile_size_frames: + return DEFAULT_SPLIT_OPERATION(dimension_size) + dim_intervals = non_causal_split(dimension_size) + modified_intervals = [ + replace(interval, end=interval.end + 1, right_ramp=0) for interval in dim_intervals.intervals[:-1] + ] + [replace(dim_intervals.intervals[-1], right_ramp=0)] + return DimensionIntervals(intervals=modified_intervals) + + return split + + +def split_by_count_temporal_causal( + num_tiles: int, overlap: int = 0, min_tile_size: int | None = None +) -> SplitOperation: + """Split a temporal dimension by count with causal handling. + Wraps :func:`split_by_count` with the same causal adjustment as + :func:`split_temporal_causal`: each tile after the first is shifted + back by 1 and its left ramp is increased by 1. + Args: + num_tiles: Number of tiles. Must be >= 1. + overlap: Overlap between adjacent tiles (default 0). + min_tile_size: Optional floor forwarded to :func:`split_by_count`. + Returns: + A split operation that divides a temporal dimension into tiles. + """ + non_causal_split = split_by_count(num_tiles, overlap, min_tile_size=min_tile_size) + + def split(dimension_size: int) -> DimensionIntervals: + dim_intervals = non_causal_split(dimension_size) + if len(dim_intervals.intervals) <= 1: + return dim_intervals + modified_intervals = [dim_intervals.intervals[0]] + [ + replace(interval, start=interval.start - 1, left_ramp=interval.left_ramp + 1) + for interval in dim_intervals.intervals[1:] + ] + return DimensionIntervals(intervals=modified_intervals) + + return split + + +def split_at_seams(boundaries: Sequence[int], num_tiles: int, overlap: int = 0) -> SplitOperation: + """Split a dimension on boundary cells whose content is already known, dropping the overlap. + ``boundaries`` are the ``K + 1`` segment edges in grid cells, starting at 0 and ending at the + last cell of the dimension. The ``K`` segments are dealt largest-first so leftover segments go to the leading tiles; + ``num_tiles`` larger than ``K`` is clamped. Each tile but the first starts ``overlap`` cells + before the boundary it resumes after. That run-up is context only: it lands in the interval's + ``left_ramp``, which :func:`identity_mapping_operation` with ``rectangular=True`` masks to zero, + so the earlier tile keeps the boundary cell and this one contributes strictly after it. + The point of cutting here is that nothing needs blending. A ramp is what a pair of tiles needs + when neither of them knows the truth at the seam; on a boundary cell both reproduce the same + known frame, so averaging them only smears it. + Args: + boundaries: Segment edges in grid cells, strictly increasing, starting at 0. + num_tiles: Number of tiles. Must be >= 1. Extra tiles beyond the segment count are dropped. + overlap: Context cells each non-first tile denoises before the cell it resumes at, in grid + units. Clamped at the start of the dimension. + Returns: + A split operation that divides a dimension on ``boundaries``. + """ + boundaries = tuple(boundaries) + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if overlap < 0: + raise ValueError(f"overlap must be >= 0, got {overlap}") + if len(boundaries) < 2 or boundaries[0] != 0: + raise ValueError(f"boundaries must start at 0 and hold at least one segment, got {list(boundaries)}") + if any(b <= a for a, b in itertools.pairwise(boundaries)): + raise ValueError(f"boundaries must be strictly increasing, got {list(boundaries)}") + n_segments = len(boundaries) - 1 + n_tiles = min(num_tiles, n_segments) + base, leftover = divmod(n_segments, n_tiles) + counts = [base + (1 if index < leftover else 0) for index in range(n_tiles)] + + def split(dim_size: int) -> DimensionIntervals: + if boundaries[-1] != dim_size - 1: + raise ValueError(f"boundaries must end at the last cell ({dim_size - 1}), got {boundaries[-1]}") + intervals: list[DimensionInterval] = [] + cursor = 0 + for tile_index, count in enumerate(counts): + resume = boundaries[cursor] + 1 + start = 0 if tile_index == 0 else max(0, resume - overlap) + cursor += count + intervals.append( + DimensionInterval( + start=start, + end=boundaries[cursor] + 1, + left_ramp=0 if tile_index == 0 else resume - start, + right_ramp=0, + ) + ) + return DimensionIntervals(intervals=intervals) + + return split + + +def split_by_count(num_tiles: int, overlap: int = 0, min_tile_size: int | None = None) -> SplitOperation: + """Split a dimension into a given number of tiles with overlap. + Computes the tile size as + ``(dim_size + overlap * (num_tiles - 1)) // num_tiles`` so that + ``num_tiles`` tiles of that size with ``overlap`` shared elements + cover the dimension evenly. Delegates to :func:`split_by_size` for + the actual interval construction. + When the total ``dim_size + overlap * (num_tiles - 1)`` is not evenly + divisible by ``num_tiles``, the first ``remainder`` tiles each absorb + one extra unit. + Args: + num_tiles: Number of tiles. Must be >= 1. + overlap: Overlap between adjacent tiles (default 0). Must be >= 0 + and less than the computed tile size. + min_tile_size: Optional floor forwarded to last-tile growth / validation. + Returns: + A split operation that divides a dimension into tiles. + """ + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if overlap < 0: + raise ValueError(f"overlap must be >= 0, got {overlap}") + if min_tile_size is not None and min_tile_size < 1: + raise ValueError(f"min_tile_size must be >= 1, got {min_tile_size}") + + def split(dim_size: int) -> DimensionIntervals: + if num_tiles > dim_size: + raise ValueError( + f"num_tiles ({num_tiles}) exceeds dim_size ({dim_size}). Cannot assign at least 1 unit per tile." + ) + if num_tiles == 1: + return DEFAULT_SPLIT_OPERATION(dim_size) + + total = dim_size + overlap * (num_tiles - 1) + tile_size = total // num_tiles + if tile_size <= overlap: + raise ValueError( + f"split_by_count produced size={tile_size} <= overlap={overlap} " + f"for dim_size={dim_size}, num_tiles={num_tiles}" + ) + remainder = total % num_tiles + + base_intervals = split_by_size(tile_size, overlap)(dim_size - remainder).intervals + + # First `remainder` tiles each absorb 1 extra unit; shift subsequent boundaries. + intervals: list[DimensionInterval] = [] + for i, iv in enumerate(base_intervals): + shift = min(i, remainder) + grow = 1 if i < remainder else 0 + intervals.append(replace(iv, start=iv.start + shift, end=iv.end + shift + grow)) + + if min_tile_size is not None: + intervals = _grow_last_tile_to_min(intervals, min_tile_size) + _validate_tile_intervals(intervals, dim_size=dim_size, min_tile_size=min_tile_size) + + return DimensionIntervals(intervals=intervals) + + return split + + +def identity_mapping_operation( + intervals: DimensionIntervals, + *, + rectangular: bool = False, +) -> tuple[list[slice], list[torch.Tensor]]: + """Map each DimensionInterval to an output region at the same position. + For every interval the output start/end matches the input start/end and a 1-D mask is built + from the interval's left_ramp and right_ramp. The default mask is trapezoidal (blend on the + ramps). ``rectangular=True`` drops the ramps outright: the overlap is context the tile denoised + but does not contribute. Pair that with a split whose ramps are one-sided, such as + :func:`split_at_seams` -- ramps on both sides of an interval would leave a hole between tiles. + """ + mask_1d = compute_rectangular_mask_1d if rectangular else compute_trapezoidal_mask_1d + out_slices: list[slice] = [] + masks: list[torch.Tensor] = [] + for iv in intervals.intervals: + out_slices.append(slice(iv.start, iv.end)) + masks.append(mask_1d(iv.end - iv.start, iv.left_ramp, iv.right_ramp)) + return out_slices, masks + + +class Tile(NamedTuple): + """ + Represents a single tile. + Attributes: + in_coords: + Tuple of slices specifying where to cut the tile from the INPUT tensor. + out_coords: + Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor. + masks_1d: + Per-dimension masks in OUTPUT units. + Untiled axes use a length-1 ones tensor (broadcasts). These are used + for separable blending (and for the dense ``blend_mask`` property). + Methods: + blend_mask: + Create a single N-D mask from the per-dimension masks. + """ + + in_coords: tuple[slice, ...] + out_coords: tuple[slice, ...] + masks_1d: tuple[torch.Tensor, ...] + + @property + def blend_mask(self) -> torch.Tensor: + num_dims = len(self.out_coords) + per_dimension_masks: list[torch.Tensor] = [] + + for dim_idx in range(num_dims): + mask_1d = self.masks_1d[dim_idx] + view_shape = [1] * num_dims + # Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply. + view_shape[dim_idx] = mask_1d.shape[0] + per_dimension_masks.append(mask_1d.view(*view_shape)) + + # Multiply per-dimension masks to form the full N-D mask (separable blending window). + combined_mask = per_dimension_masks[0] + for mask in per_dimension_masks[1:]: + combined_mask = combined_mask * mask + + return combined_mask + + +def scale_by_masks_1d(x: torch.Tensor, masks_1d: Sequence[torch.Tensor]) -> torch.Tensor: + """Multiply ``x`` by separable 1d masks with broadcasting. + ``len(masks_1d)`` must equal ``x.ndim``. Prefer float32 masks so bf16/fp16 ``x`` promotes. + Length-1 masks (untiled axes) broadcast over that dimension. + """ + if len(masks_1d) != x.ndim: + raise ValueError(f"masks_1d length {len(masks_1d)} != x.ndim {x.ndim}") + out = x + for axis, mask in enumerate(masks_1d): + view_shape = [1] * x.ndim + view_shape[axis] = -1 + out = out * mask.reshape(*view_shape) + return out + + +def masks_are_complementary( + tiles: Sequence[Tile], + full_shape: Sequence[int], + *, + atol: float = 1e-5, +) -> bool: + """Return whether per-axis 1d blend masks partition unity (sum to 1). + Checks each axis independently over the unique out-slices on that axis + (cartesian tile products would otherwise multi-count the same 1d interval). + When True, weighted accumulation needs no denominator. + """ + if not tiles: + return True + ndim = len(full_shape) + for tile in tiles: + if len(tile.out_coords) != ndim or len(tile.masks_1d) != ndim: + raise ValueError( + f"Tile out_coords/masks_1d rank {len(tile.out_coords)}/{len(tile.masks_1d)} != full_shape rank {ndim}" + ) + for axis, length in enumerate(full_shape): + # Explicit CPU float32: masks may live on CUDA; a non-CPU default device + # must not place ``acc`` on GPU (device-mismatch on ``acc[sl] +=``). + acc = torch.zeros(length, dtype=torch.float32, device="cpu") + seen: set[tuple[int | None, int | None]] = set() + for tile in tiles: + sl = tile.out_coords[axis] + key = (sl.start, sl.stop) + if key in seen: + continue + seen.add(key) + # Length-1 untiled masks broadcast over ``acc[sl]``. + acc[sl] += tile.masks_1d[axis].detach().float().cpu() + if not torch.allclose(acc, torch.ones(length, dtype=torch.float32), atol=atol, rtol=0.0): + return False + return True + + +def compute_summed_weights( + tiles: Sequence[Tile], + full_shape: Sequence[int], +) -> torch.Tensor: + """Build the dense denominator for weighted blending over ``full_shape``. + Uses separable per-axis mask broadcasts — never ``Tile.blend_mask``. + Requires concrete ``out_coords`` (``stop`` not ``None``) on every axis. + Always builds on CPU float32 so CUDA masks / a non-CPU default device cannot + place a multi-GB ``[F,H,W]`` tensor on GPU. + """ + weights = torch.zeros(*full_shape, dtype=torch.float32, device="cpu") + for tile in tiles: + masks = tuple(m.detach().float().cpu() for m in tile.masks_1d) + region_shape = tuple(s.stop - s.start for s in tile.out_coords) + region = torch.ones(region_shape, dtype=torch.float32, device="cpu") + weights[tile.out_coords] += scale_by_masks_1d(region, masks) + return weights.clamp(min=1e-8) + + +def create_tiles_from_intervals_and_mappers( + intervals: LatentIntervals, + mappers: list[MappingOperation], +) -> list[Tile]: + full_dim_input_slices: list[list[slice]] = [] + full_dim_output_slices: list[list[slice]] = [] + full_dim_masks_1d: list[list[torch.Tensor]] = [] + for axis_index in range(len(intervals.original_shape)): + dimension_intervals = intervals.dimension_intervals[axis_index] + input_slices = [slice(interval.start, interval.end) for interval in dimension_intervals.intervals] + output_slices, masks_1d = mappers[axis_index](dimension_intervals) + n_intervals = len(input_slices) + if len(output_slices) != n_intervals or len(masks_1d) != n_intervals: + raise ValueError( + f"Axis {axis_index}: mapper produced {len(output_slices)} output slices and " + f"{len(masks_1d)} masks for {n_intervals} input intervals" + ) + full_dim_input_slices.append(input_slices) + full_dim_output_slices.append(output_slices) + full_dim_masks_1d.append(masks_1d) + + return [ + Tile(in_coords=in_coord, out_coords=out_coord, masks_1d=mask_1d) + for in_coord, out_coord, mask_1d in zip( + itertools.product(*full_dim_input_slices), + itertools.product(*full_dim_output_slices), + itertools.product(*full_dim_masks_1d), + strict=True, + ) + ] + + +def create_tiles( + latent_shape: torch.Size, + splitters: list[SplitOperation], + mappers: list[MappingOperation], +) -> list[Tile]: + if len(splitters) != len(latent_shape): + raise ValueError( + f"Number of splitters must be equal to number of dimensions in latent shape, " + f"got {len(splitters)} and {len(latent_shape)}" + ) + if len(mappers) != len(latent_shape): + raise ValueError( + f"Number of mappers must be equal to number of dimensions in latent shape, " + f"got {len(mappers)} and {len(latent_shape)}" + ) + intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)] + latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals)) + return create_tiles_from_intervals_and_mappers(latent_intervals, mappers) + + +def group_tiles_by_temporal_slice(tiles: list[Tile]) -> list[list[Tile]]: + """Group consecutive tiles that share the same temporal ``out_coords`` slice. + Assumes ``tiles`` is ordered with the temporal axis varying slowest (true + for every tile list this codebase builds via ``itertools.product`` with + the temporal axis first), so equal temporal slices are always contiguous. + """ + if not tiles: + return [] + + groups = [] + current_slice = tiles[0].out_coords[2] + current_group = [] + + for tile in tiles: + tile_slice = tile.out_coords[2] + if tile_slice == current_slice: + current_group.append(tile) + else: + groups.append(current_group) + current_slice = tile_slice + current_group = [tile] + + if current_group: + groups.append(current_group) + + return groups + + +@dataclass(frozen=True) +class DimensionTilingConfig: + """Tiling parameters for a single dimension of the patchified grid. + Attributes: + num_tiles: Number of tiles along this dimension. ``1`` with ``overlap=0`` + means the axis is not tiled. + overlap: Overlap between adjacent tiles, in latent grid units. + Adjacent tiles share ``overlap`` grid cells at their + boundary, producing an overlap zone blended with + trapezoidal masks. + """ + + num_tiles: int = 1 + overlap: int = 0 + + def __post_init__(self) -> None: + if self.num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {self.num_tiles}") + if self.overlap < 0: + raise ValueError(f"overlap must be >= 0, got {self.overlap}") + + def is_tiled(self) -> bool: + """True when this axis is split into more than one tile (or has overlap).""" + return self.num_tiles > 1 or self.overlap > 0 + + @classmethod + def from_tile_size(cls, dim_size: int, tile_size: int, overlap: int = 0) -> DimensionTilingConfig: + """Create config by computing ``num_tiles`` from dimension size and tile size. + Args: + dim_size: Total length of the dimension. + tile_size: Desired tile size. + overlap: Overlap between consecutive tiles. + Returns: + A ``DimensionTilingConfig`` with the computed ``num_tiles``. + """ + split_op = split_by_size(tile_size, overlap) + intervals = split_op(dim_size) + return cls(num_tiles=len(intervals.intervals), overlap=overlap) + + +@dataclass(frozen=True) +class DimensionSizeConfig: + """Tile size and overlap for a single video axis (frames / height / width). + Mirrors :class:`DimensionTilingConfig`, but specifies tile *size* rather than + tile *count*. ``tile_size=0`` means the axis is not tiled (covers the whole + length). Axis-specific VAE pixel constraints (divisibility / minimums) are + enforced by :meth:`TileSizeConfig.validate` for tiled axes only. + """ + + tile_size: int = 0 + overlap: int = 0 + + def __post_init__(self) -> None: + if self.tile_size < 0: + raise ValueError(f"tile_size must be >= 0, got {self.tile_size}") + if self.overlap < 0: + raise ValueError(f"overlap must be >= 0, got {self.overlap}") + if self.tile_size == 0: + if self.overlap != 0: + raise ValueError("untiled axis (tile_size=0) must have overlap=0") + return + if self.overlap >= self.tile_size: + raise ValueError(f"Overlap must be less than tile size, got {self.overlap} and {self.tile_size}") + + def is_tiled(self) -> bool: + """True when this axis has a positive tile size (caller intends to split it).""" + return self.tile_size > 0 + + +@dataclass(frozen=True) +class TileCountConfig: + """Tiling layout for a ``(F, H, W)`` grid by tile *counts*. + Overlaps are in latent-grid units. Mirror of :class:`TileSizeConfig`. + Attributes: + frames: Tiling along the temporal (frames) dimension. + height: Tiling along the latent height dimension. + width: Tiling along the latent width dimension. + """ + + frames: DimensionTilingConfig = DimensionTilingConfig() + height: DimensionTilingConfig = DimensionTilingConfig() + width: DimensionTilingConfig = DimensionTilingConfig() + + def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: + """Raise if this count layout cannot tile ``video_shape`` under ``scale_factors``. + Counts/overlaps are in latent-grid units. ``video_shape.frames <= 0`` skips the + temporal axis (duration not yet known). Spatial axes always checked. + """ + check_temporal = _assert_video_on_vae_grid(scale_factors, video_shape) + latent_h = video_shape.height // scale_factors.height + latent_w = video_shape.width // scale_factors.width + _validate_count_axis(self.height, latent_h, "height") + _validate_count_axis(self.width, latent_w, "width") + if check_temporal: + latent_f = (video_shape.frames - 1) // scale_factors.time + 1 + _validate_count_axis(self.frames, latent_f, "frames") + + def to_splitters( + self, + scale_factors: SpatioTemporalScaleFactors, + min_tile_size: tuple[int, int, int] | None = None, + *, + causal_temporal: bool = True, + ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: + """Build ``(T, H, W)`` latent-grid split operations for this count layout. + ``scale_factors`` is accepted for signature parity with + :meth:`TileSizeConfig.to_splitters` and ignored — counts are already in + grid units. When ``causal_temporal`` is True (VAE encode/decode), the + frames axis uses :func:`split_by_count_temporal_causal`; otherwise plain + :func:`split_by_count`. ``min_tile_size`` is a per-axis floor in the same + units as the split. + """ + del scale_factors + min_t = min_h = min_w = None + if min_tile_size is not None: + min_t, min_h, min_w = min_tile_size + + def axis_split(cfg: DimensionTilingConfig, axis_min: int | None, *, temporal: bool) -> SplitOperation: + if not cfg.is_tiled(): + return DEFAULT_SPLIT_OPERATION + if temporal and causal_temporal: + return split_by_count_temporal_causal(cfg.num_tiles, cfg.overlap, min_tile_size=axis_min) + return split_by_count(cfg.num_tiles, cfg.overlap, min_tile_size=axis_min) + + return ( + axis_split(self.frames, min_t, temporal=True), + axis_split(self.height, min_h, temporal=False), + axis_split(self.width, min_w, temporal=False), + ) + + def video_chunks_number(self, num_frames: int) -> int: + """Number of temporal decode chunks for ``num_frames`` under this layout.""" + del num_frames + return max(1, self.frames.num_tiles) + + +@dataclass(frozen=True) +class TileSizeConfig: + """Size-based tiling layout for a ``(F, H, W)`` video — mirror of ``TileCountConfig``. + Each axis is a non-optional :class:`DimensionSizeConfig`; ``tile_size=0`` means + untiled on that axis (:meth:`DimensionSizeConfig.is_tiled`). Sizes and overlaps + are in pixel / frame units. Conversion to a split grid is an explicit + ``scale_factors`` argument to :meth:`to_splitters` (not stored on the config). + Legality vs a VAE grid is checked by :meth:`validate` (same factors decode will + pass to :meth:`to_splitters`), not at construction. + Attributes: + frames: Temporal tile size/overlap in frames. + height: Spatial height tile size/overlap in pixels. + width: Spatial width tile size/overlap in pixels. + """ + + frames: DimensionSizeConfig = DimensionSizeConfig() + height: DimensionSizeConfig = DimensionSizeConfig() + width: DimensionSizeConfig = DimensionSizeConfig() + + def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: + """Raise if this size layout is illegal for ``video_shape`` under ``scale_factors``. + Checks tile/overlap divisibility and minimums against the VAE grid, and that + the video extents are compatible with that grid. ``video_shape.frames <= 0`` + skips the temporal axis (duration not yet known); height/width always checked. + """ + check_temporal = _assert_video_on_vae_grid(scale_factors, video_shape) + _validate_size_axis(self.height, scale_factors.height, "height") + _validate_size_axis(self.width, scale_factors.width, "width") + if check_temporal: + _validate_size_axis(self.frames, scale_factors.time, "frames") + + @classmethod + def default(cls) -> TileSizeConfig: + return cls( + frames=DimensionSizeConfig(tile_size=80, overlap=24), + height=DimensionSizeConfig(tile_size=768, overlap=64), + width=DimensionSizeConfig(tile_size=768, overlap=64), + ) + + @classmethod + def from_long_side( + cls, + *, + long_side: DimensionSizeConfig, + height: int, + width: int, + scale_factors: SpatioTemporalScaleFactors, + frames: DimensionSizeConfig | None = None, + ) -> TileSizeConfig: + """Aspect-coupled construction — old single-spatial long-side behavior, explicit. + Matches main-era ``latent_tile_splitters``: scale the long-side tile in + *latent* units with ``round(size_lat * axis_lat / long_lat)``, then + multiply back by the VAE factor. Pixel-space ``round`` + ceil-snap would + bias the short axis up by almost one latent (e.g. 680 → 704 vs 672). + Both axes share ``long_side.overlap``. + """ + if height < 1 or width < 1: + raise ValueError(f"height/width must be >= 1, got {height}x{width}") + if not long_side.is_tiled(): + raise ValueError("long_side must be tiled (tile_size > 0)") + if scale_factors.height < 1 or scale_factors.width < 1: + raise ValueError(f"scale_factors height/width must be >= 1, got {scale_factors}") + span = max(height, width) + + def axis_size(axis_len: int, factor: int) -> int: + # Latent-grid round (same as main decode enable_on_axis), not pixel ceil. + axis_lat = axis_len // factor + long_lat = span // factor + size_lat = long_side.tile_size // factor + overlap_lat = long_side.overlap // factor + lower_threshold = max(2, overlap_lat + 1) + tile_lat = max(lower_threshold, round(size_lat * axis_lat / long_lat)) + tile_px = tile_lat * factor + min_legal = max(2 * factor, long_side.overlap + factor) + return max(tile_px, min_legal) + + return cls( + frames=DimensionSizeConfig() if frames is None else frames, + height=DimensionSizeConfig(tile_size=axis_size(height, scale_factors.height), overlap=long_side.overlap), + width=DimensionSizeConfig(tile_size=axis_size(width, scale_factors.width), overlap=long_side.overlap), + ) + + def to_splitters( + self, + scale_factors: SpatioTemporalScaleFactors, + min_tile_size: tuple[int, int, int] | None = None, + *, + causal_temporal: bool = True, + ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: + """Build ``(T, H, W)`` grid split ops from pixel/frame sizes via ``scale_factors``. + When ``causal_temporal`` is True, frames use :func:`split_temporal_causal`. + """ + min_t = min_h = min_w = None + if min_tile_size is not None: + min_t, min_h, min_w = min_tile_size + + def enable_size_axis( + factor: int, + axis_min: int | None, + cfg: DimensionSizeConfig, + axis_name: str, + *, + temporal: bool, + ) -> SplitOperation: + if not cfg.is_tiled(): + return DEFAULT_SPLIT_OPERATION + _validate_size_axis(cfg, factor, axis_name) + size = cfg.tile_size // factor + overlap = cfg.overlap // factor + lower_threshold = max(2, overlap + 1) + tile = max(lower_threshold, size) + if temporal and causal_temporal: + return split_temporal_causal(tile, overlap, min_tile_size=axis_min) + return split_by_size(tile, overlap, min_tile_size=axis_min) + + return ( + enable_size_axis(scale_factors.time, min_t, self.frames, "frames", temporal=True), + enable_size_axis(scale_factors.height, min_h, self.height, "height", temporal=False), + enable_size_axis(scale_factors.width, min_w, self.width, "width", temporal=False), + ) + + def video_chunks_number(self, num_frames: int, *, time_scale: int = VIDEO_SCALE_FACTORS.time) -> int: + """Number of temporal decode chunks for ``num_frames`` under this layout. + Mirrors what decode actually does: :meth:`to_splitters` converts this axis to the + latent grid and hands it to :func:`split_by_size`, so the count must be taken there + too. Doing the arithmetic in pixel units instead over-reports by one whenever the + trailing tile is absorbed -- including the common case of a tile larger than the + clip, which is a single tile but used to report two. + """ + if not self.frames.is_tiled(): + return 1 + # Same derivation as ``to_splitters.enable_size_axis``. + overlap = self.frames.overlap // time_scale + size = max(2, overlap + 1, self.frames.tile_size // time_scale) + latent_frames = (num_frames - 1) // time_scale + 1 + if latent_frames <= size: + return 1 + # Same tile count as ``split_by_size``. + return (latent_frames + size - 2 * overlap - 1) // (size - overlap) + + +TilingConfig = TileSizeConfig | TileCountConfig + + +class AutoTiling: + """Sentinel: pipeline should recommend decode tiling (DiffVAE-aware / Conv default). + Distinct from ``None``, which means untiled decode. + """ + + __slots__ = () + + def __repr__(self) -> str: + return "AUTO_TILING" + + +AUTO_TILING = AutoTiling() + + +PipelineTiling = TilingConfig | AutoTiling | None + + +def _assert_video_on_vae_grid( + scale_factors: SpatioTemporalScaleFactors, + video_shape: VideoPixelShape, +) -> bool: + """Raise if ``video_shape`` is incompatible with the VAE ``scale_factors`` grid. + Returns whether the temporal axis is known (``video_shape.frames > 0``). When + False, callers skip frames-axis checks (duration not yet resolved). + """ + if scale_factors.time < 1 or scale_factors.height < 1 or scale_factors.width < 1: + raise ValueError(f"scale_factors must be >= 1 on each axis, got {scale_factors}") + if video_shape.height < 1 or video_shape.width < 1: + raise ValueError(f"video_shape height/width must be >= 1, got {video_shape.height}x{video_shape.width}") + if video_shape.height % scale_factors.height != 0: + raise ValueError(f"video height {video_shape.height} must be divisible by scale {scale_factors.height}") + if video_shape.width % scale_factors.width != 0: + raise ValueError(f"video width {video_shape.width} must be divisible by scale {scale_factors.width}") + if video_shape.frames <= 0: + return False + if (video_shape.frames - 1) % scale_factors.time != 0: + raise ValueError(f"video frames {video_shape.frames} must satisfy (frames - 1) % {scale_factors.time} == 0") + return True + + +def _validate_size_axis(cfg: DimensionSizeConfig, factor: int, axis_name: str) -> None: + """Pixel/frame size-axis legality vs VAE ``factor``.""" + if not cfg.is_tiled(): + return + min_size = 2 * factor + if cfg.tile_size < min_size: + raise ValueError(f"{axis_name}.tile_size must be at least {min_size}, got {cfg.tile_size}") + if cfg.tile_size % factor != 0: + raise ValueError(f"{axis_name}.tile_size must be divisible by {factor}, got {cfg.tile_size}") + if cfg.overlap % factor != 0: + raise ValueError(f"{axis_name}.overlap must be divisible by {factor}, got {cfg.overlap}") + + +def _validate_count_axis(cfg: DimensionTilingConfig, latent_extent: int, axis_name: str) -> None: + """Latent count-axis legality vs latent ``extent``.""" + if not cfg.is_tiled(): + return + if cfg.num_tiles > latent_extent: + raise ValueError(f"{axis_name}.num_tiles {cfg.num_tiles} exceeds latent {axis_name} extent {latent_extent}") + # split_by_count requires overlap < tile_size; tile_size grows with extent/n. + max_overlap = latent_extent - cfg.num_tiles + if cfg.overlap > max_overlap: + raise ValueError( + f"{axis_name}.overlap {cfg.overlap} exceeds latent bound {max_overlap} " + f"for extent {latent_extent} with {cfg.num_tiles} tiles" + ) + + +def _validate_overlap( + tiling_config: TilingConfig, + *, + min_overlap_frames: int, + min_overlap_pixels: int, +) -> None: + """Raise if any tiled ``TileSizeConfig`` axis overlap is below the given floors.""" + if not isinstance(tiling_config, TileSizeConfig): + return + + for axis_name, cfg, recommended, unit in ( + ("frames", tiling_config.frames, min_overlap_frames, "frames"), + ("height", tiling_config.height, min_overlap_pixels, "px"), + ("width", tiling_config.width, min_overlap_pixels, "px"), + ): + if cfg.is_tiled() and cfg.overlap < recommended: + raise ValueError(f"{axis_name} overlap {cfg.overlap} {unit} is below the required {recommended} {unit}.") + + +def balanced_tile_split(num_tiles: int) -> tuple[int, int]: + """Factor ``num_tiles`` into ``(small, large)`` as square as possible. + ``small`` is the largest divisor not exceeding the square root, so + ``small * large == num_tiles`` and ``small <= large``. E.g. 2 -> (1, 2), + 4 -> (2, 2), 8 -> (2, 4), 16 -> (4, 4). The caller decides which tiled + dimension gets which factor. + """ + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + small = next(d for d in range(math.isqrt(num_tiles), 0, -1) if num_tiles % d == 0) + return small, num_tiles // small + + + +class DiffVAEMode(Enum): + COMBINED_COMPILE = "combined_compile" + CHUNKED_COMPILE = "chunked_compile" + CHUNKED_EAGER = "chunked_eager" + BLACKWELL_DSL = "blackwell_dsl" + + def resolve(self): + return self + + +class NAttentionKind(Enum): + TRITON = "triton" + EAGER_SDPA = "eager_sdpa" + + +@dataclass(frozen=True) +class _ResolvedAttention: + attention: NAttentionKind = NAttentionKind.EAGER_SDPA + compile_blocks: bool = False + + +def resolve_attention_for_host(mode): + del mode + return _ResolvedAttention() + + +def frames_per_yuv_gemm(height: int, width: int) -> int: + del height, width + return 2**31 - 1 + + + +def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks + and moves pixels from each block into separate channels (space-to-depth). + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks. + patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching). + For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw) + Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b c (f p) (h q) (w r) -> b (c p r q) f h w", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + + +def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from + channels back into patch_size x patch_size blocks (depth-to-space). + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x. + patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion). + For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw) + Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b (c p r q) f h w -> b c (f p) (h q) (w r)", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + + return x + + +class PerChannelStatistics(nn.Module): + """ + Per-channel statistics for normalizing and denormalizing the latent representation. + This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict. + Defaults are identity (std=1, mean=0) so models constructed without a checkpoint + do not inherit allocator garbage / NaNs from ``torch.empty``. + """ + + def __init__(self, latent_channels: int = 128): + super().__init__() + self.register_buffer("std-of-means", torch.ones(latent_channels)) + self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + +"""Keyframe (dual-stream) inputs and coordinate math for DiffVAE decode. +A keyframe-aware decode carries two streams through the decoder: the video volume +``(B, T, H, W, C)`` and a stack of keyframe *planes* ``(B, P, H, W, C)`` whose plane +axis occupies video's temporal slot. Weights are fully shared; the streams only ever +mix inside one joint attention softmax (see ``transformer/fallback_na/joint_eager.py``). +Everything here is pure coordinate/geometry math with no module state, so the eager and +triton backends can share it and therefore agree exactly on slot selection. +Deviation from upstream worth knowing: upstream carries per-sample keyframe times and +masks (``(B, n_kf)``). Here they are batch-shared 1-D ``(P,)`` tensors, because our +decode path is single-sample and ``rope_math.rot_abs_axis_impl`` takes a 1-D position +vector per axis. That keeps the RoPE call and the slot tables batch-independent. +""" + + +KEYFRAME_CONTEXT_SLOTS = 2 + + +@dataclass(frozen=True) +class DecodeKeyframes: + """Caller-facing keyframe input to a DiffVAE decode. + Attributes: + latents: ``(B, C, P, H, W)`` per-channel-normalized latents, exactly one latent + frame per keyframe. Each plane must have been encoded as a standalone + one-pixel-frame clip -- the VAE is causal, so a ``P``-frame encode would + blend planes that were never temporally adjacent. + pixel_frame_indices: ``(P,)`` int64 **global** pixel frame index of each plane. + Never rebased onto a tile. A Dist slice whose first pixel is 56 still carries + a plane at 48 as ``48``; :attr:`clip_start_frame` is how DiffVAE learns the + 8-frame gap. + clip_start_frame: first global pixel frame of the video latent in this decode + call. ``0`` for a full-clip decode. Dist sets it to the tile origin so stage + times are ``t_s(index) - t_s(clip_start)``. + """ + + latents: torch.Tensor + pixel_frame_indices: torch.Tensor + clip_start_frame: int = 0 + + def validate(self, *, num_frames: int | None = None) -> None: + """Raise if shapes/indices are inconsistent (optionally against a frame count).""" + if self.latents.ndim != 5: + raise ValueError(f"keyframe latents must be (B, C, P, H, W), got {tuple(self.latents.shape)}") + if self.pixel_frame_indices.ndim != 1: + raise ValueError(f"pixel_frame_indices must be 1-D (P,), got {tuple(self.pixel_frame_indices.shape)}") + if self.clip_start_frame < 0: + raise ValueError(f"clip_start_frame must be non-negative, got {self.clip_start_frame}") + planes = self.latents.shape[2] + if planes != self.pixel_frame_indices.shape[0]: + raise ValueError( + f"keyframe plane count {planes} != len(pixel_frame_indices) {self.pixel_frame_indices.shape[0]}" + ) + if planes == 0: + # An empty stack is a plain decode wearing a keyframe decode's costs, and every + # backend has to special-case it (the slot tables are all -1, and gathering plane 0 + # of an empty axis is an out-of-bounds read). Say so here instead. + raise ValueError("keyframe decode needs at least one plane; use decode_video() for a plain decode") + if planes and int(self.pixel_frame_indices.min()) < 0: + raise ValueError("pixel_frame_indices must be non-negative (global pixel frames)") + if num_frames is not None and num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}") + # Planes may sit outside [clip_start_frame, clip_start_frame + num_frames): Dist tiles + # keep the nearest plane on each side so |dt| matches a whole-clip decode. A far plane + # on a full clip is the same geometry -- joint attention ranks it by distance. + + def for_frame_span(self, frame_lo: int, frame_hi: int) -> "DecodeKeyframes": + """Keep the planes a decode of pixel frames ``[lo, hi]`` needs; indices stay global. + Selection is :func:`planes_for_tile`: every plane inside the span **plus the nearest + plane on each side outside it**. Those two are not optional. DiffVAE's joint attention + picks a frame's anchors by ``|dt|``, so a window ending at frame 64 whose last inside + plane is 48 still has to carry the plane at 96 -- drop it and frames near the boundary + anchor on 48 alone, which is exactly how a split decode stops matching a whole one. + ``pixel_frame_indices`` are not rewritten. :attr:`clip_start_frame` becomes ``frame_lo`` + so the decoder subtracts ``t_s(48) - t_s(56)`` rather than treating the slice as a new + clip that starts at pixel 0. + """ + keep = planes_for_tile(self.pixel_frame_indices, frame_lo, frame_hi) + return DecodeKeyframes( + latents=self.latents[:, :, keep.to(self.latents.device)], + pixel_frame_indices=self.pixel_frame_indices[keep.to(self.pixel_frame_indices.device)], + clip_start_frame=frame_lo, + ) + + def crop_spatial(self, height: slice, width: slice) -> "DecodeKeyframes": + """Crop the planes to a spatial window, with the *same* latent slices the video used. + For a decode that splits the latent across workers (see + :class:`~ltx_core.multigpu.vae.distributed_decoder.DistributedVideoDecoder`): each worker + holds a crop of the video latent, so it must hold the matching crop of every keyframe + plane. Cropping one and not the other offsets every plane from the video by the + difference, which reads as ghosting rather than as an obvious failure. + Plane count, ``pixel_frame_indices``, and :attr:`clip_start_frame` are untouched -- a + spatial split leaves every worker the full frame range. + """ + return DecodeKeyframes( + latents=self.latents[:, :, :, height, width], + pixel_frame_indices=self.pixel_frame_indices, + clip_start_frame=self.clip_start_frame, + ) + + @property + def num_planes(self) -> int: + return int(self.latents.shape[2]) + + +@dataclass(frozen=True) +class KeyframeStream: + """The keyframe half of the dual stream at one decoder stage. + Attributes: + x: ``(B, P, H, W, C)`` channels-last activations. ``H``/``W`` always match the + video stream at the same stage; ``P`` is invariant across the whole decode. + times: ``(P,)`` float32 plane position in *this stage's* temporal units and + *local to the current tile* -- the same origin the video stream's RoPE uses. + Both streams must share one origin or the joint softmax sees wrong offsets. + valid: ``(P,)`` bool. Invalid planes are masked out of every softmax and their + activations are re-zeroed after each upsample. + """ + + x: torch.Tensor + times: torch.Tensor + valid: torch.Tensor + + def masked(self) -> KeyframeStream: + """Re-zero invalid planes' activations (channels-last plane axis).""" + return KeyframeStream(x=self.x * self.valid[None, :, None, None, None], times=self.times, valid=self.valid) + + def select_planes(self, keep: torch.Tensor) -> KeyframeStream: + """Subset the plane axis, keeping ``x``/``times``/``valid`` in step. + ``keep`` is a ``(P,)`` bool mask. Used by tiled decode, which gives each tile only the + planes near it -- see :func:`planes_for_tile`. + """ + if keep.shape != (self.num_planes,): + raise ValueError(f"keep must be ({self.num_planes},) bool, got {tuple(keep.shape)}") + return KeyframeStream(x=self.x[:, keep], times=self.times[keep], valid=self.valid[keep]) + + def crop_spatial(self, height: slice, width: slice) -> KeyframeStream: + """Crop H/W with the *same* slices the video stream's tile used. + Cropping only one stream offsets every plane from the video by the difference, which + reads as ghosting rather than as an obvious failure -- the same hazard as the spatial + padding rule. + """ + return KeyframeStream(x=self.x[:, :, height, width, :], times=self.times, valid=self.valid) + + @property + def num_planes(self) -> int: + return int(self.x.shape[1]) + + +def keyframe_stage_times(pixel_frame_indices: torch.Tensor, remaining_time_stride: int) -> torch.Tensor: + """Chunk-center position of each keyframe in a stage's temporal units. + A stage whose remaining temporal upsampling is ``r`` has cells covering ``r`` pixel + frames each, except cell 0 which covers only pixel frame 0 (the causal first frame). + So ``t_s(0) = 0`` and ``t_s(f) = (f + (r - 1) / 2) / r`` -- the center of the chunk + holding ``f``. At stage 5 ``r == 1``, making the times the raw pixel indices. + Args: + pixel_frame_indices: ``(P,)`` global pixel frame index per plane. + remaining_time_stride: product of the temporal strides *still to come*. + """ + if remaining_time_stride < 1: + raise ValueError(f"remaining_time_stride must be positive, got {remaining_time_stride}") + frames = pixel_frame_indices.to(torch.float32) + center_offset = (remaining_time_stride - 1) / 2 + times = (frames + center_offset) / remaining_time_stride + return torch.where(frames == 0, torch.zeros_like(times), times) + + +def keyframe_clip_times( + pixel_frame_indices: torch.Tensor, + remaining_time_stride: int, + clip_start_frame: int, + extra_origin: float = 0.0, +) -> torch.Tensor: + """Stage times relative to a decode whose first pixel frame is ``clip_start_frame``. + ``t_s(global) - t_s(clip_start)`` is the gap joint attention should see. Rebasing the + indices onto the tile (``48 -> -8``) and then calling :func:`keyframe_stage_times` is not + the same: ``t_s`` is not linear through a fake clip start, so stages with ``r > 1`` get + the wrong ``|dt|``. + Single-GPU tiled decode uses ``clip_start_frame=0`` and passes the in-volume tile origin + as ``extra_origin``. A Dist slice whose first pixel is global 56 uses + ``clip_start_frame=56`` and ``extra_origin=0``. + """ + times = keyframe_stage_times(pixel_frame_indices, remaining_time_stride) + origin = keyframe_stage_times( + torch.as_tensor([clip_start_frame], dtype=torch.int64, device=pixel_frame_indices.device), + remaining_time_stride, + ) + return times - origin - extra_origin + + +def planes_for_tile( + pixel_frame_indices: torch.Tensor, + frame_lo: int, + frame_hi: int, + *, + clip_start_frame: int = 0, +) -> torch.Tensor: + """``(P,)`` bool: which planes a tile spanning pixel frames ``[lo, hi]`` should carry. + Every plane inside the span, **plus the nearest plane on each side outside it**. Those two + boundary planes are the point: without them a video frame at a tile edge ranks only + in-tile planes and attends to the wrong one, which is what made tiled keyframe decode + non-invariant. They arrive with negative / past-the-end tile-local times, and the + ``(|dt|, index)`` slot ranking already handles those, so nothing downstream changes. + Selection is by *value*, not position: ``pixel_frame_indices`` is not required to be + sorted. + Args: + pixel_frame_indices: ``(P,)`` global pixel frame index per plane. + frame_lo: first pixel frame in the tile, relative to ``clip_start_frame``. + frame_hi: last pixel frame in the tile (inclusive), relative to ``clip_start_frame``. + clip_start_frame: first global pixel of this latent. A full-clip decode leaves it 0. + Dist tiles keep global indices and pass the slice origin so a local ``[0, 72)`` + still selects global ``[56, 127]``. + """ + frame_lo = frame_lo + clip_start_frame + frame_hi = frame_hi + clip_start_frame + indices = pixel_frame_indices.to(torch.int64) + keep = (indices >= frame_lo) & (indices <= frame_hi) + before = indices < frame_lo + if bool(before.any()): + # Latest plane strictly before the tile. + keep[int(torch.where(before, indices, torch.full_like(indices, -1)).argmax())] = True + after = indices > frame_hi + if bool(after.any()): + # Earliest plane strictly after the tile. + sentinel = int(indices.max()) + 1 + keep[int(torch.where(after, indices, torch.full_like(indices, sentinel)).argmin())] = True + return keep + + +def remaining_time_strides(upsamples: Sequence[torch.nn.Module]) -> tuple[int, ...]: + """Remaining temporal upsampling at each stage input, plus 1 for stage 5. + For the production ladder (temporal strides ``1, 2, 2, 2``) this is + ``(8, 8, 4, 2, 1)``: stage ``i``'s blocks see the product of strides ``i..end``. + """ + strides = [int(up.stride[0]) for up in upsamples] + remaining: list[int] = [] + for index in range(len(strides)): + product = 1 + for stride in strides[index:]: + product *= stride + remaining.append(product) + remaining.append(1) + return tuple(remaining) + + +def upsample_keyframe_planes(upsample: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + """Spatially upsample keyframe planes, keeping the plane count invariant. + Each plane is folded into the batch as its own ``T=1`` clip and pushed through the + *same* upsample module as video, always with ``drop_leading_frame=True``: a temporal + stride of 2 expands ``T=1`` to 2 and the leading-frame drop takes it back to 1, + keeping phase 1. So temporal strides collapse and only ``H``/``W`` grow. + Passing ``drop_leading_frame=False`` here (as tiled video does for non-origin tiles) + would invent a second temporal plane per keyframe and is always wrong. + Args: + upsample: the video stream's ``LinearPixelShuffleUpsample`` for this stage. + x: ``(B, P, H, W, C)`` keyframe activations. + """ + planes = x.shape[1] + flat = rearrange(x, "b p h w c -> (b p) 1 h w c") + upsampled = upsample(flat, drop_leading_frame=True) + if upsampled.shape[1] != 1: + raise RuntimeError(f"isolated keyframe upsampling must preserve one temporal plane, got T={upsampled.shape[1]}") + out = rearrange(upsampled[:, 0], "(b p) h w c -> b p h w c", p=planes) + if out.shape[1] != planes: + raise RuntimeError(f"keyframe plane count changed under upsample: {planes} -> {out.shape[1]}") + return out + + +def _nearest_slots( + query_times: torch.Tensor, + candidate_times: torch.Tensor, + candidate_valid: torch.Tensor | None, + num_slots: int, +) -> torch.Tensor: + """``(Q, num_slots)`` candidate indices ranked by ``(|dt|, index)``, ``-1`` when empty. + A stable argsort on ``|dt|`` breaks ties by ascending candidate index, which is + exactly upstream's ``distances + arange * 1e-6`` tie-break. + """ + distances = (query_times[:, None] - candidate_times[None, :]).abs().to(torch.float32) + if candidate_valid is not None: + distances = distances.masked_fill(~candidate_valid[None, :], float("inf")) + order = torch.argsort(distances, dim=-1, stable=True) + take = min(num_slots, candidate_times.shape[0]) + chosen = order[:, :take] + # Drop slots that only exist because every remaining candidate was invalid. + finite = torch.gather(distances, 1, chosen).isfinite() + chosen = torch.where(finite, chosen, torch.full_like(chosen, -1)) + if take < num_slots: + pad = torch.full((chosen.shape[0], num_slots - take), -1, dtype=chosen.dtype, device=chosen.device) + chosen = torch.cat([chosen, pad], dim=1) + return chosen + + +def video_keyframe_slots( + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + video_length: int, + num_slots: int = KEYFRAME_CONTEXT_SLOTS, +) -> torch.Tensor: + """``(T, num_slots)`` keyframe plane index per video frame, ``-1`` for an empty slot. + Ranked by ``(|t_s(plane) - t|, plane)``. Deliberately independent of the temporal + kernel: the nearest planes are visible even when they lie outside ``K_t``. + """ + query = torch.arange(video_length, dtype=torch.float32, device=keyframe_times.device) + return _nearest_slots(query, keyframe_times.to(torch.float32), keyframe_valid, num_slots) + + +def keyframe_video_slots( + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + video_length: int, + num_slots: int = KEYFRAME_CONTEXT_SLOTS, +) -> torch.Tensor: + """``(P, num_slots)`` video frame index per keyframe plane, ``-1`` for an empty slot. + Ranked by ``(|t' - t_s(plane)|, t')``. Rows of invalid planes are all ``-1``. + """ + candidates = torch.arange(video_length, dtype=torch.float32, device=keyframe_times.device) + slots = _nearest_slots(keyframe_times.to(torch.float32), candidates, None, num_slots) + return torch.where(keyframe_valid[:, None], slots, torch.full_like(slots, -1)) + + +"""DiffVAE tiling helpers: schedule, pad/crop/size-floor, blend utilities. +Decode orchestration lives on ``DiffusionVideoDecoder``. This module owns the +geometry/schedule/mask pieces that tiling uses. +""" + + +ResizeAxisMode = Literal["repeat_last", "symmetric"] + + +_GIB: int = 1 << 30 + + +@dataclass(frozen=True, slots=True) +class _StageFiveBudget: + """One mode's stage-5 multiplicity and withheld reserve, with and without keyframes.""" + + coef: float + coef_keyframes: float + reserve_bytes: int + + +_BUDGET_BY_MODE: dict[DiffVAEMode, _StageFiveBudget] = { + DiffVAEMode.COMBINED_COMPILE: _StageFiveBudget(coef=11, coef_keyframes=15, reserve_bytes=2 * _GIB), + DiffVAEMode.CHUNKED_COMPILE: _StageFiveBudget(coef=7, coef_keyframes=5, reserve_bytes=2 * _GIB), + DiffVAEMode.CHUNKED_EAGER: _StageFiveBudget(coef=5, coef_keyframes=5, reserve_bytes=1 * _GIB), + DiffVAEMode.BLACKWELL_DSL: _StageFiveBudget(coef=2.5, coef_keyframes=2.5, reserve_bytes=2 * _GIB), +} + + +_DEFAULT_ELEMENT_SIZE: int = 2 # bf16 features → fp16 accumulator / bf16 stage-5 + + +_ACCUMULATOR_CHANNELS: int = 3 # RGB pixel blend buffer (decoder out_channels) + + +_MIN_MODEL_BYTES_FLOOR: int = 1 << 30 # never assume a free DiffVAE weight footprint + + +_BUDGET_SAFETY_BYTES_EAGER: int = 1 * _GIB + + +_BUDGET_SAFETY_BYTES_JOINT_MATERIALIZED: int = 2 * _GIB + + +def _falls_back_to_eager_na(mode: DiffVAEMode) -> bool: + """True when this host has no natten and the mode's NA remaps to Triton/eager. + For chunked modes that remap also switches ``compile_blocks`` off, so the peak is the eager + one; :func:`resolve_attention_for_host` is the single owner of that decision. + """ + resolved = resolve_attention_for_host(mode.resolve()) + return resolved.attention in (NAttentionKind.TRITON, NAttentionKind.EAGER_SDPA) and not resolved.compile_blocks + + +def stage5_mem_coef(mode: DiffVAEMode, *, keyframes: bool = False) -> float: + """Stage-5 working-set multiplicity for auto tiling, after host NA resolve. + Args: + mode: the decode preset. + keyframes: whether this decode carries a keyframe stream, which runs eager blocks. + """ + try: + budget = _BUDGET_BY_MODE[mode] + except KeyError as exc: + raise ValueError(f"Unsupported DiffVAEMode for tiling budget: {mode!r}") from exc + if keyframes: + return budget.coef_keyframes + if _falls_back_to_eager_na(mode): + return _BUDGET_BY_MODE[DiffVAEMode.CHUNKED_EAGER].coef + return budget.coef + + +_CONVERT_PEAK_UV_CHANNELS_X2 = 5 # uv_full (2) + pooled uv (0.5) + + +_PACK_PEAK_UV_CHANNELS_X2 = 4 # pooled uv (0.5) + packed float (1.5) + + +_PACK_PEAK_UINT8_BYTES_X2 = 3 # 1.5 B/px + + +def max_emitted_frames(*, num_frames: int, tile_frames: int, overlap_frames: int) -> int: + """Longest chunk the tiled decode yields, in pixel frames. + ``_decode_groups_with_keyframes`` yields only a group's *exclusive* span -- one + stride -- and keeps the trailing overlap as a stub for the next group. Only the + final group yields its whole buffer. Charging ``tile_frames`` for every chunk + therefore roughly doubles the estimate on long clips, which blocks layouts that + would have fit. ``+1`` covers the causal shift, which moves each group after the + first back by one frame. + """ + if tile_frames >= num_frames: + return num_frames + stride = tile_frames - overlap_frames + if stride <= 0: + return tile_frames + n_tiles = 1 + -(-(num_frames - tile_frames) // stride) + last_group = num_frames - (n_tiles - 1) * stride + return min(tile_frames, max(stride, last_group + 1)) + + +def emit_convert_bytes( + *, + tile_frames: int, + height: int, + width: int, + out_channels: int, + element_size: int, +) -> int: + """Downstream bytes a yielded chunk costs while the encoder consumes it. + The decode budget alone is not enough to size a tile: whatever ``_emit`` yields is + handed straight to the video encoder, and that peak overlaps the decode because the + decoder is a generator -- it stays suspended holding stage-4 features and the + accumulator while the consumer converts the chunk it just yielded. A layout that + decodes comfortably can therefore still die in the encoder, and because the + recommender spends spare VRAM on *larger* temporal tiles, more free memory used to + make that failure more likely rather than less. + ``out_channels`` is the width of the intermediate **YUV** tensor, not a second copy + of the emitted RGB: the chunk itself is charged by the accumulator term. On the + write-back path YUV lands in that same RGB storage, so it is not charged here at all + and only the GEMM temporary is. + """ + frames = int(tile_frames) + gemm_frames = min(frames, frames_per_yuv_gemm(height, width)) + # Write-back reuses the RGB storage, so no full YUV tensor survives into the pack. + resident_yuv_frames = frames if gemm_frames == frames else 0 + yuv_channels_x2 = 2 * int(out_channels) + + convert_peak_x2 = element_size * (yuv_channels_x2 * gemm_frames + _CONVERT_PEAK_UV_CHANNELS_X2 * frames) + pack_peak_x2 = ( + element_size * (yuv_channels_x2 * resident_yuv_frames + _PACK_PEAK_UV_CHANNELS_X2 * frames) + + _PACK_PEAK_UINT8_BYTES_X2 * frames + ) + + return int(height) * int(width) * max(convert_peak_x2, pack_peak_x2) // 2 + + +def budget_safety_bytes( + mode: DiffVAEMode, + *, + keyframes: bool = False, + joint_sdpa_materializes: bool = False, +) -> int: + """Extra bytes withheld from the recommend budget. + Args: + mode: the decode preset. + keyframes: whether this decode carries a keyframe stream. + joint_sdpa_materializes: whether the joint attention will run on torch's MATH SDPA kernel + (see ``fallback_na.joint_eager.sdpa_materializes_scores``). Ignored without keyframes, + and irrelevant when a fused joint kernel serves the decode. + """ + try: + budget = _BUDGET_BY_MODE[mode] + except KeyError as exc: + raise ValueError(f"Unsupported DiffVAEMode for tiling budget: {mode!r}") from exc + if keyframes and mode is not DiffVAEMode.BLACKWELL_DSL: + return _BUDGET_SAFETY_BYTES_JOINT_MATERIALIZED if joint_sdpa_materializes else _BUDGET_SAFETY_BYTES_EAGER + if _falls_back_to_eager_na(mode): + return _BUDGET_SAFETY_BYTES_EAGER + return budget.reserve_bytes + + +def accumulator_element_size(feature_dtype: torch.dtype) -> int: + """Bytes per accumulator element; mirrors ``_decode_temporal_group_isolated``. + ``accum_dtype = float16 if feat_s4.dtype == bfloat16 else feat_s4.dtype``. + """ + if feature_dtype is torch.bfloat16: + return 2 # stored as fp16 + return int(torch.tensor([], dtype=feature_dtype).element_size()) + + +def stage4_feature_bytes( + *, + height: int, + width: int, + num_frames: int, + upsample_strides: Sequence[Tuple[int, int, int]], + stage4_channels: int, + element_size: int = _DEFAULT_ELEMENT_SIZE, + natten_trailing_pad_latent_frames: int = 0, +) -> int: + """Resident stages-1-3 output size (full volume tiled into stage 4). + Matches ``DiffusionVideoDecoder.forward_stages_1_to_3`` after optional NATTEN + trailing latent pad: channels-last ``(B, T, H, W, C)`` at stage-4 input resolution. + """ + if stage4_channels < 1: + raise ValueError(f"stage4_channels must be >= 1, got {stage4_channels}") + if element_size < 1: + raise ValueError(f"element_size must be >= 1, got {element_size}") + if len(upsample_strides) < 3: + raise ValueError(f"need at least 3 upsample strides, got {len(upsample_strides)}") + if natten_trailing_pad_latent_frames < 0: + raise ValueError(f"natten_trailing_pad_latent_frames must be >= 0, got {natten_trailing_pad_latent_frames}") + + # Local import: types ↔ tiling cycle avoidance at module import time. + from ltx_core.types import VIDEO_SCALE_FACTORS, VideoLatentShape, VideoPixelShape # noqa: PLC0415 + + latent = VideoLatentShape.from_pixel_shape( + VideoPixelShape(batch=1, frames=int(num_frames), height=int(height), width=int(width), fps=24.0), + scale_factors=VIDEO_SCALE_FACTORS, + ) + s4_t, s4_h, s4_w = stage4_thw_from_latent( + upsample_strides[:3], + latent.frames + int(natten_trailing_pad_latent_frames), + latent.height, + latent.width, + drop_leading_frame=True, + ) + return int(s4_t) * int(s4_h) * int(s4_w) * int(stage4_channels) * int(element_size) + + +def recommended_decode_tiling_config( # noqa: PLR0913 + *, + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], + pixel_scale: SpatioTemporalScaleFactors, + min_tile_size_s4: Tuple[int, int, int], + patch_size: int, + height: int, + width: int, + num_frames: int, + mode: DiffVAEMode, + free_bytes: int, + stage5_channels: int, + stage4_channels: int, + upsample_strides: Sequence[Tuple[int, int, int]], + model_bytes: int = 0, + element_size: int = _DEFAULT_ELEMENT_SIZE, + natten_trailing_pad_latent_frames: int = 0, + out_channels: int = _ACCUMULATOR_CHANNELS, + keyframes: bool = False, + joint_sdpa_materializes: bool = False, +) -> TileSizeConfig: + """Pick DiffVAE decode tiling from stage-4/5 halos and free VRAM. + Always enables both spatial and temporal tiling (temporal-only full-frame slabs + are unsafe on Hopper / some natten builds). + Selection (size-grid, accumulator-aware): + 1. Enumerate legal tile **sizes** on the LCM of DiffVAE ``pixel_scale`` and + :data:`~ltx_core.types.VIDEO_SCALE_FACTORS` (so configs also pass + :class:`~ltx_core.tiling.TileSizeConfig` construction); derive tile + counts from :func:`~ltx_core.tiling.split_by_size` (same as decode). + 2. Drop triples whose peak-bytes estimate exceeds ``usable`` bytes + (``free - max(model, 1 GiB) - safety - stage4_feature``; safety is + 1 GiB eager / 2 GiB compiled, see :func:`budget_safety_bytes`). + Stage-4 input features stay resident for the whole tiled decode. + 3. Among feasible triples, pick minimal :func:`volumetric_overlap_waste`. + Peak-bytes estimate:: + stage4_feature_bytes(...) # hard, full volume + + H * W * (2 * tile_t) * out_channels * element_size + + stage5_tokens * stage5_channels * element_size * coef + Accumulator is full output HxW (not spatially tiled) with temporal extent + ``2 * tile_t``: current group buffer plus the still-live previous exclusive + emit / overlap stub during handoff (not merely ``tile_t + overlap_t``). + RGBx``element_size`` by default. ``element_size`` is the activation width: + production bf16 features use fp16 accumulators (2), matching + :func:`accumulator_element_size`. Stage-5 uses the same element size x + ``stage5_channels`` x ``coef``, which :func:`stage5_mem_coef` reads off the + mode and ``keyframes`` (11 / 7 / 5 / 2.5 by mode; 15 / 5 / 5 / 2.5 with a + keyframe stream, which runs eager blocks). + Args beyond the geometry: + keyframes: this decode carries a keyframe stream (joint attention, eager blocks). + joint_sdpa_materializes: the joint attention will run on torch's MATH SDPA kernel; + costs one more GiB of reserve. See :func:`budget_safety_bytes`. + """ + if height < 1 or width < 1 or num_frames < 1: + raise ValueError(f"height/width/num_frames must be >= 1, got {height}x{width}x{num_frames}") + if patch_size < 1: + raise ValueError(f"patch_size must be >= 1, got {patch_size}") + if stage5_channels < 1: + raise ValueError(f"stage5_channels must be >= 1, got {stage5_channels}") + if out_channels < 1: + raise ValueError(f"out_channels must be >= 1, got {out_channels}") + if element_size < 1: + raise ValueError(f"element_size must be >= 1, got {element_size}") + + overlap_t, overlap_hw = recommended_pixel_overlaps(tile_halos, pixel_scale) + + ft, fh, fw = pixel_scale.time, pixel_scale.height, pixel_scale.width + # Construction validates fixed 8/32/32; to_splitters uses pixel_scale - step both. + step_t = math.lcm(ft, VIDEO_SCALE_FACTORS.time) + step_h = math.lcm(fh, VIDEO_SCALE_FACTORS.height) + step_w = math.lcm(fw, VIDEO_SCALE_FACTORS.width) + min_t_px = _round_up( + # ``2 * overlap`` so left+right ramps fit (else masks are not complementary and + # decode allocates a full weights buffer ≈ another accumulator). + max(2 * ft, 2 * overlap_t, _round_up(min_tile_size_s4[0] * ft, ft), 16), + step_t, + ) + min_h_px = _round_up( + max(2 * fh, 2 * overlap_hw, _round_up(min_tile_size_s4[1] * fh, fh), 64), + step_h, + ) + min_w_px = _round_up( + max(2 * fw, 2 * overlap_hw, _round_up(min_tile_size_s4[2] * fw, fw), 64), + step_w, + ) + + model_cost = max(int(model_bytes), _MIN_MODEL_BYTES_FLOOR) + coef = stage5_mem_coef(mode, keyframes=keyframes) + s4_feat_bytes = stage4_feature_bytes( + height=height, + width=width, + num_frames=num_frames, + upsample_strides=upsample_strides, + stage4_channels=stage4_channels, + element_size=element_size, + natten_trailing_pad_latent_frames=natten_trailing_pad_latent_frames, + ) + reserve = budget_safety_bytes(mode, keyframes=keyframes, joint_sdpa_materializes=joint_sdpa_materializes) + usable = max(0, int(free_bytes) - model_cost - reserve - s4_feat_bytes) + s5_bytes_per_token = max(1.0, float(stage5_channels) * float(element_size) * coef) + acc_bytes_per_pixel = int(out_channels) * int(element_size) + + t_cands = _axis_candidates(num_frames, overlap_t, min_t_px, step_t) + h_cands = _axis_candidates(height, overlap_hw, min_h_px, step_h) + w_cands = _axis_candidates(width, overlap_hw, min_w_px, step_w) + + scored: list[tuple[float, int, int, int, int, int]] = [] + # (waste, -volume, n_t*n_h*n_w, tile_t, tile_h, tile_w) - minimize waste, then launches. + for tile_t, n_t in t_cands: + # Current group buffer + still-live emit/stub during temporal handoff. + acc_frames = 2 * int(tile_t) + acc_bytes = acc_frames * int(height) * int(width) * acc_bytes_per_pixel + # The consumer converts each yielded chunk while this decode is suspended. + downstream_bytes = emit_convert_bytes( + tile_frames=max_emitted_frames(num_frames=num_frames, tile_frames=tile_t, overlap_frames=overlap_t), + height=height, + width=width, + out_channels=out_channels, + element_size=element_size, + ) + if acc_bytes + downstream_bytes >= usable: + continue + s5_budget_bytes = usable - acc_bytes - downstream_bytes + max_s5_tokens = int(s5_budget_bytes // s5_bytes_per_token) + for tile_h, n_h in h_cands: + for tile_w, n_w in w_cands: + if stage5_tokens_for_pixel_tile(tile_t, tile_h, tile_w, patch_size=patch_size) > max_s5_tokens: + continue + waste = volumetric_overlap_waste( + num_frames=num_frames, + height=height, + width=width, + tile_frames=tile_t, + tile_height=tile_h, + tile_width=tile_w, + n_t=n_t, + n_h=n_h, + n_w=n_w, + ) + scored.append((waste, -tile_t * tile_h * tile_w, n_t * n_h * n_w, tile_t, tile_h, tile_w)) + + if not scored: + raise ValueError( + "Cannot fit a DiffVAE decode tile under the memory budget: " + f"min tile ~{min_t_px}f x {min_h_px}x{min_w_px}px " + f"(overlaps T={overlap_t}, HW={overlap_hw}), " + f"mode={mode.value}, keyframes={keyframes}, coef={coef}, stage5_channels={stage5_channels}, " + f"stage4_feature_bytes={s4_feat_bytes}, usable_bytes={usable}. " + "Reduce resolution, reduce num_frames (stage-4 features and the per-chunk " + "encode buffers both scale with it), or free GPU memory." + ) + + scored.sort() + _waste, _vol, _ntiles, tile_t, tile_h, tile_w = scored[0] + return TileSizeConfig( + frames=DimensionSizeConfig(tile_size=tile_t, overlap=overlap_t), + height=DimensionSizeConfig(tile_size=tile_h, overlap=overlap_hw), + width=DimensionSizeConfig(tile_size=tile_w, overlap=overlap_hw), + ) + + +def prepare_tile_schedule( + stage4_shape_bcthw: torch.Size, + tiling_config: TilingConfig | None, + *, + upsample3_stride: Tuple[int, int, int], + patch_size: int, + min_tile_size: Tuple[int, int, int], + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], +) -> List[Tile]: + """Build pixel-blend tiles whose ``in_coords`` land on the stage-4 input grid. + DiffVAE temporal tiling deliberately skips ConvVAE causal split/mask tricks + (``split_temporal_causal``, ``left_starts_from_0``): pixel overlap already + covers blend+halo, and interval propagation follows + :class:`~ltx_core.model.video_vae.transformer.layers.LinearPixelShuffleUpsample` + (``drop_leading_frame`` only on the origin tile) with *symmetric* trapezoid + ramps so masks stay complementary without a weight buffer. + """ + pixel_scale = stage4_to_pixel_scale_factors(upsample3_stride, patch_size) + if tiling_config is None: + return [ + Tile( + in_coords=(slice(None), slice(None), slice(None), slice(None), slice(None)), + out_coords=(slice(None), slice(None), slice(None), slice(None), slice(None)), + masks_1d=( + untiled_mask_1d(), + untiled_mask_1d(), + untiled_mask_1d(), + untiled_mask_1d(), + untiled_mask_1d(), + ), + ) + ] + + overlap_t, overlap_hw = recommended_pixel_overlaps(tile_halos, pixel_scale) + _validate_overlap(tiling_config, min_overlap_frames=overlap_t, min_overlap_pixels=overlap_hw) + # Plain split (not split_temporal_causal): no start-1 / left_ramp+1 copycat of ConvVAE. + t_split, h_split, w_split = tiling_config.to_splitters( + pixel_scale, min_tile_size=min_tile_size, causal_temporal=False + ) + st, sh, sw = upsample3_stride + + def axis_specs( + split_op: SplitOperation, + dim_len: int, + stride_component: int, + *, + propagate_causal: bool, + apply_patch: bool, + ) -> list[tuple[slice, slice, torch.Tensor]]: + if split_op is DEFAULT_SPLIT_OPERATION: + return [(slice(None), slice(None), untiled_mask_1d())] + intervals = split_op(dim_len).intervals + specs = [] + for iv in intervals: + stage5 = _propagate_interval_through_upsample_hops(iv, [stride_component], propagate_causal) + if apply_patch: + pixel = _propagate_interval_through_upsample_hops(stage5, [patch_size], causal=False) + else: + pixel = stage5 + # Symmetric ramps (left_starts_from_0=False) for partition-of-unity with + # pixel-shuffle out_coords; ConvVAE sacrificial first-sample is not used. + mask_pixel = compute_trapezoidal_mask_1d( + pixel.end - pixel.start, pixel.left_ramp, pixel.right_ramp, left_starts_from_0=False + ) + specs.append((slice(iv.start, iv.end), slice(pixel.start, pixel.end), mask_pixel)) + return specs + + # Temporal: pixel-shuffle propagate (drop-leading geometry); spatial: exact x stride. + t_specs = axis_specs(t_split, stage4_shape_bcthw[2], st, propagate_causal=True, apply_patch=False) + h_specs = axis_specs(h_split, stage4_shape_bcthw[3], sh, propagate_causal=False, apply_patch=True) + w_specs = axis_specs(w_split, stage4_shape_bcthw[4], sw, propagate_causal=False, apply_patch=True) + + tiles: List[Tile] = [] + for t_spec, h_spec, w_spec in itertools.product(t_specs, h_specs, w_specs): + t_s4, t_px, t_mask = t_spec + h_s4, h_px, h_mask = h_spec + w_s4, w_px, w_mask = w_spec + tiles.append( + Tile( + in_coords=(slice(None), t_s4, h_s4, w_s4, slice(None)), + out_coords=(slice(None), slice(None), t_px, h_px, w_px), + masks_1d=(untiled_mask_1d(), untiled_mask_1d(), t_mask, h_mask, w_mask), + ) + ) + return tiles + + +def slice_stage4_tile( + feat_s4: torch.Tensor, + tile: Tile, + *, + content_frames: int, +) -> tuple[torch.Tensor, bool, bool, tuple[int, int, int]]: + """Slice a stage-4 feature tile, extending trailing tiles to include ghost frames.""" + is_origin = tile.in_coords[1].start in (0, None) + _, stop, _ = tile.in_coords[1].indices(content_frames) + pad_trailing = stop == content_frames + _b, t_coord, h_coord, w_coord, _c = tile.in_coords + t0, t1, _ = t_coord.indices(content_frames) + h0, h1, _ = h_coord.indices(feat_s4.shape[2]) + w0, w1, _ = w_coord.indices(feat_s4.shape[3]) + content_thw = (t1 - t0, h1 - h0, w1 - w0) + if pad_trailing: + t1 = feat_s4.shape[1] + feat_tile = feat_s4[:, t0:t1, h_coord, w_coord, :] + return feat_tile, is_origin, pad_trailing, content_thw + + +@dataclass(frozen=True) +class AxisPad: + """How many elements were added (pad) or removed (crop) on each side of one axis.""" + + before: int + after: int + + +def resize_axis( + x: torch.Tensor, + dim: int, + size: int, + *, + mode: ResizeAxisMode, +) -> tuple[torch.Tensor, AxisPad]: + """Pad or crop axis ``dim`` so its length becomes ``size``. + Pad (``len < size``): + ``repeat_last`` - append copies of the last slice. + ``symmetric`` - edge-replicate first/last; leftover goes to the end + (``before = need // 2``, ``after = need - before``). + Crop (``len > size``): + ``repeat_last`` - drop from the end. + ``symmetric`` - drop from both ends with the same split rule as pad. + """ + if size < 1: + raise ValueError(f"resize_axis target size must be >= 1, got {size}") + if dim < 0: + dim += x.ndim + if not 0 <= dim < x.ndim: + raise ValueError(f"dim {dim} out of range for rank-{x.ndim} tensor") + + length = x.shape[dim] + if length == size: + return x, AxisPad(0, 0) + + if length < size: + need = size - length + if mode == "repeat_last": + last = x.narrow(dim, length - 1, 1) + expand_shape = list(x.shape) + expand_shape[dim] = need + pad = last.expand(expand_shape) + return torch.cat([x, pad], dim=dim), AxisPad(0, need) + + before = need // 2 + after = need - before + first = x.narrow(dim, 0, 1) + last = x.narrow(dim, length - 1, 1) + parts: list[torch.Tensor] = [] + if before: + expand_shape = list(x.shape) + expand_shape[dim] = before + parts.append(first.expand(expand_shape)) + parts.append(x) + if after: + expand_shape = list(x.shape) + expand_shape[dim] = after + parts.append(last.expand(expand_shape)) + return torch.cat(parts, dim=dim), AxisPad(before, after) + + need = length - size + if mode == "repeat_last": + return x.narrow(dim, 0, size).contiguous(), AxisPad(0, need) + + before = need // 2 + after = need - before + return x.narrow(dim, before, size).contiguous(), AxisPad(before, after) + + +def ensure_min_latent_shape( + latent: torch.Tensor, + min_tile_sizes: Tuple[int, int, int], +) -> tuple[torch.Tensor, tuple[AxisPad, AxisPad, AxisPad]]: + """Pad latent ``(B, C, T, H, W)`` up to ``min_tile_sizes`` if needed.""" + min_t, min_h, min_w = min_tile_sizes + t_pad = AxisPad(0, 0) + h_pad = AxisPad(0, 0) + w_pad = AxisPad(0, 0) + x = latent + if x.shape[2] < min_t: + x, t_pad = resize_axis(x, 2, min_t, mode="repeat_last") + if x.shape[3] < min_h: + x, h_pad = resize_axis(x, 3, min_h, mode="symmetric") + if x.shape[4] < min_w: + x, w_pad = resize_axis(x, 4, min_w, mode="symmetric") + return x, (t_pad, h_pad, w_pad) + + +def scale_axis_pad(pad: AxisPad, scale: int) -> AxisPad: + """Scale a latent-grid ``AxisPad`` into pixel (or other) units.""" + return AxisPad(pad.before * scale, pad.after * scale) + + +def crop_pixels_to_content( + pixels: torch.Tensor, + frames: int, + height: int, + width: int, + *, + h_pad: AxisPad | None = None, + w_pad: AxisPad | None = None, + spatial_scale: Tuple[int, int] = (1, 1), +) -> torch.Tensor: + """Crop padded decode output ``(B, C, F, H, W)`` back to the content shape. + Temporal pad is always trailing (``repeat_last``), so T is cropped from the + end. Spatial size-floor pads must pass the recorded ``h_pad`` / ``w_pad`` + (latent units) plus ``spatial_scale`` ``(H, W)`` so odd leftovers are not + re-split by a center-crop after upscaling. + """ + x, _ = resize_axis(pixels, 2, frames, mode="repeat_last") + scale_h, scale_w = spatial_scale + if h_pad is not None: + before = scale_axis_pad(h_pad, scale_h).before + if before + height > x.shape[3]: + raise ValueError(f"H crop out of range: before={before}, height={height}, got {x.shape[3]}") + x = x.narrow(3, before, height).contiguous() + else: + x, _ = resize_axis(x, 3, height, mode="symmetric") + if w_pad is not None: + before = scale_axis_pad(w_pad, scale_w).before + if before + width > x.shape[4]: + raise ValueError(f"W crop out of range: before={before}, width={width}, got {x.shape[4]}") + x = x.narrow(4, before, width).contiguous() + else: + x, _ = resize_axis(x, 4, width, mode="symmetric") + return x + + +def stage5_pixel_shape_from_stage4( + stage4_t: int, + stage4_h: int, + stage4_w: int, + *, + upsample_stride: Tuple[int, int, int], + patch_size: int, + stage5_kernel_t: int, + drop_leading_frame: bool, + pad_trailing: bool, +) -> tuple[int, int, int]: + """Pixel ``(F, H, W)`` for a stage-4-input extent (one remaining NA hop + patch).""" + st, sh, sw = upsample_stride + frames = stage4_t * st - 1 if drop_leading_frame and st == 2 else stage4_t * st + if pad_trailing: + frames = max(frames, stage5_kernel_t) + return frames, stage4_h * sh * patch_size, stage4_w * sw * patch_size + + +def pad_trailing_latent_for_natten_border(latent: torch.Tensor, n_frames: int) -> torch.Tensor: + """Replicate the last latent frame ``n_frames`` times for NATTEN last-frame border.""" + if n_frames <= 0: + return latent + padded, _ = resize_axis(latent, 2, latent.shape[2] + n_frames, mode="repeat_last") + return padded + + +def crop_trailing_context_natten_pad( + context: torch.Tensor, + *, + n_latent_frames: int, + time_scale: int, + stage5_kernel_t: int, +) -> torch.Tensor: + """Crop ghosting appendix before stage 5, leaving at least ``stage5_kernel_t``.""" + if n_latent_frames <= 0: + return context + ghost = n_latent_frames * time_scale + content_t = max(context.shape[1] - ghost, 1) + keep = min(context.shape[1], max(content_t, stage5_kernel_t)) + cropped, _ = resize_axis(context, 1, keep, mode="repeat_last") + return cropped + + +def _weight_floor(dtype: torch.dtype) -> float: + """Smallest divisor that safely guards ``buffer / weights`` in ``dtype``.""" + return max(1e-8, torch.finfo(dtype).tiny) + + +def stage4_thw_from_latent( + upsample_strides: Sequence[Tuple[int, int, int]], + latent_t: int, + latent_h: int, + latent_w: int, + *, + drop_leading_frame: bool = True, +) -> Tuple[int, int, int]: + """Stage-4 input ``(T, H, W)`` after the first three upsample hops.""" + t, h, w = latent_t, latent_h, latent_w + for st, sh, sw in upsample_strides[:3]: + t, h, w = t * st, h * sh, w * sw + if st == 2 and drop_leading_frame: + t -= 1 + return t, h, w + + +def stage4_to_pixel_scale_factors( + upsample_stride: Tuple[int, int, int], + patch_size: int, +) -> SpatioTemporalScaleFactors: + """Pixel/frame units per stage-4-input cell (last NA hop + unpatchify).""" + st, sh, sw = upsample_stride + return SpatioTemporalScaleFactors(time=st, height=sh * patch_size, width=sw * patch_size) + + +def compute_tile_min_size( + stage4_kernel: Tuple[int, int, int], + stage5_kernel: Tuple[int, int, int], + upsample3_stride: Tuple[int, int, int], +) -> Tuple[int, int, int]: + """Min stage-4-input ``(T, H, W)`` so stages 4 and 5 each see ``>= kernel``.""" + return tuple(max(stage4_kernel[a], -(-stage5_kernel[a] // upsample3_stride[a])) for a in range(3)) + + +def compute_tile_halos( + stage4_kernel: Tuple[int, int, int], + stage4_depth: int, + stage5_kernel: Tuple[int, int, int], + stage5_depth: int, + upsample3_stride: Tuple[int, int, int], +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]: + """One-sided halos in stage-4-input units for stages 4 and 5.""" + halo4 = tuple(stage4_depth * (stage4_kernel[a] // 2) for a in range(3)) + halo5 = tuple(-(-(stage5_depth * (stage5_kernel[a] // 2)) // upsample3_stride[a]) for a in range(3)) + return halo4, halo5 # type: ignore[return-value] + + +def _cumulative_upsample_strides( + upsamples: Sequence[Tuple[Tuple[int, int, int], int]], +) -> List[Tuple[int, int, int]]: + """Per-axis product of hop strides for ``upsamples[:i]`` (``cumulative[0] = (1,1,1)``).""" + cumulative = [(1, 1, 1)] + t, h, w = 1, 1, 1 + for stride, _ in upsamples: + t, h, w = t * stride[0], h * stride[1], w * stride[2] + cumulative.append((t, h, w)) + return cumulative + + +def all_stages_min_tile_size( + stage_kernels: Sequence[Tuple[int, int, int]], + upsamples: Sequence[Tuple[Tuple[int, int, int], int]], + stage5_kernel: Tuple[int, int, int], +) -> Tuple[int, int, int]: + """Per-axis latent-grid floor so every stage's NA sees dims ``>= kernel_size``.""" + cumulative = _cumulative_upsample_strides(upsamples) + mins = [1, 1, 1] + for stage_i in range(len(upsamples)): + strides = cumulative[stage_i] + for axis in range(3): + mins[axis] = max(mins[axis], -(-stage_kernels[stage_i][axis] // strides[axis])) + strides5 = cumulative[len(upsamples)] + for axis in range(3): + mins[axis] = max(mins[axis], -(-stage5_kernel[axis] // strides5[axis])) + return (mins[0], mins[1], mins[2]) + + +def pixel_tile_shape(full_shape: tuple[int, ...], out_coords: tuple[slice, ...]) -> tuple[int, ...]: + dims: list[int] = [] + for size, coord in zip(full_shape, out_coords, strict=True): + start, stop, step = coord.indices(size) + dims.append(len(range(start, stop, step))) + return tuple(dims) + + +def _round_up(value: int, multiple: int) -> int: + return -(-value // multiple) * multiple + + +def recommended_pixel_overlaps( + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], + pixel_scale: SpatioTemporalScaleFactors, +) -> Tuple[int, int]: + """Stage-4/5-safe ``(temporal_overlap_frames, spatial_overlap_pixels)``. + Shared by :func:`recommended_decode_tiling_config` (to *set* overlaps) and + :func:`~ltx_core.tiling._validate_overlap` (to reject undersized configs). + """ + + def dominant(axis: int) -> int: + return max(tile_halos[i][axis] for i in range(len(tile_halos))) + + overlap_t = _round_up(dominant(0) * pixel_scale.time, 8) + halo_hw = max(dominant(1), dominant(2)) + overlap_hw = _round_up(halo_hw * pixel_scale.height, 32) + return overlap_t, overlap_hw + + +def stage5_tokens_for_pixel_tile( + tile_frames: int, + tile_height: int, + tile_width: int, + *, + patch_size: int, +) -> int: + """Pre-unpatchify stage-5 token count for a pixel-space tile (NATTEN volume).""" + h5 = max(1, tile_height // patch_size) + w5 = max(1, tile_width // patch_size) + return tile_frames * h5 * w5 + + +def _axis_candidates(length: int, overlap: int, min_size: int, multiple: int) -> list[tuple[int, int]]: + """``(tile_size, num_tiles)`` for every legal size on ``multiple``'s grid.""" + out: list[tuple[int, int]] = [] + max_size = max(_round_up(length, multiple), min_size) + for size in range(min_size, max_size + multiple, multiple): + if size <= overlap: + continue + n = len(split_by_size(size, overlap)(length).intervals) + out.append((size, n)) + return out + + +def volumetric_overlap_waste( + *, + num_frames: int, + height: int, + width: int, + tile_frames: int, + tile_height: int, + tile_width: int, + n_t: int, + n_h: int, + n_w: int, +) -> float: + """``processed_volume / unique_volume`` (>= 1). Lower means less overlap recompute.""" + processed = n_t * n_h * n_w * tile_frames * tile_height * tile_width + unique = max(1, num_frames * height * width) + return processed / unique + + +def _propagate_interval_through_upsample_hops( + interval: DimensionInterval, + strides: Sequence[int], + causal: bool, +) -> DimensionInterval: + """Forward-propagate one interval through a sequence of upsample hops on one axis. + Mirrors :class:`~ltx_core.model.video_vae.transformer.layers.LinearPixelShuffleUpsample`: + multiply by ``stride``, and for the causal temporal axis when ``stride == 2`` apply + the duplicate-frame drop (``end -= 1``; non-origin also ``start -= 1``). + This is *not* :func:`~ltx_core.model.video_vae.video_vae.map_temporal_slice` (ConvVAE). + DiffVAE non-origin tiles run with ``drop_leading_frame=False`` and must keep length + ``tile_t * stride``; the ConvVAE ``1+(L-1)*stride`` mapping is one frame short and + shifts non-origin ``out_coords``, which breaks tiled↔untiled temporal blend even + when masks are complementary. + """ + x = interval + for stride in strides: + if stride < 1: + raise ValueError(f"upsample stride must be >= 1, got {stride}") + start = x.start * stride + end = x.end * stride + left_ramp = x.left_ramp * stride + right_ramp = x.right_ramp * stride + if causal and stride == 2: + end -= 1 + if x.start != 0: + start -= 1 + x = DimensionInterval(start=start, end=end, left_ramp=left_ramp, right_ramp=right_ramp) + return x + + +"""Shared small layers for the diffusion-VAE NA transformer stack.""" + + +class ChannelLinear(nn.Linear): + """``nn.Linear`` exposing ``in_channels``/``out_channels`` for config introspection.""" + + @property + def in_channels(self) -> int: + return self.in_features + + @property + def out_channels(self) -> int: + return self.out_features + + +class LinearPixelShuffleUpsample(nn.Module): + """Decoder-side resampler: Linear channel-expand, then channels-last PixelShuffle.""" + + def __init__( + self, + in_channels: int, + stride: tuple[int, int, int], + out_channels_reduction_factor: int = 1, + ) -> None: + super().__init__() + self.stride = stride + self.proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor + self.out_channels = self.proj_out_channels // math.prod(stride) + self.proj = nn.Linear(in_channels, self.proj_out_channels, bias=True) + + def forward(self, x: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: + """Upsample; when ``stride[0] == 2`` the pixel-shuffle produces a duplicate + leading frame that must be dropped to preserve the causal 1:2 (then + composed 1:8) frame mapping. ``drop_leading_frame`` gates that drop: it + must be ``True`` only for the chunk that contains the tensor's true + temporal origin (t=0). Tiled callers processing a later chunk in + isolation must pass ``False`` -- that chunk has no duplicate leading + frame of its own to drop, since the one duplicate frame in the full + (untiled) tensor belongs solely to the origin chunk. + """ + x = self.proj(x) + x = rearrange( + x, + "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + if self.stride[0] == 2 and drop_leading_frame: + x = x[:, 1:, :, :, :] + return x + + +class AdaLNZero(nn.Module): + """Per-block AdaLN-Zero modulation: ``t_emb`` -> 7 (scale/shift/gate) chunks. + Zero-init output projection so the block is an identity at every timestep + until the modulation pathway opens up during training. + """ + + NUM_CHUNKS: int = 7 # scale_msa, shift_msa, gate_msa, scale_mlp, shift_mlp, gate_mlp, gate_ctx + + def __init__(self, dim: int, t_emb_dim: int) -> None: + super().__init__() + self.dim = dim + self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True) + nn.init.zeros_(self.proj.weight) + nn.init.zeros_(self.proj.bias) + + def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: + h = self.proj(F.silu(t_emb)) + chunks = h.chunk(self.NUM_CHUNKS, dim=-1) + return tuple(c[:, None, None, None, :] for c in chunks) + + +def modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + """Apply AdaLN-style scale + shift modulation to a channels-last tensor.""" + return x * (1.0 + scale) + shift + + +"""Shared absolute-RoPE math helpers (no consumer policy).""" + + +DEFAULT_ABS_ROPE_NUM_TILES = 4 + + +def t_positions(t: int, device: torch.device) -> torch.Tensor: + return torch.arange(t, dtype=torch.float32, device=device) + + +def h_positions(h: int, device: torch.device) -> torch.Tensor: + return torch.arange(h, dtype=torch.float32, device=device) + + +def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: + """Default split of head_dim across (T, H, W) RoPE chunks.""" + assert head_dim % 8 == 0, f"head_dim={head_dim} must be a multiple of 8 for default split" + d_t = (head_dim // 4) // 2 * 2 + d_hw = (head_dim - d_t) // 2 + if d_hw % 2 != 0: + d_t -= 2 + d_hw = (head_dim - d_t) // 2 + assert d_t > 0 + assert d_hw > 0 + return (d_t, d_hw, d_hw) + + +def rope_inv_freqs(dim: int, base: float = 10000.0) -> torch.Tensor: + """Inverse RoPE frequencies: ``1 / base**(i/dim)`` for ``i`` in ``[0, dim, 2)``.""" + assert dim % 2 == 0, f"RoPE dim must be even, got {dim}" + exponents = np.arange(0, dim, 2, dtype=np.float64) / dim + inv_freqs = 1.0 / np.power(float(base), exponents) + return torch.from_numpy(inv_freqs).to(torch.float32) + + +def rot_abs_axis_impl( + xc: torch.Tensor, + pos: torch.Tensor, + inv: torch.Tensor, + axis: int, + *, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Absolute RoPE on one axis chunk ``xc[..., D]`` (D even) → new tensor.""" + out_dtype = xc.dtype + pairs = xc.reshape(*xc.shape[:-1], xc.shape[-1] // 2, 2) + xe = pairs[..., 0].to(compute_dtype) + xo = pairs[..., 1].to(compute_dtype) + shape = [1, 1, 1, 1, 1, inv.shape[0]] + shape[axis] = pos.shape[0] + ang = (pos[:, None] * inv[None, :]).reshape(shape) + c = ang.cos().to(compute_dtype) + s = ang.sin().to(compute_dtype) + re = xe * c - xo * s + ro = xe * s + xo * c + out = torch.stack([re, ro], dim=-1).reshape(xc.shape) + return out.to(out_dtype) if out.dtype != out_dtype else out + + +"""Opaque full-volume abs-RoPE for deterministic (pre-diffusion) NA. +Owns QKV + opaque ``custom_op`` packaging for det ``NA.forward``. Det stages +have differently shaped T/H/W; the opaque op keeps Dynamo from specializing +on each stage shape. Diffusion paths use ``diff_attn/`` RoPE — not this file. +No T/H/W origin/offset is threaded for tiled decode, and none is needed: +every attention call here is ``natten.na3d``, a local window with no +cross-tile tokens. Absolute-vs-local RoPE differs by a global phase that +cancels inside the attention softmax over that window, so the attention +output is unchanged. Since every tiled-decode call processes exactly one +tile in isolation, using each tile's local 0-based positions is identical +to using its true absolute origin. Absolute origin still matters for +whether a tile contains the latent's first frame (``drop_leading_frame``), +which is handled outside RoPE in the decoder stage / per-tile decode. +""" + + +def _apply_opaque_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + """Rotate one W-extent with raw abs-RoPE (runs inside the opaque op). + ``t_pos`` overrides the default integer ``arange`` on the first axis; the keyframe + stream passes its (possibly fractional) plane times there. + """ + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + positions_t = t_positions(t, x.device) if t_pos is None else t_pos + xt = rot_abs_axis_impl(x[..., :d_t], positions_t, inv_t, axis=1, compute_dtype=compute_dtype) + xh = rot_abs_axis_impl( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _apply_opaque_tiled_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + """Fixed-``num_tiles`` W split + per-slab rotation (body of the opaque op). + The keyframe stream shares the video stream's W extent at every stage, so passing the + same ``num_tiles`` yields identical slab boundaries and identical ``w_pos`` -- which + is what keeps the two streams' W phases comparable inside the joint softmax. + """ + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + _apply_opaque_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + t_pos=t_pos, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +@torch.library.custom_op("diffsynth_ltx25::abs_rope", mutates_args=()) +def _abs_rope_op( + x: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + """Opaque out-of-place abs-RoPE: Dynamo sees one node.""" + compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 + return _apply_opaque_tiled_rope( + x, + (d_t, d_h, d_w), + (inv_t, inv_h, inv_w), + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + + +@_abs_rope_op.register_fake +def _abs_rope_fake( + x: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + del inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 + return torch.empty(x.shape, device=x.device, dtype=x.dtype) + + +def _apply_opaque_abs_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"compute_dtype must be float32 or bfloat16, got {compute_dtype}") + d_t, d_h, d_w = rope_split + inv_t, inv_h, inv_w = inv_freqs + return _abs_rope_op( + x, + inv_t, + inv_h, + inv_w, + d_t, + d_h, + d_w, + num_tiles, + compute_dtype == torch.bfloat16, + ) + + +@torch.library.custom_op("diffsynth_ltx25::abs_rope_at_t", mutates_args=()) +def _abs_rope_at_t_op( + x: torch.Tensor, + t_pos: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + """Opaque abs-RoPE with caller-supplied (possibly fractional) first-axis positions.""" + compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 + return _apply_opaque_tiled_rope( + x, + (d_t, d_h, d_w), + (inv_t, inv_h, inv_w), + num_tiles=num_tiles, + compute_dtype=compute_dtype, + t_pos=t_pos, + ) + + +@_abs_rope_at_t_op.register_fake +def _abs_rope_at_t_fake( + x: torch.Tensor, + t_pos: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + del t_pos, inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 + return torch.empty(x.shape, device=x.device, dtype=x.dtype) + + +def _rope_config(attn: object, x: torch.Tensor) -> tuple[tuple[torch.Tensor, ...], int, torch.dtype]: + """``(inv_freqs, num_tiles, compute_dtype)`` read off an attention module.""" + inv_freqs = ( + attn.rope_inv_t.to(device=x.device), # type: ignore[attr-defined] + attn.rope_inv_h.to(device=x.device), # type: ignore[attr-defined] + attn.rope_inv_w.to(device=x.device), # type: ignore[attr-defined] + ) + num_tiles = getattr(attn, "rope_num_tiles", DEFAULT_ABS_ROPE_NUM_TILES) + compute_dtype = getattr(attn, "rope_compute_dtype", torch.float32) + return inv_freqs, num_tiles, compute_dtype + + +def _det_project_qkv(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Shared Q/K/V proj + norm + scale for :func:`det_qkv_rope` and ``_at_times``.""" + q, k, v = attn.project_qkv(x) # type: ignore[attr-defined] + q = attn.q_norm(q) # type: ignore[attr-defined] + k = attn.k_norm(k) # type: ignore[attr-defined] + q = q * attn.scale # type: ignore[attr-defined] + return q, k, v + + +def det_qkv_rope_at_times( + attn: object, + x: torch.Tensor, + t_pos: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V proj + norm/scale + abs-RoPE with explicit first-axis positions. + The keyframe-stream counterpart of :func:`det_qkv_rope`. ``t_pos`` is ``(P,)`` float + stage times *in the same origin* the video stream's RoPE uses, so the joint softmax + sees true relative offsets: absolute-vs-local RoPE only cancels as a global phase + when every token in the softmax shares one origin, which no longer holds once + keyframe tokens join a video window. + """ + if t_pos.ndim != 1 or t_pos.shape[0] != x.shape[1]: + raise ValueError(f"t_pos must be ({x.shape[1]},) to match the plane axis, got {tuple(t_pos.shape)}") + q, k, v = _det_project_qkv(attn, x) + + inv_freqs, num_tiles, compute_dtype = _rope_config(attn, x) + d_t, d_h, d_w = attn.rope_dim_split # type: ignore[attr-defined] + positions = t_pos.to(device=x.device, dtype=torch.float32) + rotated = [ + _abs_rope_at_t_op( + tensor, + positions, + *inv_freqs, + d_t, + d_h, + d_w, + num_tiles, + compute_dtype == torch.bfloat16, + ) + for tensor in (q, k) + ] + return rotated[0], rotated[1], v + + +def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V proj + norm/scale + opaque full-volume abs-RoPE.""" + q, k, v = _det_project_qkv(attn, x) + inv_freqs, num_tiles, compute_dtype = _rope_config(attn, x) + q = _apply_opaque_abs_rope( + q, + attn.rope_dim_split, # type: ignore[attr-defined] + inv_freqs, + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + k = _apply_opaque_abs_rope( + k, + attn.rope_dim_split, # type: ignore[attr-defined] + inv_freqs, + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + return q, k, v + + + +class QKVProjections(nn.Module): + """Checkpoint-fused QKV weights executed as the target's three projections.""" + + def __init__(self, dim: int) -> None: + super().__init__() + linear = nn.Linear(dim, dim * 3, bias=True) + self.weight = linear.weight + self.bias = linear.bias + + def forward(self, x): + weights = self.weight.chunk(3, dim=0) + biases = self.bias.chunk(3, dim=0) + return tuple(F.linear(x, weight, bias) for weight, bias in zip(weights, biases, strict=True)) + + +DEFAULT_SWIGLU_TILE_SIZE: Final[int] = 16_384 +DEFAULT_SWIGLU_TILES: Final[int] = 4 + + +@dataclass(frozen=True) +class SwiGLUTileSpec: + num_tiles: int | None = None + tile_size: int | None = None + + def __post_init__(self) -> None: + if (self.num_tiles is None) == (self.tile_size is None): + raise ValueError("Provide exactly one of num_tiles or tile_size") + if self.num_tiles is not None and self.num_tiles < 1: + raise ValueError("num_tiles must be >= 1") + if self.tile_size is not None and self.tile_size < 1: + raise ValueError("tile_size must be >= 1") + + @classmethod + def by_count(cls, num_tiles: int = DEFAULT_SWIGLU_TILES): + return cls(num_tiles=num_tiles) + + @classmethod + def by_size(cls, tile_size: int = DEFAULT_SWIGLU_TILE_SIZE): + return cls(tile_size=tile_size) + + +DEFAULT_SWIGLU_TILE_SPEC = SwiGLUTileSpec(tile_size=DEFAULT_SWIGLU_TILE_SIZE) + + +def _swiglu_chunk(x, w_gate, w_up, w_down): + return F.linear(F.silu(F.linear(x, w_gate)) * F.linear(x, w_up), w_down) + + +def swiglu_tiled(x, w_gate, w_up, w_down, tile, *, use_triton=None): + del use_triton + if x.numel() == 0: + return x + leading = x.shape[:-1] + flat = x.reshape(-1, x.shape[-1]) + if tile.tile_size is not None: + chunk_size = tile.tile_size + else: + chunk_size = max(1, math.ceil(flat.shape[0] / tile.num_tiles)) + output = torch.cat([_swiglu_chunk(chunk, w_gate, w_up, w_down) for chunk in flat.split(chunk_size)], dim=0) + return output.reshape(*leading, output.shape[-1]) + + +def plain_mlp(x, mlp, norm, tile): + y = norm(x) + if y.numel() == 0: + return x + return x + swiglu_tiled(y, mlp.w_gate.weight, mlp.w_up.weight, mlp.w_down.weight, tile) + + +class SwiGLU(nn.Module): + def __init__(self, dim: int, hidden_dim: int, tile: SwiGLUTileSpec = DEFAULT_SWIGLU_TILE_SPEC) -> None: + super().__init__() + self.w_up = nn.Linear(dim, hidden_dim, bias=False) + self.w_gate = nn.Linear(dim, hidden_dim, bias=False) + self.w_down = nn.Linear(hidden_dim, dim, bias=False) + self.tile = tile + + def forward(self, x): + return swiglu_tiled(x, self.w_gate.weight, self.w_up.weight, self.w_down.weight, self.tile) + + +def configure_swiglu_tile(module_root, *, num_tiles=None, tile_size=None): + if num_tiles is None and tile_size is None: + return + tile = SwiGLUTileSpec(num_tiles=num_tiles, tile_size=tile_size) + for module in module_root.modules(): + if isinstance(module, SwiGLU): + module.tile = tile + + + +"""Limited-workspace 3D neighborhood attention (NATTEN ``na3d`` semantics) in pure torch. +Vendored from comfy-kitchen ``backends/eager/na.py`` (Apache-2.0) for DiffVAE hosts +without natten or Triton. Queries are tiled; tiles that share window geometry stack +into batched ``scaled_dot_product_attention`` calls with one additive mask per group. +""" + + +NA_SCORE_BUDGET = 2**25 + + +NA_KV_STACK_BUDGET = 2**28 + + +def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], list[int]]: + """Per-index (start, end) of the attended window along one axis.""" + starts: list[int] = [] + ends: list[int] = [] + if causal: + for i in range(length): + starts.append(max(0, i - kernel + 1)) + ends.append(i + 1) + else: + kernel = min(kernel, length) + lo = length - kernel + half = kernel // 2 + for i in range(length): + start = min(max(i - half, 0), lo) + starts.append(start) + ends.append(start + kernel) + return starts, ends + + +def _pick_tiles(dims: tuple[int, int, int], kernels: list[int]) -> list[int]: + """Per-axis query-tile lengths keeping one tile's [Nq, Nk] under budget.""" + tiles = list(dims) + + def cost(ts: list[int]) -> int: + nq = math.prod(ts) + nk = math.prod(min(d, t + k - 1) for t, k, d in zip(ts, kernels, dims, strict=True)) + return nq * nk + + while cost(tiles) > NA_SCORE_BUDGET and max(tiles) > 1: + i = max(range(3), key=lambda a: tiles[a] / kernels[a]) + if tiles[i] <= 1: + break + tiles[i] = max(1, (tiles[i] + 1) // 2) + return tiles + + +def _group_mask( + rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group.""" + bools = [] + for starts, ends in rel_bounds: + st = torch.tensor(starts, device=device) + en = torch.tensor(ends, device=device) + kj = torch.arange(int(en.max()), device=device) + bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) + visible = ( + bools[0][:, None, None, :, None, None] + & bools[1][None, :, None, None, :, None] + & bools[2][None, None, :, None, None, :] + ) + nq = visible.shape[0] * visible.shape[1] * visible.shape[2] + nk = visible.shape[3] * visible.shape[4] * visible.shape[5] + mask = torch.zeros((nq, nk), dtype=dtype, device=device) + mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) + return mask.reshape(1, 1, nq, nk) + + +def na3d( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int] | tuple[int, ...], + is_causal: list[bool] | None = None, + scale: float | None = None, +) -> torch.Tensor: + """3D neighborhood attention over ``(B, T, H, W, NH, HD)`` tensors. + ``scale`` defaults to ``head_dim**-0.5``. Pass ``scale=1.0`` when Q is already scaled. + """ + batch, t, h, w, nh, hd = q.shape + dims = (t, h, w) + causal = [False, False, False] if is_causal is None else list(is_causal) + kernels = [k_ if c else min(k_, d) for k_, c, d in zip(kernel_size, causal, dims, strict=True)] + if scale is None: + scale = hd**-0.5 + device = q.device + if scale != 1.0: + q = q * scale + + bounds = [_window_bounds(d, k_, c) for d, k_, c in zip(dims, kernels, causal, strict=True)] + tile_t, tile_h, tile_w = _pick_tiles(dims, [min(k_, d) for k_, d in zip(kernels, dims, strict=True)]) + + groups: dict[ + tuple[ + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + ], + list[tuple[tuple[slice, slice, slice], tuple[slice, slice, slice]]], + ] = {} + for t0 in range(0, t, tile_t): + t1 = min(t0 + tile_t, t) + rt0, rt1 = bounds[0][0][t0], bounds[0][1][t1 - 1] + rel_t = ( + tuple(s - rt0 for s in bounds[0][0][t0:t1]), + tuple(e - rt0 for e in bounds[0][1][t0:t1]), + ) + for h0 in range(0, h, tile_h): + h1 = min(h0 + tile_h, h) + rh0, rh1 = bounds[1][0][h0], bounds[1][1][h1 - 1] + rel_h = ( + tuple(s - rh0 for s in bounds[1][0][h0:h1]), + tuple(e - rh0 for e in bounds[1][1][h0:h1]), + ) + for w0 in range(0, w, tile_w): + w1 = min(w0 + tile_w, w) + rw0, rw1 = bounds[2][0][w0], bounds[2][1][w1 - 1] + rel_w = ( + tuple(s - rw0 for s in bounds[2][0][w0:w1]), + tuple(e - rw0 for e in bounds[2][1][w0:w1]), + ) + groups.setdefault((rel_t, rel_h, rel_w), []).append( + ( + (slice(t0, t1), slice(h0, h1), slice(w0, w1)), + (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), + ) + ) + + out = torch.empty((batch, t, h, w, nh, hd), device=device, dtype=v.dtype) + for rel, tiles in groups.items(): + mask = _group_mask(rel, q.dtype, device) + nq, nk = mask.shape[2], mask.shape[3] + g_max = max(1, NA_KV_STACK_BUDGET // max(1, batch * nh * nk * hd * 2)) if device.type == "cuda" else 1 + qs0, _ = tiles[0] + tq = qs0[0].stop - qs0[0].start + th = qs0[1].stop - qs0[1].start + tw = qs0[2].stop - qs0[2].start + for c0 in range(0, len(tiles), g_max): + chunk = tiles[c0 : c0 + g_max] + g = len(chunk) + q_s = torch.stack([q[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) + k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nq, hd) + k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + o = functional.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) + o = o.view(g, batch, nh, tq, th, tw, hd).permute(0, 1, 3, 4, 5, 2, 6) + for i, (qs, _) in enumerate(chunk): + out[:, qs[0], qs[1], qs[2]] = o[i] + + return out + + +"""Pure-torch joint (video + keyframe) 3D neighborhood attention. +Computes, for one softmax per query: +* **video query** at ``(t, h, w)``: its local ``Kt x Kh x Kw`` video window, clamped to the + volume and masked where it hangs over the edge, **plus** the whole ``Kh x Kw`` window at the + same ``(h, w)`` on each of the ``num_slots`` nearest keyframe planes. Keyframe visibility does + not depend on ``Kt`` -- a plane far outside the temporal radius is still visible. +* **keyframe query** on plane ``i``: the ``Kh x Kw`` window on its own plane (there is no + plane-to-plane attention), plus the same window on each of the nearest video frames. +Which planes and frames are "nearest" comes from :func:`video_keyframe_slots` / +:func:`keyframe_video_slots`, so every backend agrees on visibility. +No Triton, no natten, no ``torch.compile``: this is the backend that always exists -- CPU, +macOS/MPS, Windows without a built extra. +Structure +--------- +Everything is arranged so the arithmetic happens inside ``F.scaled_dot_product_attention``: +**Query bricks.** Queries are grouped into ``(bt, bh, bw)`` bricks and many bricks ride one SDPA +call as its batch dimension, so a frame costs a handful of launches rather than thousands. All +queries in a brick share one gathered key slab, of extent ``(bt + Kt - 1, bh + Kh - 1, bw + Kw - 1)``. +**Brick shape.** Wasted work is ``Nk / keys_actually_visible`` and *grows* with the brick, so the +spatial face stays small and square (square minimizes the slab at a fixed query count). Depth is +the exception: the gather is the larger cost and it scales with ``Nk / Nq``, which *falls* with +depth, so :data:`DEFAULT_BRICK_DEPTH` frames deep beats one frame deep despite doing more +arithmetic. Both defaults sit on measured plateaus at the production stage-5 shape. +**One shared, 2D-broadcast mask.** The visible-key pattern is a property of the brick geometry, +identical for every brick, so it is built once and passed as ``(1, 1, Nq, Nk)``. That shape is +load-bearing: torch keeps the memory-efficient backend and expands neither the mask nor the +scores, whereas a pre-expanded ``(G, NH, Nq, Nk)`` bias halves throughput and costs gigabytes. +**Per-key validity rides in the keys.** Out-of-volume positions, empty slots (``-1``) and invalid +planes are data-dependent, so folding them into the mask would make it per-brick. Instead ``K`` +carries one extra channel holding ``0`` for a live key and :data:`_DEAD` for a dead one, against a +constant ``1`` channel on ``Q``. Q arrives pre-scaled, so with ``scale=1.0`` that adds exactly the +bias to the score. The channel count is then rounded up to :data:`_HEAD_DIM_ALIGN`. +**Head-major staging.** Key slabs are gathered from a ``(B, NH, A, Hp, Wp, C)`` copy rather than +from the caller's channels-last layout, so the gather's innermost contiguous run is ``ew * C`` +instead of ``C``. +**Runs of constant slot row.** A brick spanning several frames shares one keyframe key slab, so it +must not straddle a change of visible planes. ``T`` is cut into maximal runs of identical slot rows +and bricks are tiled inside a run -- which also means one plane gather per run, entering the slab +view with a **zero** group stride. +Both loops are budgeted by :data:`DEFAULT_WORKSPACE_BYTES`: frames per staging pass, then +``(bricks, brick rows)`` per SDPA call. Peak transient memory is therefore bounded by that budget +and not by the volume, which is what lets this sit next to a decoder that has its own memory plan. +The gather is the floor: SDPA needs materialized ``(G, NH, Nk, HD)`` keys, so every key is copied +``Nk / Nq`` times. Only a fused neighborhood kernel avoids that. +""" + + +_DEAD = -1.0e4 + + +_HEAD_DIM_ALIGN = 8 + + +DEFAULT_BRICK_QUERIES = 64 + + +DEFAULT_BRICK_DEPTH = 4 + + +DEFAULT_WORKSPACE_BYTES = 256 * 1024**2 + + +_STAGING_FACTOR_FUSED = 4.75 + + +_STAGING_FACTOR_MATERIALIZED = 22.1 + + +def sdpa_materializes_scores(device: torch.device) -> bool: + """Whether torch's SDPA will materialize this backend's score block on ``device``. + CUDA takes the memory-efficient (or cuDNN) kernel for this backend's broadcast mask and + aligned head dim. Everywhere else the math path runs and the ``(G, NH, Nq, Nk)`` scores land + in memory -- notably MPS. Auto tiling reads this to size its reserve. + """ + return device.type != "cuda" + + +def staging_factor(device: torch.device) -> float: + """The peak-to-staging multiplier for this host, from the table above.""" + return _STAGING_FACTOR_MATERIALIZED if sdpa_materializes_scores(device) else _STAGING_FACTOR_FUSED + + +def _key_channels(head_dim: int) -> int: + """Key/query channel count: ``head_dim``, the bias channel, and alignment padding.""" + return -(-(head_dim + 1) // _HEAD_DIM_ALIGN) * _HEAD_DIM_ALIGN + + +def _window(kernel: int) -> tuple[int, int]: + """``(lo, hi)`` halo for one axis: the offsets ``range(-k // 2, k - k // 2)`` reach.""" + lo = kernel // 2 + return lo, kernel - lo - 1 + + +def pick_brick( + time: int, + height: int, + width: int, + target: int = DEFAULT_BRICK_QUERIES, + depth: int = DEFAULT_BRICK_DEPTH, +) -> tuple[int, int, int]: + """``(bt, bh, bw)``: ``depth`` frames deep with the squarest ~``target``-query face. + Square minimizes the key slab ``(bh + Kh - 1) * (bw + Kw - 1)`` at a fixed query count, which + is exactly the wasted-work term. + """ + side = max(1, round(math.sqrt(target))) + return min(depth, time), min(side, height), min(side, width) + + +class _Geometry: + """Brick decomposition of a volume, plus the padding and slab extents it implies.""" + + def __init__( + self, + height: int, + width: int, + kernel: tuple[int, int, int], + brick: tuple[int, int, int], + ) -> None: + kernel_t, kernel_h, kernel_w = kernel + lo_h, hi_h = _window(kernel_h) + lo_w, hi_w = _window(kernel_w) + self.height, self.width = height, width + self.brick = brick + self.kernel = kernel + self.grid = (-(-height // brick[1]), -(-width // brick[2])) + # Slab extents: T grows with the brick depth, H/W with the spatial face. + self.span_t = brick[0] + kernel_t - 1 + self.span = (brick[1] + kernel_h - 1, brick[2] + kernel_w - 1) + # Halo, plus enough to cover the last (partial) brick's slab. + self.pad_h = (lo_h, hi_h + self.grid[0] * brick[1] - height) + self.pad_w = (lo_w, hi_w + self.grid[1] * brick[2] - width) + self.pad_t = _window(kernel_t) + self.queries = brick[0] * brick[1] * brick[2] + self.footprint = self.span[0] * self.span[1] + self.padded_height = height + sum(self.pad_h) + self.padded_width = width + sum(self.pad_w) + + def row_extent(self, rows: int) -> int: + """Padded ``H`` extent a group of ``rows`` brick rows needs from the staged volume.""" + return (rows - 1) * self.brick[1] + self.span[0] + + +class _Schedule: + """How the nested loops are cut so transient memory stays inside the budget. + ``group_axis`` counts *bricks* along the volume's leading axis (frames for the video pass, + planes for the keyframe pass); ``stage_axis`` counts them per staging pass. + """ + + def __init__( + self, + geometry: _Geometry, + blocks: int, + heads: int, + head_dim: int, + axis_bricks: int, + element_size: int, + workspace_bytes: int, + factor: float, + ) -> None: + channels = _key_channels(head_dim) + # One (brick along the axis, brick row) pair's worth of gathered keys and values, plus its + # score block on backends that materialize one. + keys = blocks * geometry.footprint + pair_bytes = geometry.grid[1] * heads * keys * (channels + head_dim) * element_size + # ``factor`` folds in whatever the selected SDPA kernel allocates on top of the staging, + # chiefly a materialized score block. See :data:`_STAGING_FACTOR_FUSED`. + pairs = max(1, int(workspace_bytes / max(pair_bytes * factor, 1.0))) + if pairs >= geometry.grid[0]: + self.group_axis = min(axis_bricks, max(1, pairs // geometry.grid[0])) + self.group_rows = geometry.grid[0] + else: + self.group_axis = 1 + self.group_rows = pairs + staged = geometry.padded_height * geometry.padded_width * heads * (channels + head_dim) * element_size + per_axis_brick = staged * geometry.brick[0] + self.stage_axis = min(axis_bricks, max(self.group_axis, workspace_bytes // max(per_axis_brick, 1))) + + +def _banded(queries: int, span: int, kernel: int, device: torch.device) -> torch.Tensor: + """``(queries, span)`` bool: key ``i`` is visible to query ``j`` iff ``j <= i < j + kernel``.""" + key = torch.arange(span, device=device)[None, :] + query = torch.arange(queries, device=device)[:, None] + return (key >= query) & (key < query + kernel) + + +def _joint_mask(geometry: _Geometry, num_slots: int, device: torch.device) -> torch.Tensor: + """``(1, 1, Nq, Nk)`` visibility, shared by every brick. + Query order is ``(jt, jh, jw)``; key order is the video slab ``(it, p, r)`` followed by the + keyframe slab ``(slot, p, r)``. Keyframe keys carry no temporal condition -- a plane is visible + to every frame in the brick, which is what makes the run grouping legal. + """ + brick_t, brick_h, brick_w = geometry.brick + kernel_t, kernel_h, kernel_w = geometry.kernel + spatial = ( + _banded(brick_h, geometry.span[0], kernel_h, device)[:, None, :, None] + & _banded(brick_w, geometry.span[1], kernel_w, device)[None, :, None, :] + ).reshape(brick_h * brick_w, geometry.footprint) + temporal = _banded(brick_t, geometry.span_t, kernel_t, device) + video = (temporal[:, None, :, None] & spatial[None, :, None, :]).reshape( + geometry.queries, geometry.span_t * geometry.footprint + ) + planes = ( + spatial[None, :, None, :] + .expand(brick_t, brick_h * brick_w, num_slots, geometry.footprint) + .reshape(geometry.queries, num_slots * geometry.footprint) + ) + return torch.cat([video, planes], dim=1)[None, None].contiguous() + + +def _stage( + x: torch.Tensor, + geometry: _Geometry, + pad_t: tuple[int, int], + *, + with_bias_channel: bool, +) -> torch.Tensor: + """``(B, A, H, W, NH, HD)`` -> padded head-major ``(B, NH, A + pad, Hp, Wp, C)``. + Head-major so a brick slab's innermost ``(ew, C)`` block is contiguous in both source and + destination; channels-last staging makes the same gather markedly slower. + """ + batch, axis, height, width, heads, head_dim = x.shape + channels = _key_channels(head_dim) if with_bias_channel else head_dim + out = x.new_zeros((batch, heads, axis + sum(pad_t), geometry.padded_height, geometry.padded_width, channels)) + if with_bias_channel: + out[..., head_dim] = _DEAD + live = out[ + :, + :, + pad_t[0] : pad_t[0] + axis, + geometry.pad_h[0] : geometry.pad_h[0] + height, + geometry.pad_w[0] : geometry.pad_w[0] + width, + ] + live[..., :head_dim] = x.permute(0, 4, 1, 2, 3, 5) + if with_bias_channel: + live[..., head_dim] = 0.0 + return out + + +def _slabs( + staged: torch.Tensor, + geometry: _Geometry, + bricks: int, + rows: int, + blocks: int, + *, + group_stride: int, +) -> torch.Tensor: + """Overlapping brick slabs as a *view*: ``(B, bricks, rows, Gw, NH, blocks, eh, ew, C)``. + ``staged`` is head-major ``(B, NH, A, Hp, Wp, C)``, already sliced to this group's first brick + and brick row, so the view inherits its storage offset. ``group_stride`` is how far consecutive + bricks advance along ``A``: the brick depth for the sliding video window, and **zero** for the + keyframe planes, which every brick in a run shares. + """ + batch, heads = staged.shape[0], staged.shape[1] + stride_b, stride_nh, stride_a, stride_h, stride_w, _ = staged.stride() + return staged.as_strided( + (batch, bricks, rows, geometry.grid[1], heads, blocks, *geometry.span, staged.shape[-1]), + ( + stride_b, + group_stride * stride_a, + geometry.brick[1] * stride_h, + geometry.brick[2] * stride_w, + stride_nh, + stride_a, + stride_h, + stride_w, + 1, + ), + ) + + +def _query_bricks(x: torch.Tensor, geometry: _Geometry, bricks: int, rows: int) -> torch.Tensor: + """``(B, A, h, W, NH, HD)`` -> ``(B * bricks * rows * Gw, NH, Nq, C)``, unit channel set. + ``x`` is this group's slice, so ``A`` may be short of ``bricks * bt`` and ``h`` short of + ``rows * bh`` at a volume edge; the shortfall is zero-padded here and cropped by + :func:`_unbrick`. + """ + batch, axis, height, width, heads, head_dim = x.shape + brick_t, brick_h, brick_w = geometry.brick + pad_t, pad_h, pad_w = bricks * brick_t - axis, rows * brick_h - height, geometry.grid[1] * brick_w - width + if pad_t or pad_h or pad_w: + x = F.pad(x, (0, 0, 0, 0, 0, pad_w, 0, pad_h, 0, pad_t)) + bricked = ( + x.reshape(batch, bricks, brick_t, rows, brick_h, geometry.grid[1], brick_w, heads, head_dim) + .permute(0, 1, 3, 5, 7, 2, 4, 6, 8) + .reshape(batch * bricks * rows * geometry.grid[1], heads, geometry.queries, head_dim) + ) + out = bricked.new_zeros((*bricked.shape[:-1], _key_channels(head_dim))) + out[..., :head_dim] = bricked + out[..., head_dim] = 1.0 + return out + + +def _unbrick( + attended: torch.Tensor, + geometry: _Geometry, + batch: int, + bricks: int, + rows: int, + extent: tuple[int, int], +) -> torch.Tensor: + """Inverse of :func:`_query_bricks`, cropping to ``extent`` frames/rows and the real width.""" + brick_t, brick_h, brick_w = geometry.brick + heads, head_dim = attended.shape[1], attended.shape[3] + plane = ( + attended.reshape(batch, bricks, rows, geometry.grid[1], heads, brick_t, brick_h, brick_w, head_dim) + .permute(0, 1, 5, 2, 6, 3, 7, 4, 8) + .reshape(batch, bricks * brick_t, rows * brick_h, geometry.grid[1] * brick_w, heads, head_dim) + ) + return plane[:, : extent[0], : extent[1], : geometry.width] + + +def _with_null(slots: torch.Tensor, null_index: int) -> torch.Tensor: + """Map empty slots (``-1``) onto the appended null row, which biases itself out.""" + return torch.where(slots < 0, torch.full_like(slots, null_index), slots) + + +def _append_null(keys: torch.Tensor, values: torch.Tensor, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + """Append one all-dead key plane (and a zero value plane) along the staged plane axis.""" + shape = (keys.shape[0], keys.shape[1], 1, *keys.shape[3:]) + null_key = keys.new_zeros(shape) + null_key[..., head_dim] = _DEAD + null_value = values.new_zeros((*shape[:-1], values.shape[-1])) + return torch.cat([keys, null_key], dim=2), torch.cat([values, null_value], dim=2) + + +def _slot_runs(slots: torch.Tensor) -> list[tuple[int, int]]: + """Maximal ``[start, stop)`` runs of leading-axis positions whose slot row is identical. + A brick spanning several frames shares one keyframe key slab, so it may not straddle a change + of slot row. At production keyframe spacing these runs are ~16 frames long, so the constraint + costs little; carrying every frame's slots in the slab instead would more than give back what + brick depth wins. + """ + rows = slots.tolist() + runs: list[tuple[int, int]] = [] + start = 0 + for index in range(1, len(rows)): + if rows[index] != rows[start]: + runs.append((start, index)) + start = index + runs.append((start, len(rows))) + return runs + + +def _attend_group( + query_slice: torch.Tensor, + key_views: tuple[torch.Tensor, ...], + value_views: tuple[torch.Tensor, ...], + geometry: _Geometry, + shape: tuple[int, int], + mask: torch.Tensor, +) -> torch.Tensor: + """Gather one ``(bricks x brick rows)`` block's keys, attend, and un-brick the result. + ``key_views`` / ``value_views`` are the strided slab views to concatenate along the key axis, + in order. ``shape`` is ``(bricks, rows)``. + """ + bricks, rows = shape + batch = query_slice.shape[0] + heads, head_dim = query_slice.shape[4], query_slice.shape[5] + blocks = sum(view.shape[5] for view in key_views) + channels = _key_channels(head_dim) + keys = query_slice.new_empty((batch, bricks, rows, geometry.grid[1], heads, blocks, *geometry.span, channels)) + values = query_slice.new_empty((batch, bricks, rows, geometry.grid[1], heads, blocks, *geometry.span, head_dim)) + start = 0 + for key_view, value_view in zip(key_views, value_views, strict=True): + stop = start + key_view.shape[5] + keys[:, :, :, :, :, start:stop].copy_(key_view) + values[:, :, :, :, :, start:stop].copy_(value_view) + start = stop + count = batch * bricks * rows * geometry.grid[1] + attended = F.scaled_dot_product_attention( + _query_bricks(query_slice, geometry, bricks, rows), + keys.view(count, heads, blocks * geometry.footprint, channels), + values.view(count, heads, blocks * geometry.footprint, head_dim), + attn_mask=mask, + scale=1.0, + ) + return _unbrick(attended, geometry, batch, bricks, rows, (query_slice.shape[1], query_slice.shape[2])) + + +def _row_groups(geometry: _Geometry, schedule: _Schedule) -> list[tuple[int, int, slice, slice]]: + """``(row, rows, staged H slice, output H slice)`` per brick-row group.""" + brick_h = geometry.brick[1] + groups = [] + for row in range(0, geometry.grid[0], schedule.group_rows): + rows = min(schedule.group_rows, geometry.grid[0] - row) + groups.append( + ( + row, + rows, + slice(row * brick_h, row * brick_h + geometry.row_extent(rows)), + slice(row * brick_h, min((row + rows) * brick_h, geometry.height)), + ) + ) + return groups + + +def _video_query_pass( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + slots: torch.Tensor, + geometry: _Geometry, + workspace_bytes: int, + factor: float, +) -> torch.Tensor: + """Video queries: the local ``Kt x Kh x Kw`` window plus the nearest keyframe planes.""" + time, heads, head_dim = q.shape[1], q.shape[4], q.shape[5] + brick_t = geometry.brick[0] + lo_t, hi_t = geometry.pad_t + num_slots = slots.shape[1] + blocks = geometry.span_t + num_slots + + plane_keys, plane_values = _append_null( + _stage(keyframe_k, geometry, (0, 0), with_bias_channel=True), + _stage(keyframe_v, geometry, (0, 0), with_bias_channel=False), + head_dim, + ) + slot_table = _with_null(slots, keyframe_k.shape[1]) + mask = _joint_mask(geometry, num_slots, q.device) + schedule = _Schedule( + geometry, + blocks, + heads, + head_dim, + -(-time // brick_t), + q.element_size(), + workspace_bytes, + factor, + ) + rows_groups = _row_groups(geometry, schedule) + + out = torch.empty_like(q) + for run_start, run_stop in _slot_runs(slot_table): + # One plane gather per run: every brick inside it sees the same slots. + planes = plane_keys.index_select(2, slot_table[run_start]) + plane_vals = plane_values.index_select(2, slot_table[run_start]) + run_bricks = -(-(run_stop - run_start) // brick_t) + for staged_brick in range(0, run_bricks, schedule.stage_axis): + staged_bricks = min(schedule.stage_axis, run_bricks - staged_brick) + first = run_start + staged_brick * brick_t + last = first + staged_bricks * brick_t # exclusive; may reach past the run or T + source = slice(max(0, first - lo_t), min(time, last + hi_t)) + pad_t = (max(0, lo_t - first), max(0, last + hi_t - time)) + window_keys = _stage(k[:, source], geometry, pad_t, with_bias_channel=True) + window_values = _stage(v[:, source], geometry, pad_t, with_bias_channel=False) + + for brick in range(staged_brick, staged_brick + staged_bricks, schedule.group_axis): + count = min(schedule.group_axis, staged_brick + staged_bricks - brick) + start = run_start + brick * brick_t + stop = min(start + count * brick_t, run_stop) + offset = (brick - staged_brick) * brick_t + for _, rows, key_rows, out_rows in rows_groups: + tile = _attend_group( + q[:, start:stop, out_rows], + ( + _slabs( + window_keys[:, :, offset:, key_rows], + geometry, + count, + rows, + geometry.span_t, + group_stride=brick_t, + ), + _slabs(planes[:, :, :, key_rows], geometry, count, rows, num_slots, group_stride=0), + ), + ( + _slabs( + window_values[:, :, offset:, key_rows], + geometry, + count, + rows, + geometry.span_t, + group_stride=brick_t, + ), + _slabs(plane_vals[:, :, :, key_rows], geometry, count, rows, num_slots, group_stride=0), + ), + geometry, + (count, rows), + mask, + ) + out[:, start:stop, out_rows] = tile + return out + + +def _keyframe_query_pass( + keyframe_q: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + slots: torch.Tensor, + keyframe_valid: torch.Tensor, + geometry: _Geometry, + workspace_bytes: int, + factor: float, +) -> torch.Tensor: + """Keyframe queries: own plane only (``d_t == 0``) plus the nearest video frames. + Runs one plane per brick: planes have no temporal window, so depth would buy nothing, and each + plane's video slots differ anyway. + """ + planes_total, heads, head_dim = keyframe_q.shape[1], keyframe_q.shape[4], keyframe_q.shape[5] + num_slots = slots.shape[1] + blocks = 1 + num_slots + time = k.shape[1] + flat = _Geometry(geometry.height, geometry.width, (1, *geometry.kernel[1:]), (1, *geometry.brick[1:])) + + # Only the frames some plane actually points at get staged -- at most ``P * num_slots`` of them, + # against the whole volume if this staged ``k`` wholesale. ``unique`` doubles as the remap: slot + # rows are rewritten to index the compacted stack. + wanted, inverse = torch.unique(_with_null(slots, time).reshape(-1), return_inverse=True) + frame_keys = _stage(k.index_select(1, wanted.clamp(max=time - 1)), flat, (0, 0), with_bias_channel=True) + frame_values = _stage(v.index_select(1, wanted.clamp(max=time - 1)), flat, (0, 0), with_bias_channel=False) + # An empty slot clamped onto a real frame above; kill it here instead of appending a null row. + frame_keys[:, :, wanted == time, ..., head_dim] = _DEAD + own_keys = _stage(keyframe_k, flat, (0, 0), with_bias_channel=True) + own_values = _stage(keyframe_v, flat, (0, 0), with_bias_channel=False) + own_keys[:, :, ~keyframe_valid, ..., head_dim] = _DEAD + slot_table = inverse.reshape(planes_total, num_slots) + mask = _joint_mask(flat, num_slots, keyframe_q.device) + schedule = _Schedule( + flat, + blocks, + heads, + head_dim, + planes_total, + keyframe_q.element_size(), + workspace_bytes, + factor, + ) + rows_groups = _row_groups(flat, schedule) + + out = torch.empty_like(keyframe_q) + for start in range(0, planes_total, schedule.group_axis): + stop = min(start + schedule.group_axis, planes_total) + count = stop - start + picked = slot_table[start:stop].reshape(-1) + frames = frame_keys.index_select(2, picked) + frame_vals = frame_values.index_select(2, picked) + for _, rows, key_rows, out_rows in rows_groups: + tile = _attend_group( + keyframe_q[:, start:stop, out_rows], + ( + _slabs(own_keys[:, :, start:, key_rows], flat, count, rows, 1, group_stride=1), + _slabs(frames[:, :, :, key_rows], flat, count, rows, num_slots, group_stride=num_slots), + ), + ( + _slabs(own_values[:, :, start:, key_rows], flat, count, rows, 1, group_stride=1), + _slabs(frame_vals[:, :, :, key_rows], flat, count, rows, num_slots, group_stride=num_slots), + ), + flat, + (count, rows), + mask, + ) + out[:, start:stop, out_rows] = tile + # An invalid plane sees nothing; zero it rather than shipping the uniform mean. + return out * keyframe_valid[None, :, None, None, None, None] + + +def joint_na3d( # noqa: PLR0913 + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keyframe_q: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + kernel_size: tuple[int, int, int], + num_slots: int = KEYFRAME_CONTEXT_SLOTS, + brick: tuple[int, int, int] | None = None, + workspace_bytes: int = DEFAULT_WORKSPACE_BYTES, + factor: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Joint neighborhood attention over a video volume and a keyframe plane stack. + Args: + q, k, v: ``(B, T, H, W, NH, HD)`` video stream. ``q`` must arrive pre-scaled by + ``head_dim ** -0.5``, as the shared attention module does. + keyframe_q, keyframe_k, keyframe_v: ``(B, P, H, W, NH, HD)`` keyframe stream. + keyframe_times: ``(P,)`` float32 plane times, same origin as the video RoPE. + keyframe_valid: ``(P,)`` bool. + kernel_size: ``(Kt, Kh, Kw)``. + num_slots: cross-stream slots per query. + brick: query brick ``(bt, bh, bw)``; defaults to :func:`pick_brick`. + workspace_bytes: transient budget bounding the staged window and key/value block. + factor: peak-to-staging multiplier; :func:`staging_factor` supplies it when omitted. + Returns: + ``(video_out, keyframe_out)``, each shaped like its stream's ``q``. + """ + time, height, width = q.shape[1], q.shape[2], q.shape[3] + video_slots = video_keyframe_slots(keyframe_times, keyframe_valid, time, num_slots) + keyframe_slots = keyframe_video_slots(keyframe_times, keyframe_valid, time, num_slots) + geometry = _Geometry(height, width, kernel_size, brick if brick is not None else pick_brick(time, height, width)) + if factor is None: + factor = staging_factor(q.device) + return ( + _video_query_pass(q, k, v, keyframe_k, keyframe_v, video_slots, geometry, workspace_bytes, factor), + _keyframe_query_pass( + keyframe_q, + keyframe_k, + keyframe_v, + k, + v, + keyframe_slots, + keyframe_valid, + geometry, + workspace_bytes, + factor, + ), + ) + + + +class EagerNAAttention: + def __call__(self, attn, q, k, v): + return na3d(q, k, v, kernel_size=attn.kernel_size, scale=1.0) + + +class EagerJointNAAttention: + def __call__(self, attn, q, k, v, keyframe_q, keyframe_k, keyframe_v, keyframe_times, keyframe_valid): + return joint_na3d( + q, k, v, keyframe_q, keyframe_k, keyframe_v, + keyframe_times, keyframe_valid, kernel_size=attn.kernel_size, + ) + + + +"""3D Neighborhood Attention via NATTEN + absolute RoPE prelude. +Parameter shell shared by det ``NABlock`` and both diff-attn roles. +Diffusion AdaLN residuals live in pathway packages (each owns its RoPE); +det stages use ``det_attn_rope`` from :meth:`NeighborhoodAttention3D.forward`. +``attention_function`` selects the NA backend (NATTEN, Triton/eager fallback, or CuTe DSL). +""" + + +if TYPE_CHECKING: + from ltx_core.model.video_vae.keyframes import KeyframeStream + + +try: + import natten + + _NATTEN_AVAILABLE = True +except ImportError: # pragma: no cover + natten = None # type: ignore[assignment] + _NATTEN_AVAILABLE = False + + +def natten_available() -> bool: + return _NATTEN_AVAILABLE + + +class NAAttentionCallable(Protocol): + """A windowed 3D neighborhood-attention backend. + Q/K/V arrive as ``(B, T, H, W, NH, HD)``, already normed, scaled and RoPE'd; + the return is ``(B, T, H, W, NH*HD)`` or anything reshapeable to it. The owning + module is passed so a backend can read configuration (``kernel_size``, softmax + bound, …). Backend-specific settings (NATTEN's kernel pin) live on the callable. + """ + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: ... + + +class JointNAAttentionCallable(Protocol): + """A windowed 3D NA backend that also carries a keyframe plane stack. + Same conventions as :class:`NAAttentionCallable` -- Q/K/V already normed, scaled and + absolutely RoPE'd -- with a second stream shaped ``(B, P, H, W, NH, HD)`` whose plane + axis sits in video's temporal slot. Both streams' RoPE must share one origin, so + ``keyframe_times`` are tile-local. Returns one output per stream. + NATTEN and the CuTe DSL kernel cannot express a joint window, so this is a separate + slot from ``attention_function`` rather than a widening of it: it keeps the shipping + keyframe-less hot path untouched, and it is immune to the install-order hazard that + ``configure_natten_backend`` creates by overwriting ``attention_function`` wholesale. + """ + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keyframe_q: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + +class NattenAttention(NAAttentionCallable): + """``natten.na3d``, the default backend. + ``backend`` pins ``na3d``'s own kernel choice (e.g. ``"cutlass-fna"``); ``None`` + leaves NATTEN's auto-pick (hopper-fna on H100, etc.). + """ + + def __init__(self, backend: str | None = None) -> None: + self._backend = backend + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + if not _NATTEN_AVAILABLE: + raise ImportError( + "natten is required for NeighborhoodAttention3D. " + "Install with: uv sync --package ltx-core --extra natten " + '(or: uv pip install "natten==0.21.7+torch2130cu132" -f https://whl.natten.org; ' + "requires torch==2.13.0+cu132)" + ) + # scale=1.0: callers already applied ``attn.scale`` to Q. + # RMSNorm under bf16 autocast can leave Q/K in float32 while V stays bf16; + # natten requires a uniform dtype (same cast pattern as flash-attn paths). + if q.dtype != v.dtype or k.dtype != v.dtype: + q = q.to(dtype=v.dtype) + k = k.to(dtype=v.dtype) + return natten.na3d(q, k, v, kernel_size=attn.kernel_size, scale=1.0, backend=self._backend) + + +class NeighborhoodAttention3D(nn.Module): + """3D Neighborhood Attention with absolute RoPE + pluggable NA backend. + Q/K receive absolute RoPE; attention is ``attention_function`` (NATTEN by + default; Triton or eager SDPA when natten is missing; CuTe DSL via + DiffVAE BLACKWELL_DSL install). Relative gather-based NA is not used as a + production gather path on this branch. + NATTEN shifts its window inward at grid boundaries instead of + clamp-and-mask; interior positions match the gather reference closely, + boundary positions may differ slightly. + """ + + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + head_dim: int = 64, + rope_dim_split: tuple[int, int, int] | None = None, + rope_base: float = 10000.0, + ) -> None: + super().__init__() + assert dim % head_dim == 0, f"dim={dim} not divisible by head_dim={head_dim}" + self.dim = dim + self.num_heads = dim // head_dim + self.head_dim = head_dim + self.kernel_size = tuple(kernel_size) + self.scale = head_dim**-0.5 + + if rope_dim_split is None: + rope_dim_split = default_rope_dim_split(head_dim) + assert sum(rope_dim_split) == head_dim, f"rope_dim_split={rope_dim_split} must sum to head_dim={head_dim}" + self.rope_dim_split = rope_dim_split + self.rope_base = rope_base + self.rope_num_tiles = DEFAULT_ABS_ROPE_NUM_TILES + self.rope_compute_dtype = torch.float32 + # Kept for the chunked opaque residual (string arg); callable is the swap surface. + self.natten_backend: str | None = None + self.attention_function: NAAttentionCallable = EagerNAAttention() + # Separate slot, installed for every mode; never NATTEN/DSL. Only the keyframe + # decode path reads it, so keyframe-less decode keeps NATTEN when it is installed. + self.joint_attention_function: JointNAAttentionCallable | None = EagerJointNAAttention() + + self.register_buffer("rope_inv_t", rope_inv_freqs(rope_dim_split[0], rope_base), persistent=False) + self.register_buffer("rope_inv_h", rope_inv_freqs(rope_dim_split[1], rope_base), persistent=False) + self.register_buffer("rope_inv_w", rope_inv_freqs(rope_dim_split[2], rope_base), persistent=False) + + self.qkv = QKVProjections(dim) + self.proj = nn.Linear(dim, dim, bias=True) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + + # W-chunking configuration (consumed by ``chunked.attn``). + self.w_chunks = 1 # 1 = no chunking + + def project_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V as owned contiguous ``(B,T,H,W,NH,HD)`` tensors.""" + batch, t, h, w, _ = x.shape + q, k, v = self.qkv(x) + shape = (batch, t, h, w, self.num_heads, self.head_dim) + return q.view(shape), k.view(shape), v.view(shape) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Det-stage NA: opaque abs-RoPE via ``det_attn_rope`` + ``attention_function``. + ``x``/output: (B, T, H, W, C) — channels-last. RoPE positions are local + 0-based (see ``det_attn_rope`` module docstring for why that is + equivalent under tiled decode). + """ + batch, t, h, w, _ = x.shape + kt, kh, kw = self.kernel_size + if t < kt or h < kh or w < kw: + raise ValueError( + f"3D neighborhood attention requires spatial dims >= kernel_size; " + f"got (T,H,W)=({t},{h},{w}) vs kernel={self.kernel_size}" + ) + + q, k, v = det_qkv_rope(self, x) + # natten's CUTLASS kernel silently produces wrong output if inputs are non-contiguous. + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = self.attention_function(self, q, k, v) + out = out.reshape(batch, t, h, w, self.dim) + return self.proj(out) + + def forward_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + ) -> tuple[torch.Tensor, KeyframeStream]: + """Dual-stream det NA: one joint softmax over video and keyframe planes. + Unlike :meth:`forward` there is no ``dims >= kernel_size`` floor: the joint window + is clamp-and-mask, so an undersized axis simply masks its out-of-range offsets. + """ + if self.joint_attention_function is None: + raise RuntimeError( + "keyframe decode needs joint_attention_function installed; build the decoder " + "through apply_diffvae_config / apply_diffvae_mode" + ) + batch, t, h, w, _ = x.shape + planes = keyframes.x.shape[1] + + q, k, v = det_qkv_rope(self, x) + keyframe_q, keyframe_k, keyframe_v = det_qkv_rope_at_times(self, keyframes.x, keyframes.times) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + keyframe_q = keyframe_q.contiguous() + keyframe_k = keyframe_k.contiguous() + keyframe_v = keyframe_v.contiguous() + + out, keyframe_out = self.joint_attention_function( + self, + q, + k, + v, + keyframe_q, + keyframe_k, + keyframe_v, + keyframes.times, + keyframes.valid, + ) + out = self.proj(out.reshape(batch, t, h, w, self.dim)) + keyframe_out = self.proj(keyframe_out.reshape(batch, planes, h, w, self.dim)) + return out, dataclasses.replace(keyframes, x=keyframe_out) + + +def configure_w_chunks(module_root: nn.Module, w_chunks: int = 1) -> None: + """Set W-chunking on ``NeighborhoodAttention3D`` under ``module_root``. + When ``w_chunks > 1``, also sets ``rope_num_tiles=1`` so ``chunked.attn`` + owns the W axis (RoPE W-tiling would double-split). Pass only the diffusion + residual subtree — det-stage attention must keep its default RoPE tiling. + """ + for module in module_root.modules(): + if isinstance(module, NeighborhoodAttention3D): + module.w_chunks = w_chunks + if w_chunks > 1: + module.rope_num_tiles = 1 + + +"""NABlock and DiffusionNABlock parameter shells for DiffVAE. +Pathway subclasses live in ``chunked/`` and ``combined/``; ``apply`` installs +them via ``__class__`` swap (same pattern as ``Fp8CastLinear``). The shell owns +weights + shared AdaLN helpers only — no pathway forward. +""" + + +if TYPE_CHECKING: + from ltx_core.model.video_vae.keyframes import KeyframeStream + + +__all__ = [ + "DiffusionNABlock", + "NABlock", +] + + +class NABlock(nn.Module): + """Pre-norm transformer block: NA -> SwiGLU MLP with residual adds.""" + + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + head_dim: int = 64, + mlp_ratio: float = 4.0, + rope_dim_split: tuple[int, int, int] | None = None, + ) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(dim, eps=1e-6) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) + self.norm2 = nn.RMSNorm(dim, eps=1e-6) + hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 + self.mlp = SwiGLU(dim, hidden) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Channels-last in/out: (B, T, H, W, C).""" + x = x + self.attn(self.norm1(x)) + x = plain_mlp(x, self.mlp, self.norm2, self.mlp.tile) + return x + + def forward_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + ) -> tuple[torch.Tensor, KeyframeStream]: + """Dual-stream block: video ``(B,T,H,W,C)`` and keyframe planes ``(B,P,H,W,C)``. + Every weight is shared with :meth:`forward`; the streams meet only inside the + joint attention softmax. Invalid planes are deliberately *not* re-zeroed here -- + the decoder re-zeroes after each upsample instead, matching upstream, so a masked + plane's hidden state may drift within a stage. It is masked out of every softmax + regardless, so this is cosmetic; reproducing it keeps us comparable. + """ + attn_out, keyframe_attn = self.attn.forward_with_keyframes( + self.norm1(x), + dataclasses.replace(keyframes, x=self.norm1(keyframes.x)), + ) + x = x + attn_out + keyframe_x = keyframes.x + keyframe_attn.x + x = plain_mlp(x, self.mlp, self.norm2, self.mlp.tile) + keyframe_x = plain_mlp(keyframe_x, self.mlp, self.norm2, self.mlp.tile) + return x, dataclasses.replace(keyframes, x=keyframe_x) + + +class DiffusionNABlock(nn.Module): + """Parameter shell for diffusion NA + SwiGLU with shared AdaLN-Zero. + Mode-specific subclasses (:class:`~ltx_core.model.video_vae.transformer.combined.block.CombinedDiffusionNABlock`, + :class:`~ltx_core.model.video_vae.transformer.chunked.block.ChunkedDiffusionNABlock`) + are installed via ModuleOps ``__class__`` swap and own the forward path. + Not a ``Protocol``: must be a concrete ``nn.Module`` so checkpoint load and + ``__class__`` swap keep one parameter identity. + """ + + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + context_channels: int, + head_dim: int = 64, + mlp_ratio: float = 4.0, + rope_dim_split: tuple[int, int, int] | None = None, + ) -> None: + super().__init__() + self.context_channels = context_channels + self.context_proj = nn.Linear(context_channels, dim, bias=True) + self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim)) + + self.norm1 = nn.RMSNorm(dim, eps=1e-6) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) + self.norm2 = nn.RMSNorm(dim, eps=1e-6) + hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 + self.mlp = SwiGLU(dim, hidden) + self.attn.proj.reset_parameters() + + def _modulation( + self, modulation: tuple[torch.Tensor, ...] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) + ] + return scale_msa, shift_msa, scale_mlp, shift_mlp + + +"""Combined context residual: project context half of ``context_and_x`` into ``x``.""" + + +def combined( + context_and_x: torch.Tensor, + w_proj: torch.Tensor, + b_proj: torch.Tensor | None, +) -> torch.Tensor: + """Split ``context_and_x`` via ``w_proj.shape[1]`` (context channels), add projected ctx. + Returns the updated ``x`` half (not the full concatenated buffer). + ``w_proj`` is ``context_proj.weight`` with shape ``(dim, context_channels)``. + """ + context_channels = w_proj.shape[1] + latent_context = context_and_x[..., :context_channels] + x = context_and_x[..., context_channels:] + return x + F.linear(latent_context, w_proj, b_proj) + + +inject_context = combined + + +"""Combined* diffusion AdaLN residual attention (full-volume NA + nested RoPE). +Owns nested full-volume abs-RoPE for the Combined / ``w_chunks==1`` path. +Does not share a residual body with ``chunked`` — only the NA module weights. +""" + + +_rot_abs_axis = torch.compiler.nested_compile_region(rot_abs_axis_impl) + + +def _apply_nested_abs_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + """Rotate one W-extent with nested per-axis abs-RoPE. + ``t_pos`` overrides the integer ``arange`` on the first axis; the keyframe stream passes + its fractional plane times there so both streams' RoPE shares one origin. + """ + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + positions_t = t_positions(t, x.device) if t_pos is None else t_pos + xt = _rot_abs_axis(x[..., :d_t], positions_t, inv_t, axis=1, compute_dtype=compute_dtype) + xh = _rot_abs_axis( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = _rot_abs_axis(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _apply_nested_full_volume_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + """Fixed-``num_tiles`` W split + nested per-slab rotation (Dynamo-safe). + Both streams share the same W extent at stage 5, so the same ``num_tiles`` gives + identical slab boundaries and ``w_pos`` -- required for their W phases to be comparable + inside the joint softmax. + """ + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + _apply_nested_abs_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + t_pos=t_pos, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +def _qkv_nested_rope( + attn: NeighborhoodAttention3D, + x: torch.Tensor, + t_pos: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V proj + norm/scale + nested abs-RoPE. + ``t_pos`` overrides the integer first-axis positions; the keyframe stream passes + its (possibly fractional) plane times there. + """ + q, k, v = attn.project_qkv(x) + q = attn.q_norm(q) * attn.scale + k = attn.k_norm(k) + inv_freqs = ( + attn.rope_inv_t.to(device=x.device), + attn.rope_inv_h.to(device=x.device), + attn.rope_inv_w.to(device=x.device), + ) + positions = None if t_pos is None else t_pos.to(device=x.device, dtype=torch.float32) + q = _apply_nested_full_volume_rope( + q, + attn.rope_dim_split, + inv_freqs, + num_tiles=attn.rope_num_tiles, + compute_dtype=attn.rope_compute_dtype, + t_pos=positions, + ) + k = _apply_nested_full_volume_rope( + k, + attn.rope_dim_split, + inv_freqs, + num_tiles=attn.rope_num_tiles, + compute_dtype=attn.rope_compute_dtype, + t_pos=positions, + ) + return q, k, v + + +def full_with_keyframes( + x: torch.Tensor, + keyframe_x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Dual-stream AdaLN residual attention: one joint softmax, shared weights. + Plain tensors rather than a ``KeyframeStream`` here: this is the stage-5 hot path, and + rebuilding a dataclass per block inside a compiled region buys nothing. + Both streams take the *same* ``scale``/``shift``. Upstream computes a separate keyframe + modulation, but the two are identical unless per-frame timestep conditioning supplies a + conditioning mask, which we deliberately do not implement yet. + No ``dims >= kernel_size`` floor, unlike :func:`full`: the joint window is + clamp-and-mask, so undersized axes simply mask their out-of-range offsets. + """ + if attn.joint_attention_function is None: + raise RuntimeError( + "keyframe decode needs joint_attention_function installed; build the decoder " + "through apply_diffvae_config / apply_diffvae_mode" + ) + batch, t, h, w, _ = x.shape + planes = keyframe_x.shape[1] + + y = norm(x) * (1.0 + scale) + shift + keyframe_y = norm(keyframe_x) * (1.0 + scale) + shift + + q, k, v = _qkv_nested_rope(attn, y) + keyframe_q, keyframe_k, keyframe_v = _qkv_nested_rope(attn, keyframe_y, keyframe_times) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + keyframe_q = keyframe_q.contiguous() + keyframe_k = keyframe_k.contiguous() + keyframe_v = keyframe_v.contiguous() + + out, keyframe_out = attn.joint_attention_function( + attn, + q, + k, + v, + keyframe_q, + keyframe_k, + keyframe_v, + keyframe_times, + keyframe_valid, + ) + x = x + attn.proj(out.reshape(batch, t, h, w, attn.dim)) + keyframe_x = keyframe_x + attn.proj(keyframe_out.reshape(batch, planes, h, w, attn.dim)) + return x, keyframe_x + + +def full( + x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """``x + NA(modulate(norm(x)))`` with nested full-volume abs-RoPE.""" + y = norm(x) * (1.0 + scale) + shift + batch, t, h, w, _ = y.shape + kt, kh, kw = attn.kernel_size + if t < kt or h < kh or w < kw: + raise ValueError( + f"3D neighborhood attention requires spatial dims >= kernel_size; " + f"got (T,H,W)=({t},{h},{w}) vs kernel={attn.kernel_size}" + ) + + q, k, v = _qkv_nested_rope(attn, y) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = attn.attention_function(attn, q, k, v) + out = out.reshape(batch, t, h, w, attn.dim) + return x + attn.proj(out) + + +residual_attn = full +residual_attn_with_keyframes = full_with_keyframes + + +"""Combined pathway MLP: out-of-place AdaLN SwiGLU residual.""" + + +def residual_mlp( + x: torch.Tensor, + mlp: nn.Module, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, + tile: SwiGLUTileSpec, +) -> torch.Tensor: + """Combined*: ``x + swiglu_tiled(modulate(norm(x), scale, shift))``.""" + y = modulate(norm(x), scale, shift) + if y.numel() == 0: + return x + return x + swiglu_tiled(y, mlp.w_gate.weight, mlp.w_up.weight, mlp.w_down.weight, tile) + + +"""CombinedDiffusionNABlock: context_and_x inject + full-volume attn + residual MLP.""" + + +class CombinedDiffusionNABlock(DiffusionNABlock): + """Combined-context diffusion block: ``forward`` / ``forward_combined``.""" + + def forward_combined_with_keyframes( + self, + context_and_x: torch.Tensor, + keyframe_context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Dual-stream block; returns the two updated x halves (not the concat buffers). + Each stream gets its own ``context_proj(context)`` injection from its own + ``[context | x]`` buffer, then both meet in one joint attention softmax, then each + runs the shared MLP. Invalid keyframe planes are re-zeroed on the way out -- unlike + the deterministic ``NABlock``, which leaves that to the decoder's post-upsample + masking. Both asymmetries are upstream's. + """ + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) + keyframe_x = inject_context(keyframe_context_and_x, self.context_proj.weight, self.context_proj.bias) + x, keyframe_x = residual_attn_with_keyframes( + x, + keyframe_x, + self.attn, + self.norm1, + scale_msa, + shift_msa, + keyframe_times, + keyframe_valid, + ) + x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + keyframe_x = residual_mlp(keyframe_x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + return x, keyframe_x * keyframe_valid[None, :, None, None, None] + + def forward_combined( + self, + context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) + x = residual_attn(x, self.attn, self.norm1, scale_msa, shift_msa) + x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + return x + + def forward( + self, + context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + return self.forward_combined(context_and_x, modulation) + + +"""Diffusion (NATTEN) video VAE decoder.""" + + +logger: logging.Logger = logging.getLogger(__name__) + + +_L_STAGE_CHANNELS: Tuple[int, ...] = (1024, 512, 256, 256, 128) + + +_L_STAGE_DEPTHS: Tuple[int, ...] = (4, 6, 4, 2, 2) + + +_L_UPSAMPLES: Tuple[Tuple[Tuple[int, int, int], int], ...] = ( + ((1, 2, 2), 2), # compress_space x2 + ((2, 1, 1), 2), # compress_time x2 + ((2, 2, 2), 1), # compress_all x1 (channel-preserving) + ((2, 2, 2), 2), # compress_all x2 +) + + +_L_STAGE_KERNELS: Tuple[Tuple[int, int, int], ...] = ( + (3, 7, 7), + (3, 7, 7), + (3, 5, 5), + (3, 5, 5), + (3, 3, 3), +) + + +_DIFF_STAGE5_KERNEL_DEFAULT: Tuple[int, int, int] = (3, 7, 7) + + +_DIFF_STAGE5_DEPTH_DEFAULT: int = 8 + + +_DIFF_STAGE_DEPTHS_DEFAULT: Tuple[int, ...] = (*_L_STAGE_DEPTHS[:-1], _DIFF_STAGE5_DEPTH_DEFAULT) + + +class DiffusionVideoDecoder(nn.Module, Disposable, VideoDecoder): + """Diffusion-based video VAE decoder (Neighborhood-Attention backbone). + Minimal port of the reference ``NADiffusionDecoder``. + Stages 1-4 deterministically upsample the latent into a context volume + (same NA-upsample path as the non-diffusion NA decoder). Stage 5 runs + ``DiffusionNABlock``s that denoise the patchified noised pixels ``x_t``, + guided by that context via AdaLN-Zero scale/shift (ungated residuals; + legacy static gates are folded into Linear weights at load time). + Last-frame NATTEN window-shift is mitigated by temporarily replicating the + last latent frame ``(stage1_K_t // 2) * 2`` times through stages 1-4, then + cropping that appendix from context before stage 5 - but only down to + ``max(original_context_T, stage5_kernel[0])`` so undersized clips (e.g. a + single latent frame) still satisfy NATTEN's kernel floor. Latents / tiles + below ``stage_min_tile_sizes`` are edge-padded first via ``diffusion_tiling``; + leftover pad is cropped from the final pixels. + """ + + def __init__( # noqa: PLR0913 + self, + in_channels: int = 128, + out_channels: int = 3, + patch_size: int = 4, + head_dim: int = 64, + rope_dim_split: Tuple[int, int, int] | None = None, + stage_channels: Tuple[int, ...] = _L_STAGE_CHANNELS, + stage_depths: Tuple[int, ...] = _DIFF_STAGE_DEPTHS_DEFAULT, + stage_kernels: Tuple[Tuple[int, int, int], ...] = _L_STAGE_KERNELS, + upsamples: Tuple[Tuple[Tuple[int, int, int], int], ...] = _L_UPSAMPLES, + stage5_kernel: Tuple[int, int, int] = _DIFF_STAGE5_KERNEL_DEFAULT, + stage5_channels: int | None = None, + t_emb_dim: int = 384, + default_num_inference_steps: int = 2, + timestep_scale_multiplier: float = 1.0, + model_output_type: Literal["v", "x0"] = "v", + ) -> None: + super().__init__() + assert len(stage_channels) == len(stage_depths) == len(stage_kernels) + assert len(upsamples) == len(stage_channels) - 1 + for c in stage_channels: + assert c % head_dim == 0, f"stage_channels {stage_channels} must each be a multiple of head_dim={head_dim}" + + self.patch_size = patch_size + self.register_buffer( + "default_inference_timesteps", + torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps, device="cpu"), + persistent=False, + ) + self.out_channels = out_channels + self.stage_channels = stage_channels + self.stage_depths = stage_depths + self.base_channels = stage_channels[-1] + self.causal = False + self.timestep_conditioning = True + self.video_downscale_factors = SpatioTemporalScaleFactors.default() + self.stage5_kernel: Tuple[int, int, int] = tuple(stage5_kernel) # type: ignore[assignment] + # NATTEN last-frame border workaround: replicate last latent frame + # ``(K_t // 2) * 2`` times through stages 1-4, then crop the appendix + # off context before stage 5 down to at least ``stage5_kernel[0]``. + self._natten_trailing_pad_latent_frames = (stage_kernels[0][0] // 2) * 2 + + # Encoder output is per-channel normalized; undo before conv_in (same as ConvVideoDecoder). + self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) + + self.conv_in = ChannelLinear(in_channels, stage_channels[0], bias=True) + # Keyframe-stream tag, added to un-normalized keyframe latents before the shared + # ``conv_in`` and nowhere else. It is the only keyframe-specific weight in the + # whole feature. Checkpoints predating the keyframe training have no such key, so + # ``video_decoder_sd_ops_for_checkpoint`` synthesizes zeros -- a missing key would + # otherwise leave the parameter on the meta device under ``strict=False`` load. + self.type_emb = nn.Parameter(torch.zeros(in_channels)) + + self.det_stages = nn.ModuleList() + self.upsamples = nn.ModuleList() + n_det_stages = len(stage_channels) - 1 + for stage_i in range(n_det_stages): + c = stage_channels[stage_i] + depth = stage_depths[stage_i] + kernel = stage_kernels[stage_i] + self.det_stages.append( + nn.ModuleList( + [ + NABlock(dim=c, kernel_size=kernel, head_dim=head_dim, rope_dim_split=rope_dim_split) + for _ in range(depth) + ] + ) + ) + stride, reduction = upsamples[stage_i] + self.upsamples.append( + LinearPixelShuffleUpsample(in_channels=c, stride=stride, out_channels_reduction_factor=reduction) + ) + + self.t_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=t_emb_dim, size_emb_dim=0) + + c_ctx = stage_channels[-1] + self.context_channels = c_ctx + c5 = stage5_channels if stage5_channels is not None else c_ctx + d5 = stage_depths[-1] + assert c5 % head_dim == 0, f"stage5_channels {c5} must be a multiple of head_dim={head_dim}" + noised_pixel_channels = out_channels * (patch_size**2) + + # Latent-grid floor so stages 1-3 (full volume) never undershoot NA. + self.stage_min_tile_sizes: Tuple[int, int, int] = all_stages_min_tile_size( + stage_kernels, upsamples, stage5_kernel + ) + # Stage-4-input tile floor / overlap halos (only stages 4-5 are tiled). + up3_stride = upsamples[3][0] + self.tile_min_sizes: Tuple[int, int, int] = compute_tile_min_size( + stage_kernels[3], stage5_kernel, up3_stride + ) + self.tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]] = compute_tile_halos( + stage_kernels[3], + stage_depths[3], + stage5_kernel, + stage_depths[-1], + up3_stride, + ) + self.conv_in_x_t = ChannelLinear(noised_pixel_channels, c5, bias=True) + + # Shared AdaLN-Zero (7-chunk for shape compat; gate slots unused in block). + self.shared_adaln = AdaLNZero(dim=c5, t_emb_dim=t_emb_dim) + + self.diff_blocks = nn.ModuleList( + [ + CombinedDiffusionNABlock( + dim=c5, + kernel_size=stage5_kernel, + context_channels=c_ctx, + head_dim=head_dim, + rope_dim_split=rope_dim_split, + ) + for _ in range(d5) + ] + ) + + self.norm_out = nn.RMSNorm(c5, eps=1e-6) + self.conv_out = ChannelLinear(c5, noised_pixel_channels, bias=True) + + self.timestep_scale_multiplier = timestep_scale_multiplier + self.model_output_type = model_output_type + # Set True by ``compile_diffusion_decoder`` so decode marks T/H/W dynamic. + self.mark_dynamic_shapes = False + # When True, skip stage-4 upsample and inject via deferred sequential upsample+proj. + # Default False = combined pathway (``CombinedDiffusionNABlock``). Chunked DiffVAE + # modes flip this via ``apply_diffvae_config``. + self.deferred_stage4_upsample = False + # Remaining temporal upsampling per stage input, plus 1 for stage 5: the divisor in + # ``keyframe_stage_times``. (8, 8, 4, 2, 1) for the production ladder. + self._keyframe_time_strides: Tuple[int, ...] = remaining_time_strides(self.upsamples) + + def _run_det_stage(self, x: torch.Tensor, stage_i: int, drop_leading_frame: bool) -> torch.Tensor: + """One deterministic stage: NA blocks + upsample.""" + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + for block in self.det_stages[stage_i]: + x = block(x) + return self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame) + + def forward_stages_1_to_3( + self, + z_noisy: torch.Tensor, + drop_leading_frame: bool = True, + ) -> torch.Tensor: + """Stages 1-3 on a full (or already ghost-padded) latent → stage-4 input feature. + Output is channels-last ``(B, T, H, W, C)`` at stage-4 input resolution. + Callers that want NATTEN trailing ghosting should pad the latent first via + ``pad_trailing_latent_for_natten_border``. + """ + z_noisy = self.per_channel_statistics.un_normalize(z_noisy) + x = z_noisy.permute(0, 2, 3, 4, 1) + x = self.conv_in(x) + for stage_i in range(3): + x = self._run_det_stage(x, stage_i, drop_leading_frame) + return x + + def _keyframe_stream_from_latents( + self, + keyframes: DecodeKeyframes, + *, + valid: torch.Tensor | None = None, + ) -> KeyframeStream: + """Keyframe latents to a stage-1-input stream: un-normalize, tag, ``conv_in``, mask. + Stages 1-3 always run on the whole volume, so times stay in the global stage-1 + origin. Tile-local rebasing happens later, at stage 4. + Unlike upstream we un-normalize first, because our ``conv_in`` consumes + un-normalized latents (``forward_stages_1_to_3`` does the same for video) while + upstream un-normalizes above the decoder. ``type_emb`` still lands in exactly the + same place: on the latents, immediately before the shared ``conv_in``. + """ + latents = self.per_channel_statistics.un_normalize(keyframes.latents) + x = latents.permute(0, 2, 3, 4, 1) + x = x + self.type_emb.to(dtype=x.dtype, device=x.device).view(1, 1, 1, 1, -1) + x = self.conv_in(x) + planes = x.shape[1] + if valid is None: + valid = torch.ones(planes, dtype=torch.bool, device=x.device) + times = keyframe_clip_times( + keyframes.pixel_frame_indices, + self._keyframe_time_strides[0], + keyframes.clip_start_frame, + ) + return KeyframeStream(x=x, times=times.to(device=x.device), valid=valid.to(device=x.device)).masked() + + def _run_det_stage_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + stage_i: int, + drop_leading_frame: bool, + pixel_frame_indices: torch.Tensor, + next_time_origin: float, + clip_start_frame: int = 0, + ) -> tuple[torch.Tensor, KeyframeStream]: + """One deterministic stage over both streams: joint NA blocks + upsample. + The keyframe upsample is spatial-only and always drops its leading frame; the video + stream's ``drop_leading_frame`` is a tiling property and must not leak into it. + Times are rebuilt from the *next* stage's remaining stride after upsampling. + ``next_time_origin`` is in the **next** stage's temporal units, because that is the + scale the times it rebases are expressed in. Zero everywhere except the stage-4 hop of + a tiled decode, whose next stage is 5 and whose origin is therefore a pixel frame. + ``clip_start_frame`` is the first global pixel of *this* video latent (0 for a full + clip; Dist's tile origin for a slice). + """ + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + for block in self.det_stages[stage_i]: + x, keyframes = block.forward_with_keyframes(x, keyframes) + x = self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame) + keyframe_x = upsample_keyframe_planes(self.upsamples[stage_i], keyframes.x) + next_times = keyframe_clip_times( + pixel_frame_indices, + self._keyframe_time_strides[stage_i + 1], + clip_start_frame, + extra_origin=next_time_origin, + ) + return x, KeyframeStream( + x=keyframe_x, + times=next_times.to(device=keyframe_x.device), + valid=keyframes.valid, + ).masked() + + def forward_stages_1_to_3_with_keyframes( + self, + z_noisy: torch.Tensor, + keyframes: DecodeKeyframes, + drop_leading_frame: bool = True, + *, + keyframe_valid: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, KeyframeStream]: + """Dual-stream stages 1-3: video latent + keyframe planes to stage-4 inputs. + Keyframe counterpart of :meth:`forward_stages_1_to_3`. ``z_noisy`` and + ``keyframes.latents`` must already carry identical spatial padding -- the pad is + applied symmetrically, so padding only one stream would offset every keyframe plane + from the video by half the pad and read as ghosting rather than a failure. + Times are relative to :attr:`DecodeKeyframes.clip_start_frame`. For a full-clip decode + that is 0, so they match global ``t_s``. For a Dist slice they are ``t_s(index) - + t_s(clip_start)`` from stage 1, because this path's "whole volume" *is* the slice. + Additional in-volume tile origins are applied in :meth:`forward_stage_4_with_keyframes`. + """ + keyframes.validate() + if z_noisy.shape[-2:] != keyframes.latents.shape[-2:]: + raise ValueError( + f"keyframe latents must share the video latent's H/W (identical padding), got " + f"{tuple(keyframes.latents.shape[-2:])} vs {tuple(z_noisy.shape[-2:])}" + ) + stream = self._keyframe_stream_from_latents(keyframes, valid=keyframe_valid) + x = self.per_channel_statistics.un_normalize(z_noisy).permute(0, 2, 3, 4, 1) + x = self.conv_in(x) + for stage_i in range(3): + x, stream = self._run_det_stage_with_keyframes( + x, + stream, + stage_i, + drop_leading_frame, + keyframes.pixel_frame_indices, + 0.0, + clip_start_frame=keyframes.clip_start_frame, + ) + return x, stream + + def forward_stage_4_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + pixel_frame_indices: torch.Tensor, + drop_leading_frame: bool = True, + pad_trailing: bool = True, + *, + stage4_time_origin: float = 0.0, + pixel_time_origin: float = 0.0, + clip_start_frame: int = 0, + ) -> tuple[torch.Tensor, KeyframeStream]: + """Dual-stream stage 4 to stage-5 context. Keyframe counterpart of + :meth:`forward_stage_4`. + The ghost-pad crop applies to the video stream only: the trailing replicate is a + temporal-border workaround and the keyframe planes have no temporal extent to pad. + On the deferred (chunked) pathway neither stream is upsampled here: each stage-5 + block folds ``upsamples[3]`` into its own context inject. The returned stream's + ``times`` are nonetheless the *stage-5* times, because that is where they are + consumed -- the same asymmetry the video stream already has, whose returned + ``x`` is a pre-upsample feature rather than stage-5 context. + **Two origins, at two scales.** The video stream's RoPE is tile-local 0-based at every + stage, so keyframe times must be rebased to whatever the tile's frame 0 is -- and this + method spans two different temporal resolutions. ``stage4_time_origin`` is the tile's + start in stage-4 input units (for the blocks); ``pixel_time_origin`` is its first + global pixel frame (for stage 5). They are taken separately from the tile rather than + derived from one another: ``drop_leading_frame`` and the causal first frame make + ``pixel_origin == stride_t * stage4_origin`` an off-by-one trap, not an identity. Both + are 0.0 for an untiled full-clip decode. ``clip_start_frame`` is subtracted in stage + units as well, so a Dist slice whose first pixel is 56 still sees ``t_s(48) - t_s(56)``. + """ + # Rebuild from global indices rather than trusting the caller's stream: stages 1-3 of a + # full-clip decode are global, and Dist has already folded clip_start into clip times. + keyframes = dataclasses.replace( + keyframes, + times=keyframe_clip_times( + pixel_frame_indices, + self._keyframe_time_strides[3], + clip_start_frame, + extra_origin=stage4_time_origin, + ).to(device=keyframes.x.device), + ) + if self.deferred_stage4_upsample: + return self._forward_stage_4_deferred_with_keyframes( + x, keyframes, pixel_frame_indices, pad_trailing, pixel_time_origin, clip_start_frame + ) + x, keyframes = self._run_det_stage_with_keyframes( + x, + keyframes, + 3, + drop_leading_frame, + pixel_frame_indices, + pixel_time_origin, + clip_start_frame=clip_start_frame, + ) + if pad_trailing: + x = crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time, + stage5_kernel_t=self.stage5_kernel[0], + ) + return x, keyframes + + def _forward_stage_4_deferred_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + pixel_frame_indices: torch.Tensor, + pad_trailing: bool, + pixel_time_origin: float, + clip_start_frame: int = 0, + ) -> tuple[torch.Tensor, KeyframeStream]: + """Stage-4 blocks only, both streams, for the deferred (chunked) pathway. + Mirrors :meth:`forward_stage_4`'s deferred branch: no ``upsamples[3]`` on either + stream, ghost cropped at pre-upsample temporal resolution (video only). + """ + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + for block in self.det_stages[3]: + x, keyframes = block.forward_with_keyframes(x, keyframes) + if pad_trailing: + up_t = int(self.upsamples[3].stride[0]) + x = crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time // up_t, + stage5_kernel_t=max(1, -(-self.stage5_kernel[0] // up_t)), + ) + stage5_times = keyframe_clip_times( + pixel_frame_indices, + self._keyframe_time_strides[4], + clip_start_frame, + extra_origin=pixel_time_origin, + ) + return x, KeyframeStream( + x=keyframes.x, + times=stage5_times.to(device=keyframes.x.device), + valid=keyframes.valid, + ).masked() + + def forward_stage_4( + self, + x: torch.Tensor, + drop_leading_frame: bool = True, + pad_trailing: bool = True, + ) -> torch.Tensor: + """Stage 4 on a stage-4-input feature tile → stage-5 context (or pre-upsample feat). + ``x`` is channels-last. When ``pad_trailing``, soft-crop the ghosting + appendix before returning (ghost pad must already be present upstream). + When ``deferred_stage4_upsample`` is set, runs NA blocks only (no + ``upsamples[3]``) and crops ghost at pre-upsample temporal resolution. + """ + if self.deferred_stage4_upsample: + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + for block in self.det_stages[3]: + x = block(x) + if pad_trailing: + up_t = int(self.upsamples[3].stride[0]) + x = crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time // up_t, + stage5_kernel_t=max(1, -(-self.stage5_kernel[0] // up_t)), + ) + return x + + x = self._run_det_stage(x, 3, drop_leading_frame) + if pad_trailing: + x = crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time, + stage5_kernel_t=self.stage5_kernel[0], + ) + return x + + def _context_and_x_for_diff_step(self, context: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor: + """Build block-ready ``[context | conv_in_x_t(patched x)]`` for ``forward_diff_step``.""" + noised_pixels_patched = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1) + x = self.conv_in_x_t(noised_pixels_patched.permute(0, 2, 3, 4, 1)) + return torch.cat([context, x], dim=-1) + + def _keyframe_context_and_x_for_diff_step( + self, + keyframe_context: torch.Tensor, + keyframe_x_t: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> torch.Tensor: + """Keyframe ``[context | conv_in_x_t(patched x)]``, mask-zeroed. + ``keyframe_x_t`` is ``(B, C_pix, P, H_pix, W_pix)`` -- the keyframe planes' own noised + pixels, one pixel frame per plane, through the *shared* ``conv_in_x_t``. + """ + patched = patchify(keyframe_x_t, patch_size_hw=self.patch_size, patch_size_t=1) + x = self.conv_in_x_t(patched.permute(0, 2, 3, 4, 1)) + x = x * keyframe_valid[None, :, None, None, None] + return torch.cat([keyframe_context, x], dim=-1) + + def _x_for_diff_step(self, x_t: torch.Tensor) -> torch.Tensor: + """Conv-processed noised pixels only (deferred-context path).""" + noised_pixels_patched = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1) + return self.conv_in_x_t(noised_pixels_patched.permute(0, 2, 3, 4, 1)) + + def _keyframe_x_for_diff_step(self, keyframe_x_t: torch.Tensor, keyframe_valid: torch.Tensor) -> torch.Tensor: + """Mask-zeroed keyframe noised pixels only (deferred-context path). + Contiguous by construction: the chunked pathway mutates this buffer in place. + """ + patched = patchify(keyframe_x_t, patch_size_hw=self.patch_size, patch_size_t=1) + x = self.conv_in_x_t(patched.permute(0, 2, 3, 4, 1)) + return (x * keyframe_valid[None, :, None, None, None]).contiguous() + + def forward_diff_step( + self, + context_and_x: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + """One stage-5 diffusion step. Returns the model prediction in pixel space. + ``context_and_x`` is ``[latent_context | conv_in_x_t(x)]`` (channels-last), built + at the call site via ``_context_and_x_for_diff_step``. That single buffer is + reused across ``diff_blocks``: each block writes its output x-half back + with ``copy_`` (no per-block ``cat``). One-tensor layout keeps Dynamo + T/H/W symbols identical under ``mark_dynamic``. + """ + x_half = context_and_x[..., self.context_channels :] + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x_half.dtype) + modulation = self.shared_adaln(t_emb) + + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(context_and_x, dim) + + for block in self.diff_blocks: + x_half.copy_(block.forward_combined(context_and_x, modulation)) + return self._pixels_from_stage5(x_half) + + def forward_diff_step_with_keyframes( + self, + context_and_x: torch.Tensor, + keyframe_context_and_x: torch.Tensor, + t: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """One dual-stream stage-5 step. Returns ``(video_pred, keyframe_pred)`` in pixel space. + The keyframe stream at stage 5 is a genuine second *pixel* diffusion stream -- its own + noised pixels through the shared ``conv_in_x_t``, its own per-block + ``context_proj(keyframe_context)``, the same AdaLN modulation -- not a zero tensor and + not the context. It is evolved through the same Euler loop as video so the hidden + state the joint attention reads sits at the noise level it was trained to see, then + discarded: callers use the video prediction only. + Both buffers follow ``forward_diff_step``'s ``[context | x]`` layout and the same + ``copy_``-into-a-view discipline. Video T/H/W stay dynamic; the keyframe plane axis is + specialized, since ``keyframe_times`` / ``keyframe_valid`` pin it. + """ + x_half = context_and_x[..., self.context_channels :] + keyframe_half = keyframe_context_and_x[..., self.context_channels :] + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x_half.dtype) + modulation = self.shared_adaln(t_emb) + + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(context_and_x, dim) + # Keyframe dim 1 is the plane count, not T, and keyframe_times / keyframe_valid pin it. + # Marking it dynamic and then specializing to P raises ConstraintViolationError. + for dim in (2, 3): + torch._dynamo.mark_dynamic(keyframe_context_and_x, dim) + + for block in self.diff_blocks: + x_out, keyframe_out = block.forward_combined_with_keyframes( + context_and_x, + keyframe_context_and_x, + modulation, + keyframe_times, + keyframe_valid, + ) + x_half.copy_(x_out) + keyframe_half.copy_(keyframe_out) + + return self._pixels_from_stage5(x_half), self._pixels_from_stage5(keyframe_half) + + def _pixels_from_stage5(self, x: torch.Tensor) -> torch.Tensor: + """Shared stage-5 tail: ``norm_out`` -> ``conv_out`` -> channels-first -> unpatchify.""" + x = self.norm_out(x) + x = self.conv_out(x) + x = x.permute(0, 4, 1, 2, 3).contiguous() + return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1) + + def forward_diff_step_deferred( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + t: torch.Tensor, + *, + drop_leading_frame: bool = True, + ) -> torch.Tensor: + """Stage-5 step with deferred context: only ``x`` + low-res ``stage4_feat``. + Marks T/H/W dynamic on both tensors. CHUNKED blocks upsample then + ``context_proj`` on the host before attn+mlp; BLACKWELL_DSL + (``DSLDiffusionBlockChain``) folds that hop into the fused kernel and never + materialises full-resolution context. ``drop_leading_frame`` must match the + flag used for this tile's stage-4 path (origin tile vs non-origin). + """ + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype) + modulation = self.shared_adaln(t_emb) + + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + torch._dynamo.mark_dynamic(stage4_feat, dim) + + from ltx_core.model.video_vae.transformer.dsl_kernels import DSLDiffusionBlockChain # noqa: PLC0415 + + if isinstance(self.diff_blocks, DSLDiffusionBlockChain): + # Ping-pong fused launches; same deferred (x, stage4_feat) contract. + x = self.diff_blocks(x, stage4_feat, modulation, drop_leading_frame=drop_leading_frame) + else: + for block in self.diff_blocks: + x = block.forward_x_ctx(x, stage4_feat, modulation, drop_leading_frame=drop_leading_frame) + + x = self.norm_out(x) + x = self.conv_out(x) + x = x.permute(0, 4, 1, 2, 3).contiguous() + return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1) + + def forward_diff_step_deferred_with_keyframes( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + keyframe_x: torch.Tensor, + keyframe_stage4_feat: torch.Tensor, + t: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + *, + drop_leading_frame: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Dual-stream stage-5 step with deferred context. Keyframe counterpart of + :meth:`forward_diff_step_deferred`. + Each stream carries its own pre-upsample stage-4 feature and injects it per block, + so full-resolution context is never materialised for either. ``drop_leading_frame`` + is the video stream's tiling property; the keyframe inject always collapses its + temporal stride. + """ + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype) + modulation = self.shared_adaln(t_emb) + + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + torch._dynamo.mark_dynamic(stage4_feat, dim) + # Plane count, not T -- see the note in the combined path above. + for dim in (2, 3): + torch._dynamo.mark_dynamic(keyframe_x, dim) + torch._dynamo.mark_dynamic(keyframe_stage4_feat, dim) + + # The DSL chain drives itself, so both streams recycle output buffers instead of + # allocating a volume per block per stream; the chunked blocks expose the same + # ``forward_x_ctx_with_keyframes``, so the fallback loop drives those. + from ltx_core.model.video_vae.transformer.dsl_kernels import DSLDiffusionBlockChain # noqa: PLC0415 + + if isinstance(self.diff_blocks, DSLDiffusionBlockChain): + x, keyframe_x = self.diff_blocks.forward_x_ctx_with_keyframes( + x, + stage4_feat, + keyframe_x, + keyframe_stage4_feat, + modulation, + keyframe_times, + keyframe_valid, + drop_leading_frame=drop_leading_frame, + ) + else: + for block in self.diff_blocks: + x, keyframe_x = block.forward_x_ctx_with_keyframes( + x, + stage4_feat, + keyframe_x, + keyframe_stage4_feat, + modulation, + keyframe_times, + keyframe_valid, + drop_leading_frame=drop_leading_frame, + ) + + return self._pixels_from_stage5(x), self._pixels_from_stage5(keyframe_x) + + def _euler_step( + self, x_t: torch.Tensor, model_out: torch.Tensor, t_now: torch.Tensor, t_next: torch.Tensor + ) -> torch.Tensor: + """One reverse-diffusion Euler update: advance ``x_t`` from ``t_now`` to + ``t_next`` given the model's prediction at ``t_now``. + """ + compute_dtype = x_t.dtype + dt = (t_now - t_next).view(-1, *([1] * (x_t.ndim - 1))).to(torch.float32) + x_t_fp32 = x_t.to(torch.float32) + v_pred = model_out if self.model_output_type == "v" else to_velocity(x_t_fp32, t_now, model_out) + return (x_t_fp32 - dt * v_pred).to(compute_dtype) + + def _decode_one_tile( + self, + feat_tile: torch.Tensor, + x_t_tile_init: torch.Tensor, + *, + is_origin: bool, + timestep: torch.Tensor, + pad_trailing: bool, + ) -> torch.Tensor: + """Run stage 4 + diffusion on one stage-4 feature tile (isolation).""" + context_tile = self.forward_stage_4( + feat_tile, + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + ) + + x_t = x_t_tile_init + _, num_steps = timestep.shape + for i in range(num_steps - 1): + t_now = timestep[:, i] + t_next = timestep[:, i + 1] + if self.deferred_stage4_upsample: + x = self._x_for_diff_step(x_t) + model_out = self.forward_diff_step_deferred(x, context_tile, t_now, drop_leading_frame=is_origin).to( + torch.float32 + ) + else: + context_and_x = self._context_and_x_for_diff_step(context_tile, x_t) + model_out = self.forward_diff_step(context_and_x, t_now).to(torch.float32) + x_t = self._euler_step(x_t, model_out, t_now, t_next) + + t_now = timestep[:, -1] + if self.deferred_stage4_upsample: + x = self._x_for_diff_step(x_t) + model_out = self.forward_diff_step_deferred(x, context_tile, t_now, drop_leading_frame=is_origin) + else: + context_and_x = self._context_and_x_for_diff_step(context_tile, x_t) + model_out = self.forward_diff_step(context_and_x, t_now) + if self.model_output_type == "x0": + return model_out + return self._euler_step(x_t, model_out.to(torch.float32), t_now, torch.zeros_like(t_now)) + + def _stage5_canvas_from_context( + self, + context_tile: torch.Tensor, + *, + drop_leading_frame: bool, + ) -> tuple[int, int, int]: + """``(T, H_pix, W_pix)`` of the stage-5 pixel canvas a context tile implies. + On the combined pathway the context is already at stage-5 resolution and + ``_context_and_x_for_diff_step`` patchifies ``x_t`` by ``patch_size``, so the canvas + is just ``(T, H * patch_size, W * patch_size)``. On the deferred pathway the tile is + still pre-upsample: each stage-5 block folds ``upsamples[3]`` into its inject, so + apply that stride here -- including the leading-frame drop the fold performs when the + temporal stride is 2. + """ + t, h, w = context_tile.shape[1], context_tile.shape[2], context_tile.shape[3] + if self.deferred_stage4_upsample: + # Context is still pre-upsample, so this is the same geometry as a stage-4 + # input. Ghost crop already ran, so do not re-apply the kernel-T floor. + return stage5_pixel_shape_from_stage4( + t, + h, + w, + upsample_stride=tuple(self.upsamples[3].stride), # type: ignore[arg-type] + patch_size=self.patch_size, + stage5_kernel_t=self.stage5_kernel[0], + drop_leading_frame=drop_leading_frame, + pad_trailing=False, + ) + return t, h * self.patch_size, w * self.patch_size + + def _decode_one_tile_with_keyframes( # noqa: PLR0913 + self, + feat_tile: torch.Tensor, + keyframes: KeyframeStream, + pixel_frame_indices: torch.Tensor, + *, + is_origin: bool, + timestep: torch.Tensor, + pad_trailing: bool, + generator: torch.Generator | None, + compute_dtype: torch.dtype, + x_t_tile_init: torch.Tensor | None = None, + stage4_time_origin: float = 0.0, + pixel_time_origin: float = 0.0, + clip_start_frame: int = 0, + ) -> torch.Tensor: + """Stage 4 + dual-stream diffusion on one stage-4 feature tile. + Both streams are Euler-stepped together; only the video pixels are returned. The + keyframe pixel stream exists so the hidden state the joint attention reads stays at + the noise level it was trained on, and is discarded here (upstream exposes it only + through its explicit per-step entry points, which the trainer uses). + Noise is sized from the stage-5 context rather than from a re-derived tile geometry; + see :meth:`_stage5_canvas_from_context` for the deferred-pathway correction. The + keyframe canvas differs only in its frame count -- one pixel frame per plane, since + keyframe upsampling collapses its temporal stride. + ``x_t_tile_init`` lets a tiled decode share one global noise field across tiles (edge + policy applied by the caller, as the plain path does); ``None`` draws fresh noise. The + keyframe stream always draws its own -- its planes are not part of the video canvas. + """ + context_tile, keyframes = self.forward_stage_4_with_keyframes( + feat_tile, + keyframes, + pixel_frame_indices, + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + stage4_time_origin=stage4_time_origin, + pixel_time_origin=pixel_time_origin, + clip_start_frame=clip_start_frame, + ) + + batch = context_tile.shape[0] + canvas_t, canvas_h, canvas_w = self._stage5_canvas_from_context(context_tile, drop_leading_frame=is_origin) + randn_device = generator.device if generator is not None else feat_tile.device + + def _noise(frames: int) -> torch.Tensor: + return torch.randn( + (batch, self.out_channels, frames, canvas_h, canvas_w), + dtype=compute_dtype, + generator=generator, + device=randn_device, + ).to(feat_tile.device) + + x_t = _noise(canvas_t) if x_t_tile_init is None else x_t_tile_init + keyframe_x_t = _noise(keyframes.num_planes) + + def _step(t_now: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if self.deferred_stage4_upsample: + return self.forward_diff_step_deferred_with_keyframes( + self._x_for_diff_step(x_t), + context_tile, + self._keyframe_x_for_diff_step(keyframe_x_t, keyframes.valid), + keyframes.x, + t_now, + keyframes.times, + keyframes.valid, + drop_leading_frame=is_origin, + ) + context_and_x = self._context_and_x_for_diff_step(context_tile, x_t) + keyframe_context_and_x = self._keyframe_context_and_x_for_diff_step( + keyframes.x, keyframe_x_t, keyframes.valid + ) + return self.forward_diff_step_with_keyframes( + context_and_x, keyframe_context_and_x, t_now, keyframes.times, keyframes.valid + ) + + _, num_steps = timestep.shape + for i in range(num_steps - 1): + t_now = timestep[:, i] + t_next = timestep[:, i + 1] + video_out, keyframe_out = _step(t_now) + x_t = self._euler_step(x_t, video_out.to(torch.float32), t_now, t_next) + keyframe_x_t = self._euler_step(keyframe_x_t, keyframe_out.to(torch.float32), t_now, t_next) + + t_now = timestep[:, -1] + video_out, _ = _step(t_now) + if self.model_output_type == "x0": + return video_out + return self._euler_step(x_t, video_out.to(torch.float32), t_now, torch.zeros_like(t_now)) + + def _decode_temporal_group_isolated_with_keyframes( # noqa: PLR0913 + self, + tiles: List[Tile], + feat_s4: torch.Tensor, + stream: KeyframeStream, + pixel_frame_indices: torch.Tensor, + content_s4_frames: int, + x_t_init: torch.Tensor | None, + timestep: torch.Tensor, + full_video_shape: VideoLatentShape, + curr_temporal_slice: slice, + generator: torch.Generator | None, + *, + complementary: bool, + clip_start_frame: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Decode one temporal group's tiles with keyframes and blend into a group buffer. + Keyframe counterpart of :meth:`_decode_temporal_group_isolated`. Each tile carries + only the planes near it -- those inside its pixel-frame span plus the nearest plane on + each *side* of it (:func:`planes_for_tile`) -- and the two origins that plane set has + to be rebased against come from the tile, never from each other. + Tile ``out_coords`` are local to this latent. Dist slices keep global + ``pixel_frame_indices`` and pass ``clip_start_frame`` so selection lines up. + """ + group_temporal_len = curr_temporal_slice.stop - curr_temporal_slice.start + group_shape = full_video_shape._replace(frames=group_temporal_len) + full_torch_shape = full_video_shape.to_torch_shape() + accum_dtype = torch.float16 if feat_s4.dtype == torch.bfloat16 else feat_s4.dtype + buffer = torch.zeros(group_shape.to_torch_shape(), device=feat_s4.device, dtype=accum_dtype) + weights: torch.Tensor | None = None if complementary else torch.zeros_like(buffer) + local_temporal_slice = slice(0, group_temporal_len) + + compute_dtype = feat_s4.dtype + up3_stride = tuple(self.upsamples[3].stride) + + for tile_index, tile in enumerate(tiles): + feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( + feat_s4, tile, content_frames=content_s4_frames + ) + # Two origins at two scales -- see forward_stage_4_with_keyframes. + stage4_origin = tile.in_coords[1].indices(content_s4_frames)[0] + pixel_lo, pixel_hi, _ = tile.out_coords[2].indices(full_torch_shape[2]) + + # ``out_coords`` are local to this latent. Dist slices keep global indices and + # ``clip_start_frame`` as the origin, so a local ``[0, 72)`` still has to select + # global ``[56, 127]``. + keep = planes_for_tile(pixel_frame_indices, pixel_lo, pixel_hi - 1, clip_start_frame=clip_start_frame) + if not bool(keep.any()): + raise RuntimeError( + f"tile covering pixel frames [{pixel_lo + clip_start_frame}, " + f"{pixel_hi - 1 + clip_start_frame}] selected no keyframe planes " + f"out of {int(pixel_frame_indices.shape[0])}; planes_for_tile always keeps at least one" + ) + tile_stream = stream.select_planes(keep.to(stream.valid.device)).crop_spatial( + tile.in_coords[2], tile.in_coords[3] + ) + # Not debug-only: the decode below needs this tile's plane positions. + tile_indices = pixel_frame_indices[keep.to(pixel_frame_indices.device)] + if logger.isEnabledFor(logging.INFO): + logger.info( + "keyframe decode: tile %d/%d frames [%d, %d), stage-4 extent %dx%dx%d, %d of %d planes at %s", + tile_index + 1, + len(tiles), + pixel_lo, + pixel_hi, + feat_tile.shape[1], + feat_tile.shape[2], + feat_tile.shape[3], + tile_stream.num_planes, + int(pixel_frame_indices.shape[0]), + tile_indices.tolist(), + ) + + x_t_tile_init: torch.Tensor | None = None + if x_t_init is not None: + stage5_f, stage5_h, stage5_w = stage5_pixel_shape_from_stage4( + content_thw[0], + content_thw[1], + content_thw[2], + upsample_stride=up3_stride, # type: ignore[arg-type] + patch_size=self.patch_size, + stage5_kernel_t=self.stage5_kernel[0], + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + ) + # Same edge policy as the plain path: expand/crop the shared noise field + # rather than drawing fresh noise, since NA mixes padded values inward. + x_t_tile_init = x_t_init[tile.out_coords] + x_t_tile_init, _ = resize_axis(x_t_tile_init, 2, stage5_f, mode="repeat_last") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 3, stage5_h, mode="symmetric") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 4, stage5_w, mode="symmetric") + + pixel_tile = self._decode_one_tile_with_keyframes( + feat_tile, + tile_stream, + tile_indices, + is_origin=is_origin, + timestep=timestep, + pad_trailing=pad_trailing, + generator=generator, + compute_dtype=compute_dtype, + x_t_tile_init=x_t_tile_init, + stage4_time_origin=float(stage4_origin), + pixel_time_origin=float(pixel_lo), + clip_start_frame=clip_start_frame, + ) + content_pixel_shape = pixel_tile_shape(full_torch_shape, tile.out_coords) + pixel_tile = crop_pixels_to_content( + pixel_tile, + content_pixel_shape[2], + content_pixel_shape[3], + content_pixel_shape[4], + ).to(buffer.dtype) + + masks = tuple(m.to(device=buffer.device, dtype=torch.float32) for m in tile.masks_1d) + local_coords = ( + tile.out_coords[0], + tile.out_coords[1], + local_temporal_slice, + tile.out_coords[3], + tile.out_coords[4], + ) + buffer[local_coords] += scale_by_masks_1d(pixel_tile, masks) + if weights is not None: + strength = torch.ones(pixel_tile.shape, device=buffer.device, dtype=buffer.dtype) + weights[local_coords] += scale_by_masks_1d(strength, masks) + + return buffer, weights + + def _decode_groups_with_keyframes( # noqa: PLR0913, PLR0915 + self, + feat_s4: torch.Tensor, + stream: KeyframeStream, + pixel_frame_indices: torch.Tensor, + latent: torch.Tensor, + tiling_config: TilingConfig, + timestep: torch.Tensor, + generator: torch.Generator | None, + *, + content_pixel: VideoLatentShape, + h_pad: AxisPad | None, + w_pad: AxisPad | None, + as_fhwc: bool, + clip_start_frame: int = 0, + ) -> Iterator[torch.Tensor]: + """Tiled keyframe decode, streaming one temporal group at a time. + Same shape as :meth:`_decode_pixels`: only the trailing overlap of the previous group + is retained between iterations, and a group's exclusive frames are yielded before the + next group decodes. That keeps residency at roughly two tile extents rather than a + whole video, which matters more here than on the plain path -- a keyframe decode also + carries a second pixel stream through stage 5. + """ + full_video_shape = ( + VideoLatentShape.from_torch_shape(latent.shape) + .upscale(self.video_downscale_factors) + ._replace(channels=self.out_channels) + ) + target_shape = full_video_shape.to_torch_shape() + strides = [tuple(u.stride) for u in self.upsamples] + s4_t, s4_h, s4_w = stage4_thw_from_latent( + strides, latent.shape[2], latent.shape[3], latent.shape[4], drop_leading_frame=True + ) + tiles = prepare_tile_schedule( + torch.Size([latent.shape[0], latent.shape[1], s4_t, s4_h, s4_w]), + tiling_config, + upsample3_stride=tuple(self.upsamples[3].stride), # type: ignore[arg-type] + patch_size=self.patch_size, + min_tile_size=self.tile_min_sizes, + tile_halos=self.tile_halos, + ) + complementary = masks_are_complementary(tiles, target_shape) + groups = group_tiles_by_temporal_slice(tiles) + group_slices = [slice(*group[0].out_coords[2].indices(target_shape[2])[:2]) for group in groups] + + single_step_x0 = timestep.shape[1] == 1 and self.model_output_type == "x0" + x_t_init: torch.Tensor | None = None + if not single_step_x0: + randn_device = generator.device if generator is not None else latent.device + x_t_init = torch.randn( + tuple(target_shape), dtype=latent.dtype, generator=generator, device=randn_device + ).to(latent.device) + + logger.info( + "keyframe decode: %d tile(s) in %d temporal group(s), %d frames at %dx%d, %d planes", + len(tiles), + len(groups), + content_pixel.frames, + content_pixel.height, + content_pixel.width, + int(pixel_frame_indices.shape[0]), + ) + + scaled_h_pad = scale_axis_pad(h_pad, self.video_downscale_factors.height) + scaled_w_pad = scale_axis_pad(w_pad, self.video_downscale_factors.width) + overlap_stub: torch.Tensor | None = None + overlap_stub_weights: torch.Tensor | None = None + + def _emit(buf: torch.Tensor, wts: torch.Tensor | None, global_start: int) -> torch.Tensor | None: + """Finalize, crop to content, and lay out one emitted run of frames.""" + if global_start >= content_pixel.frames or buf.shape[2] < 1: + return None + frames_keep = min(buf.shape[2], content_pixel.frames - global_start) + if frames_keep < 1: + return None + chunk = buf[:, :, :frames_keep] + if wts is not None: + floor = _weight_floor(wts.dtype) + chunk = chunk / wts[:, :, :frames_keep].clamp(min=floor) + chunk = crop_pixels_to_content( + chunk.to(latent.dtype), + frames_keep, + content_pixel.height, + content_pixel.width, + h_pad=scaled_h_pad, + w_pad=scaled_w_pad, + ) + return chunk[0].permute(1, 2, 3, 0).contiguous() if as_fhwc else chunk + + for group_index, group in enumerate(groups): + curr_temporal_slice = group_slices[group_index] + logger.info( + "keyframe decode: group %d/%d, frames [%d, %d)", + group_index + 1, + len(groups), + curr_temporal_slice.start, + curr_temporal_slice.stop, + ) + buffer, weights = self._decode_temporal_group_isolated_with_keyframes( + group, + feat_s4, + stream, + pixel_frame_indices, + s4_t, + x_t_init, + timestep, + full_video_shape, + curr_temporal_slice, + generator, + complementary=complementary, + clip_start_frame=clip_start_frame, + ) + + if overlap_stub is not None: + overlap_len = int(overlap_stub.shape[2]) + if overlap_len > 0: + overlap_stub += buffer[:, :, :overlap_len] + buffer[:, :, :overlap_len] = overlap_stub + if not complementary: + assert overlap_stub_weights is not None + assert weights is not None + overlap_stub_weights += weights[:, :, :overlap_len] + weights[:, :, :overlap_len] = overlap_stub_weights + overlap_stub = None + overlap_stub_weights = None + + if group_index + 1 < len(groups): + next_start = group_slices[group_index + 1].start + exclusive_len = min(max(0, next_start - curr_temporal_slice.start), buffer.shape[2]) + emitted = _emit( + buffer[:, :, :exclusive_len], + None if weights is None else weights[:, :, :exclusive_len], + curr_temporal_slice.start, + ) + if emitted is not None: + yield emitted + # Retain only the trailing overlap for the next group's handoff. + overlap_stub = buffer[:, :, exclusive_len:].clone() + if not complementary: + assert weights is not None + overlap_stub_weights = weights[:, :, exclusive_len:].clone() + del buffer, weights + else: + emitted = _emit(buffer, weights, curr_temporal_slice.start) + if emitted is not None: + yield emitted + + def _decode_pixels_with_keyframes( + self, + latent: torch.Tensor, + keyframes: DecodeKeyframes, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + *, + as_fhwc: bool = False, + ) -> Iterator[torch.Tensor]: + """Keyframe-aware decode, yielding one chunk of ``(B, C, F, H, W)`` in ``[-1, 1]``. + Stages 1-3 run once on the whole volume for both streams -- so the keyframe stream + reaches stage 4 with *global* times -- then stages 4-5 run per tile with pixel blend. + With ``tiling_config=None`` that is a single tile and the whole thing is one pass. + """ + content_shape = VideoLatentShape.from_torch_shape(latent.shape) + content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + keyframes.validate(num_frames=content_pixel.frames) + + latent, (_t_pad, h_pad, w_pad) = ensure_min_latent_shape(latent, self.stage_min_tile_sizes) + # Same spatial floor for the keyframe planes, with the plane axis pinned by a + # temporal minimum of 1. The pad is symmetric, so padding only one stream would + # offset every plane from the video by half of it. + _min_t, min_h, min_w = self.stage_min_tile_sizes + keyframe_latents, (_, keyframe_h_pad, keyframe_w_pad) = ensure_min_latent_shape( + keyframes.latents, (1, min_h, min_w) + ) + if (keyframe_h_pad, keyframe_w_pad) != (h_pad, w_pad): + raise RuntimeError( + f"keyframe spatial pad {(keyframe_h_pad, keyframe_w_pad)} != video pad {(h_pad, w_pad)}; " + "the two streams must share one spatial origin" + ) + padded_keyframes = dataclasses.replace(keyframes, latents=keyframe_latents) + + # Ghost pad is a temporal-border workaround for the video stream; keyframe planes have + # no temporal extent to pad. The appendix is cropped off context before stage 5. + latent_padded = pad_trailing_latent_for_natten_border( + latent, self._natten_trailing_pad_latent_frames + ) + feat_s4, stream = self.forward_stages_1_to_3_with_keyframes( + latent_padded, padded_keyframes, drop_leading_frame=True + ) + + batch = latent.shape[0] + timestep = self.default_inference_timesteps.to(latent.device).unsqueeze(0).expand(batch, -1) + if tiling_config is not None: + yield from self._decode_groups_with_keyframes( + feat_s4, + stream, + keyframes.pixel_frame_indices, + latent, + tiling_config, + timestep, + generator, + content_pixel=content_pixel, + h_pad=h_pad, + w_pad=w_pad, + as_fhwc=as_fhwc, + clip_start_frame=keyframes.clip_start_frame, + ) + return + + logger.info("keyframe decode: untiled, %d frames, %d planes", content_pixel.frames, stream.num_planes) + pixels = self._decode_one_tile_with_keyframes( + feat_s4, + stream, + keyframes.pixel_frame_indices, + is_origin=True, + timestep=timestep, + pad_trailing=True, + generator=generator, + compute_dtype=latent.dtype, + clip_start_frame=keyframes.clip_start_frame, + ) + pixels = crop_pixels_to_content( + pixels, + content_pixel.frames, + content_pixel.height, + content_pixel.width, + h_pad=scale_axis_pad(h_pad, self.video_downscale_factors.height), + w_pad=scale_axis_pad(w_pad, self.video_downscale_factors.width), + ).to(latent.dtype) + if as_fhwc: + yield pixels[0].permute(1, 2, 3, 0).contiguous() + else: + yield pixels + + def _decode_video_with_keyframes( + self, + latent: torch.Tensor, + keyframes: DecodeKeyframes, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + """Keyframe-aware decode, yielding float chunk(s) ``[f, h, w, c]`` in ``[0, 1]``. + Implementation of :meth:`decode_video` when ``keyframes`` is set. ``keyframes`` carries + single-frame latents plus their global pixel frame indices; every video token then + attends to the nearest planes through a joint neighborhood-attention window. + """ + + def to_rgb(frames: torch.Tensor) -> torch.Tensor: + return frames.add_(1).mul_(0.5).clamp_(0, 1) + + for chunk in self._decode_pixels_with_keyframes( + latent, keyframes, tiling_config, generator=generator, as_fhwc=True + ): + yield to_rgb(chunk) + + def _decode_temporal_group_isolated( + self, + tiles: List[Tile], + feat_s4: torch.Tensor, + content_s4_frames: int, + x_t_init: torch.Tensor | None, + timestep: torch.Tensor, + full_video_shape: VideoLatentShape, + curr_temporal_slice: slice, + generator: torch.Generator | None, + *, + complementary: bool, + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Decode every tile of one temporal group in isolation and blend.""" + group_temporal_len = curr_temporal_slice.stop - curr_temporal_slice.start + group_shape = full_video_shape._replace(frames=group_temporal_len) + full_torch_shape = full_video_shape.to_torch_shape() + accum_dtype = torch.float16 if feat_s4.dtype == torch.bfloat16 else feat_s4.dtype + buffer = torch.zeros(group_shape.to_torch_shape(), device=feat_s4.device, dtype=accum_dtype) + weights: torch.Tensor | None = None if complementary else torch.zeros_like(buffer) + local_temporal_slice = slice(0, group_temporal_len) + + compute_dtype = feat_s4.dtype + randn_device = generator.device if generator is not None else feat_s4.device + up3_stride = tuple(self.upsamples[3].stride) + + for tile in tiles: + feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( + feat_s4, tile, content_frames=content_s4_frames + ) + content_pixel_shape = pixel_tile_shape(full_torch_shape, tile.out_coords) + stage5_f, stage5_h, stage5_w = stage5_pixel_shape_from_stage4( + content_thw[0], + content_thw[1], + content_thw[2], + upsample_stride=up3_stride, # type: ignore[arg-type] + patch_size=self.patch_size, + stage5_kernel_t=self.stage5_kernel[0], + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + ) + + if x_t_init is None: + x_t_tile_init = torch.randn( + (content_pixel_shape[0], content_pixel_shape[1], stage5_f, stage5_h, stage5_w), + dtype=compute_dtype, + generator=generator, + device=randn_device, + ).to(feat_s4.device) + else: + # Expand/crop to stage-5 canvas with the same edge policy as latent + # size-floor / ghost pad (not fresh noise - NA mixes padded values + # into kept pixels near the boundary). + x_t_tile_init = x_t_init[tile.out_coords] + x_t_tile_init, _ = resize_axis(x_t_tile_init, 2, stage5_f, mode="repeat_last") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 3, stage5_h, mode="symmetric") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 4, stage5_w, mode="symmetric") + + pixel_tile = self._decode_one_tile( + feat_tile, + x_t_tile_init, + is_origin=is_origin, + timestep=timestep, + pad_trailing=pad_trailing, + ) + pixel_tile = crop_pixels_to_content( + pixel_tile, + content_pixel_shape[2], + content_pixel_shape[3], + content_pixel_shape[4], + ).to(buffer.dtype) + + masks = tuple(m.to(device=buffer.device, dtype=torch.float32) for m in tile.masks_1d) + local_coords = ( + tile.out_coords[0], + tile.out_coords[1], + local_temporal_slice, + tile.out_coords[3], + tile.out_coords[4], + ) + buffer[local_coords] += scale_by_masks_1d(pixel_tile, masks) + if weights is not None: + strength = torch.ones(pixel_tile.shape, device=buffer.device, dtype=buffer.dtype) + weights[local_coords] += scale_by_masks_1d(strength, masks) + + return buffer, weights + + def _decode_pixels( # noqa: PLR0912, PLR0915 + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + *, + as_fhwc: bool = False, + ) -> Iterator[torch.Tensor]: + """Decode latent to pixels, yielding temporal chunks. + Default yields raw ``(B, C, F, H, W)`` in ``[-1, 1]``. With ``as_fhwc=True`` + (used by :meth:`decode_video`), each chunk is materialized once as + contiguous ``[F, H, W, C]`` still in ``[-1, 1]`` - layout copy only; + range mapping stays in ``to_rgb``. + Stages 1-3 run once on the full volume; stages 4-5 run per tile with + pixel blend (one tile / one group when untiled or no real split). + Across temporal groups only the trailing overlap is retained between + iterations; exclusive frames are yielded before the next group decodes. + Peak residency is ~two tile extents (current buffer + still-live emit / + overlap stub), not a single ``tile + overlap`` slab. + """ + content_shape = VideoLatentShape.from_torch_shape(latent.shape) + content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + + latent, (_t_pad, h_pad, w_pad) = ensure_min_latent_shape(latent, self.stage_min_tile_sizes) + spatial_scale = (self.video_downscale_factors.height, self.video_downscale_factors.width) + work_shape = VideoLatentShape.from_torch_shape(latent.shape) + full_video_shape = work_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + target_shape = full_video_shape.to_torch_shape() + + strides = [tuple(u.stride) for u in self.upsamples] + s4_t, s4_h, s4_w = stage4_thw_from_latent( + strides, latent.shape[2], latent.shape[3], latent.shape[4], drop_leading_frame=True + ) + tiles = prepare_tile_schedule( + torch.Size([latent.shape[0], latent.shape[1], s4_t, s4_h, s4_w]), + tiling_config, + upsample3_stride=tuple(self.upsamples[3].stride), # type: ignore[arg-type] + patch_size=self.patch_size, + min_tile_size=self.tile_min_sizes, + tile_halos=self.tile_halos, + ) + + latent_padded = pad_trailing_latent_for_natten_border( + latent, self._natten_trailing_pad_latent_frames + ) + if self.mark_dynamic_shapes: + for dim in (2, 3, 4): + torch._dynamo.mark_dynamic(latent_padded, dim) + + feat_s4 = self.forward_stages_1_to_3(latent_padded, drop_leading_frame=True) + + batch = latent.shape[0] + timestep = self.default_inference_timesteps.to(latent.device).unsqueeze(0).expand(batch, -1) + single_step_x0 = timestep.shape[1] == 1 and self.model_output_type == "x0" + + x_t_init: torch.Tensor | None = None + if not single_step_x0: + compute_dtype = latent.dtype + randn_device = generator.device if generator is not None else latent.device + x_t_init = torch.randn( + tuple(target_shape), dtype=compute_dtype, generator=generator, device=randn_device + ).to(latent.device) + + complementary = masks_are_complementary(tiles, target_shape) + groups = group_tiles_by_temporal_slice(tiles) + group_slices = [slice(*group[0].out_coords[2].indices(target_shape[2])[:2]) for group in groups] + + # Keep only the trailing temporal overlap of the previous group (not the full + # chunk). Exclusive frames are yielded before the next group is decoded; the + # consumer may still hold that emit while the next buffer is live (~2x tile). + overlap_stub: torch.Tensor | None = None + overlap_stub_weights: torch.Tensor | None = None + + def _finalize(buf: torch.Tensor, wts: torch.Tensor | None) -> torch.Tensor: + if complementary: + return buf.to(latent.dtype) + assert wts is not None + wts = wts.clamp(min=_weight_floor(wts.dtype)) + return (buf / wts).to(latent.dtype) + + def _narrow_content_cfhw(t: torch.Tensor, frames_keep: int) -> torch.Tensor: + """Spatial/temporal content crop as views (no ``.contiguous()``).""" + x = t[:, :, :frames_keep] + th, tw = content_pixel.height, content_pixel.width + scale_h, scale_w = spatial_scale + if h_pad is not None: + before = scale_axis_pad(h_pad, scale_h).before + x = x.narrow(3, before, th) + else: + need = x.shape[3] - th + if need > 0: + x = x.narrow(3, need // 2, th) + elif need < 0: + x, _ = resize_axis(x, 3, th, mode="symmetric") + if w_pad is not None: + before = scale_axis_pad(w_pad, scale_w).before + x = x.narrow(4, before, tw) + else: + need = x.shape[4] - tw + if need > 0: + x = x.narrow(4, need // 2, tw) + elif need < 0: + x, _ = resize_axis(x, 4, tw, mode="symmetric") + return x + + def _crop_emit(buf: torch.Tensor, wts: torch.Tensor | None, global_start: int) -> torch.Tensor | None: + if global_start >= content_pixel.frames or buf.shape[2] < 1: + return None + frames_keep = min(buf.shape[2], content_pixel.frames - global_start) + if frames_keep < 1: + return None + if not as_fhwc: + chunk = _finalize(buf[:, :, :frames_keep], None if wts is None else wts[:, :, :frames_keep]) + return crop_pixels_to_content( + chunk, + frames_keep, + content_pixel.height, + content_pixel.width, + h_pad=h_pad, + w_pad=w_pad, + spatial_scale=spatial_scale, + ) + + # One materialize: contiguous FHWC in latent.dtype, still [-1, 1]. + # CFHW→FHWC cannot be inplace; range mapping is left to to_rgb. + cfhw = _narrow_content_cfhw(buf, frames_keep) + src = cfhw[0] # C, F, H, W (view into accumulator) + video = torch.empty( + src.shape[1], + src.shape[2], + src.shape[3], + src.shape[0], + dtype=latent.dtype, + device=src.device, + ) + video.copy_(src.permute(1, 2, 3, 0)) + if not complementary: + assert wts is not None + w_cfhw = _narrow_content_cfhw(wts, frames_keep) + wview = w_cfhw[0].permute(1, 2, 3, 0) + # Inplace floor on exclusive weight region only (discarded after emit). + wview.clamp_min_(_weight_floor(w_cfhw.dtype)) + video.div_(wview) + return video + + for gi, group in enumerate(groups): + curr_temporal_slice = group_slices[gi] + buffer, weights = self._decode_temporal_group_isolated( + group, + feat_s4, + s4_t, + x_t_init, + timestep, + full_video_shape, + curr_temporal_slice, + generator=generator, + complementary=complementary, + ) + + if overlap_stub is not None: + overlap_len = int(overlap_stub.shape[2]) + if overlap_len > 0: + # Stub is exactly the region overlapping this group (cloned when + # the previous group finished); blend then write back into buffer. + overlap_stub += buffer[:, :, :overlap_len] + if complementary: + buffer[:, :, :overlap_len] = overlap_stub + else: + assert overlap_stub_weights is not None + assert weights is not None + overlap_stub_weights += weights[:, :, :overlap_len] + buffer[:, :, :overlap_len] = overlap_stub + weights[:, :, :overlap_len] = overlap_stub_weights + overlap_stub = None + overlap_stub_weights = None + + if gi + 1 < len(groups): + next_start = group_slices[gi + 1].start + exclusive_len = min(max(0, next_start - curr_temporal_slice.start), buffer.shape[2]) + emitted = _crop_emit( + buffer[:, :, :exclusive_len], + None if weights is None else weights[:, :, :exclusive_len], + curr_temporal_slice.start, + ) + if emitted is not None: + yield emitted + # Retain only the trailing overlap for the next handoff. + overlap_stub = buffer[:, :, exclusive_len:].clone() + if not complementary: + assert weights is not None + overlap_stub_weights = weights[:, :, exclusive_len:].clone() + del buffer, weights + else: + emitted = _crop_emit(buffer, weights, curr_temporal_slice.start) + if emitted is not None: + yield emitted + + def forward( + self, + sample: torch.Tensor, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Decode via ``_decode_pixels`` with ``tiling_config=None`` (single full tile).""" + return next(self._decode_pixels(sample, tiling_config=None, generator=generator)) + + def tiled_decode( + self, + latent: torch.Tensor, + tiling_config: TilingConfig, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + """Tiled decode: stages 1-3 once, stages 4-5 per tile, pixel blend.""" + yield from self._decode_pixels(latent, tiling_config, generator=generator) + + def decode_video( + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + *, + keyframes: DecodeKeyframes | None = None, + ) -> Iterator[torch.Tensor]: + """Decode latent video, yielding float chunk(s) ``[f, h, w, c]`` in ``[0, 1]``. + Untiled and tiled both go through ``_decode_pixels``. Tiled decode may yield + multiple times when ``tiling_config.frames`` splits the video. + With ``keyframes`` this is :meth:`_decode_video_with_keyframes`; the argument exists on + every decoder so a caller can pass planes without first asking which VAE it holds. + Layout is packed once to contiguous FHWC on emit; ``to_rgb`` only does + inplace ``[-1, 1]→[0, 1]`` (no second realloc). + """ + if keyframes is not None: + yield from self._decode_video_with_keyframes(latent, keyframes, tiling_config, generator) + return + + def to_rgb(frames: torch.Tensor) -> torch.Tensor: + return frames.add_(1).mul_(0.5).clamp_(0, 1) + + for chunk in self._decode_pixels(latent, tiling_config, generator=generator, as_fhwc=True): + yield to_rgb(chunk) + + def decode_single_frames( + self, + latents: Sequence[torch.Tensor], + generator: torch.Generator | Sequence[torch.Generator | None] | None = None, + ) -> Iterator[torch.Tensor]: + """Decode each latent as its own one-frame clip, yielding one RGB tensor per latent.""" + yield from iter_decoded_single_frames(self, latents, generator) + + + +class LTX25DiffusionVideoDecoder(DiffusionVideoDecoder): + """DiffSynth-facing decoder with one full/tiled/keyframe interface.""" + + def forward(self, sample, generator=None, keyframes=None): + return self.decode(sample, generator=generator, keyframes=keyframes) + + def auto_tiling_config(self, latent, keyframes=None): + pixel_shape = ( + VideoLatentShape.from_torch_shape(latent.shape) + .upscale(self.video_downscale_factors) + ._replace(channels=self.out_channels) + ) + device = latent.device + if device.type == "cuda": + # Cached allocator blocks from a previous decode would otherwise make the + # free-memory query report a budget of zero for back-to-back decodes. + torch.cuda.empty_cache() + free_bytes = torch.cuda.mem_get_info(device.index)[0] + else: + free_bytes = 0 + if free_bytes <= 0: + return None + + # Budget estimate must not read weight dtype/device: parameters may be meta + # or disk-offloaded here, so assume bf16 storage for the footprint estimate. + model_bytes = sum(parameter.numel() for parameter in self.parameters()) * 2 + upsample_strides = [tuple(upsample.stride) for upsample in self.upsamples] + element_size = accumulator_element_size(latent.dtype) + return recommended_decode_tiling_config( + tile_halos=self.tile_halos, + pixel_scale=stage4_to_pixel_scale_factors(upsample_strides[3], self.patch_size), + min_tile_size_s4=self.tile_min_sizes, + patch_size=self.patch_size, + height=pixel_shape.height, + width=pixel_shape.width, + num_frames=pixel_shape.frames, + mode=DiffVAEMode.CHUNKED_EAGER, + free_bytes=free_bytes, + stage5_channels=self.stage_channels[-1], + stage4_channels=self.stage_channels[3], + upsample_strides=upsample_strides, + model_bytes=model_bytes, + element_size=element_size, + natten_trailing_pad_latent_frames=self._natten_trailing_pad_latent_frames, + keyframes=keyframes is not None, + ) + + def decode( + self, + latent, + tiled=False, + tile_size_in_pixels=None, + tile_overlap_in_pixels=None, + tile_size_in_frames=None, + tile_overlap_in_frames=None, + generator=None, + keyframes=None, + ): + tiling_config = None + if tiled is True or tiled is AUTO_TILING: + if tiled is AUTO_TILING or tile_size_in_pixels is None or tile_overlap_in_pixels is None or tile_size_in_frames is None or tile_overlap_in_frames is None: + tiling_config = self.auto_tiling_config(latent, keyframes=keyframes) + if tiling_config is None: + raise ValueError("Automatic DiffVAE tiling requires a CUDA device with queryable free memory.") + else: + if isinstance(tile_size_in_pixels, Sequence) and not isinstance(tile_size_in_pixels, (str, bytes)): + tile_height, tile_width = tile_size_in_pixels + else: + tile_height = tile_width = int(tile_size_in_pixels) + if isinstance(tile_overlap_in_pixels, Sequence) and not isinstance(tile_overlap_in_pixels, (str, bytes)): + overlap_height, overlap_width = tile_overlap_in_pixels + else: + overlap_height = overlap_width = int(tile_overlap_in_pixels) + tiling_config = TileSizeConfig( + frames=DimensionSizeConfig(int(tile_size_in_frames), int(tile_overlap_in_frames)), + height=DimensionSizeConfig(int(tile_height), int(overlap_height)), + width=DimensionSizeConfig(int(tile_width), int(overlap_width)), + ) + iterator = ( + self._decode_pixels_with_keyframes(latent, keyframes, tiling_config, generator=generator) + if keyframes is not None + else self._decode_pixels(latent, tiling_config, generator=generator) + ) + chunks = list(iterator) + if not chunks: + raise RuntimeError("Diffusion decoder produced no output chunks") + return torch.cat(chunks, dim=2) + diff --git a/diffsynth/models/ltx25_diffvae/NOTICE.md b/diffsynth/models/ltx25_diffvae/NOTICE.md deleted file mode 100644 index 34a20f287..000000000 --- a/diffsynth/models/ltx25_diffvae/NOTICE.md +++ /dev/null @@ -1,9 +0,0 @@ -# LTX-2.5 DiffVAE eager port notice - -The Python sources in this directory are a dependency-closed, in-tree port of -selected `ltx_core` DiffVAE decoder sources from `Lightricks/LTX-2` revision -`400fd31054597515f47125691032c04b1c3ee24e`. Original source headers are -retained. The port deliberately excludes NATTEN, Triton, Blackwell DSL, and -runtime imports from the installed `ltx_core` package. It uses the upstream -pure-PyTorch eager tiled-SDPA implementation as its neighborhood-attention -backend. diff --git a/diffsynth/models/ltx25_diffvae/__init__.py b/diffsynth/models/ltx25_diffvae/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/diffsynth/models/ltx25_diffvae/model/__init__.py b/diffsynth/models/ltx25_diffvae/model/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/diffsynth/models/ltx25_diffvae/model/transformer/__init__.py b/diffsynth/models/ltx25_diffvae/model/transformer/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py b/diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py deleted file mode 100644 index 87d9845ca..000000000 --- a/diffsynth/models/ltx25_diffvae/model/transformer/timestep_embedding.py +++ /dev/null @@ -1,115 +0,0 @@ -import math - -import torch - - -def get_timestep_embedding( - timesteps: torch.Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -) -> torch.Tensor: - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) - exponent = exponent / (half_dim - downscale_freq_shift) - - emb = torch.exp(exponent) - emb = timesteps[:, None].float() * emb[None, :] - - emb = scale * emb - - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) - - if flip_sin_to_cos: - emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) - - if embedding_dim % 2 == 1: - emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(torch.nn.Module): - def __init__( - self, - in_channels: int, - time_embed_dim: int, - out_dim: int | None = None, - post_act_fn: str | None = None, - cond_proj_dim: int | None = None, - sample_proj_bias: bool = True, - ): - super().__init__() - - self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) - - if cond_proj_dim is not None: - self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) - else: - self.cond_proj = None - - self.act = torch.nn.SiLU() - time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim - - self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) - - if post_act_fn is None: - self.post_act = None - - def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class Timesteps(torch.nn.Module): - def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.scale = scale - - def forward(self, timesteps: torch.Tensor) -> torch.Tensor: - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - scale=self.scale, - ) - return t_emb - - -class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): - def __init__( - self, - embedding_dim: int, - size_emb_dim: int, - ): - super().__init__() - - self.outdim = size_emb_dim - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - def forward( - self, - timestep: torch.Tensor, - hidden_dtype: torch.dtype, - ) -> torch.Tensor: - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) - return timesteps_emb diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py b/diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py deleted file mode 100644 index 1b979e5a5..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/diffusion_video_decoder.py +++ /dev/null @@ -1,208 +0,0 @@ -from __future__ import annotations - -from typing import Literal - -import torch -from torch import nn - -from diffsynth.models.ltx25_diffvae.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings -from diffsynth.models.ltx25_diffvae.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer import ( - AdaLNZero, - ChannelLinear, - CombinedDiffusionNABlock, - LinearPixelShuffleUpsample, - NABlock, -) - - -class DiffusionVideoDecoder(nn.Module): - def __init__( - self, - in_channels: int = 128, - out_channels: int = 3, - patch_size: int = 4, - head_dim: int = 64, - rope_dim_split: tuple[int, int, int] | None = None, - stage_channels: tuple[int, ...] = (1024, 512, 256, 256, 128), - stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8), - stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5), (3, 7, 7)), - upsamples: tuple[tuple[tuple[int, int, int], int], ...] = (((1, 2, 2), 2), ((2, 1, 1), 2), ((2, 2, 2), 1), ((2, 2, 2), 2)), - stage5_kernel: tuple[int, int, int] | None = None, - stage5_channels: int | None = None, - t_emb_dim: int = 384, - default_num_inference_steps: int = 1, - timestep_scale_multiplier: float = 1.0, - model_output_type: Literal["v", "x0"] = "x0", - ) -> None: - super().__init__() - if len(stage_channels) != len(stage_depths) or len(stage_channels) != len(stage_kernels): - raise ValueError("stage_channels, stage_depths, and stage_kernels must have the same length") - if len(upsamples) != len(stage_channels) - 1: - raise ValueError("one fewer upsample than decoder stages is required") - if any(channels % head_dim for channels in stage_channels): - raise ValueError("every stage channel count must be divisible by head_dim") - - self.patch_size = patch_size - self.out_channels = out_channels - self.stage_kernels = stage_kernels - self.upsample_strides = tuple(stride for stride, _ in upsamples) - self.stage5_kernel = tuple(stage5_kernel or stage_kernels[-1]) - self.model_output_type = model_output_type - self.timestep_scale_multiplier = timestep_scale_multiplier - self.register_buffer( - "default_inference_timesteps", - torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps), - persistent=False, - ) - - self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) - self.conv_in = ChannelLinear(in_channels, stage_channels[0], bias=True) - self.det_stages = nn.ModuleList() - self.upsamples = nn.ModuleList() - for channels, depth, kernel, (stride, reduction) in zip( - stage_channels[:-1], stage_depths[:-1], stage_kernels[:-1], upsamples, strict=True - ): - self.det_stages.append( - nn.ModuleList( - [NABlock(channels, kernel, head_dim=head_dim, rope_dim_split=rope_dim_split) for _ in range(depth)] - ) - ) - self.upsamples.append(LinearPixelShuffleUpsample(channels, stride, reduction)) - - context_channels = stage_channels[-1] - diffusion_channels = stage5_channels or context_channels - if diffusion_channels % head_dim: - raise ValueError("stage5_channels must be divisible by head_dim") - self.context_channels = context_channels - self.t_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=t_emb_dim, size_emb_dim=0) - self.conv_in_x_t = ChannelLinear(out_channels * patch_size**2, diffusion_channels, bias=True) - self.shared_adaln = AdaLNZero(dim=diffusion_channels, t_emb_dim=t_emb_dim) - self.diff_blocks = nn.ModuleList( - [ - CombinedDiffusionNABlock( - diffusion_channels, - self.stage5_kernel, - context_channels, - head_dim=head_dim, - rope_dim_split=rope_dim_split, - ) - for _ in range(stage_depths[-1]) - ] - ) - self.norm_out = nn.RMSNorm(diffusion_channels, eps=1e-6) - self.conv_out = ChannelLinear(diffusion_channels, out_channels * patch_size**2, bias=True) - - self.min_latent_shape = self._minimum_latent_shape() - - self._trailing_latent_frames = (stage_kernels[0][0] // 2) * 2 - - def _minimum_latent_shape(self) -> tuple[int, int, int]: - cumulative = [1, 1, 1] - minimum = [1, 1, 1] - for kernel, stride in zip(self.stage_kernels[:-1], self.upsample_strides, strict=True): - for axis in range(3): - minimum[axis] = max(minimum[axis], -(-kernel[axis] // cumulative[axis])) - cumulative[axis] *= stride[axis] - for axis in range(3): - minimum[axis] = max(minimum[axis], -(-self.stage5_kernel[axis] // cumulative[axis])) - return tuple(minimum) - - @staticmethod - def _pad_axis(x: torch.Tensor, axis: int, size: int, *, trailing_only: bool) -> tuple[torch.Tensor, int]: - current = x.shape[axis] - if current >= size: - return x, 0 - missing = size - current - before = 0 if trailing_only else missing // 2 - after = missing - before - pieces = [] - if before: - pieces.append(x.narrow(axis, 0, 1).expand(*x.shape[:axis], before, *x.shape[axis + 1 :])) - pieces.append(x) - if after: - pieces.append(x.narrow(axis, current - 1, 1).expand(*x.shape[:axis], after, *x.shape[axis + 1 :])) - return torch.cat(pieces, dim=axis), before - - def _pad_to_minimum(self, latent: torch.Tensor) -> tuple[torch.Tensor, int, int]: - latent, _ = self._pad_axis(latent, 2, self.min_latent_shape[0], trailing_only=True) - latent, h_before = self._pad_axis(latent, 3, self.min_latent_shape[1], trailing_only=False) - latent, w_before = self._pad_axis(latent, 4, self.min_latent_shape[2], trailing_only=False) - return latent, h_before, w_before - - def _run_det_stage(self, x: torch.Tensor, stage_index: int, drop_leading_frame: bool) -> torch.Tensor: - for block in self.det_stages[stage_index]: - x = block(x) - return self.upsamples[stage_index](x, drop_leading_frame=drop_leading_frame) - - def forward_stages_1_to_3(self, z_noisy: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: - x = self.per_channel_statistics.un_normalize(z_noisy).permute(0, 2, 3, 4, 1) - x = self.conv_in(x) - for stage_index in range(3): - x = self._run_det_stage(x, stage_index, drop_leading_frame) - return x - - def forward_stage_4( - self, - x: torch.Tensor, - drop_leading_frame: bool = True, - pad_trailing: bool = True, - ) -> torch.Tensor: - x = self._run_det_stage(x, 3, drop_leading_frame) - if pad_trailing and self._trailing_latent_frames: - ghost_frames = self._trailing_latent_frames * 8 - keep = min(x.shape[1], max(x.shape[1] - ghost_frames, self.stage5_kernel[0])) - x = x[:, :keep] - return x - - def _context_and_x_for_diff_step(self, context: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor: - pixels = patchify(x_t, patch_size_hw=self.patch_size).permute(0, 2, 3, 4, 1) - return torch.cat([context, self.conv_in_x_t(pixels)], dim=-1) - - def forward_diff_step(self, context_and_x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: - x = context_and_x[..., self.context_channels :] - modulation = self.shared_adaln(self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype)) - for block in self.diff_blocks: - x = block(context_and_x, modulation) - context_and_x[..., self.context_channels :].copy_(x) - x = self.conv_out(self.norm_out(x)).permute(0, 4, 1, 2, 3).contiguous() - return unpatchify(x, patch_size_hw=self.patch_size) - - def _euler_step(self, x_t: torch.Tensor, model_out: torch.Tensor, t_now: torch.Tensor, t_next: torch.Tensor) -> torch.Tensor: - dt = (t_now - t_next).view(-1, *([1] * (x_t.ndim - 1))).to(torch.float32) - if self.model_output_type == "x0": - model_out = (x_t.to(torch.float32) - model_out.to(torch.float32)) / t_now.view( - -1, *([1] * (x_t.ndim - 1)) - ).to(torch.float32) - return (x_t.to(torch.float32) - dt * model_out.to(torch.float32)).to(x_t.dtype) - - def forward(self, sample: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: - frames = (sample.shape[2] - 1) * 8 + 1 - height, width = sample.shape[3] * 32, sample.shape[4] * 32 - latent, h_before, w_before = self._pad_to_minimum(sample) - trailing = latent[:, :, -1:].expand(-1, -1, self._trailing_latent_frames, -1, -1) - context = self.forward_stages_1_to_3(torch.cat([latent, trailing], dim=2)) - context = self.forward_stage_4(context) - - pixel_shape = (sample.shape[0], self.out_channels, context.shape[1], context.shape[2] * self.patch_size, context.shape[3] * self.patch_size) - x_t = torch.randn(pixel_shape, dtype=sample.dtype, device=sample.device, generator=generator) - timesteps = self.default_inference_timesteps.to(sample.device).expand(sample.shape[0], -1) - for index in range(timesteps.shape[1] - 1): - prediction = self.forward_diff_step(self._context_and_x_for_diff_step(context, x_t), timesteps[:, index]) - x_t = self._euler_step(x_t, prediction, timesteps[:, index], timesteps[:, index + 1]) - prediction = self.forward_diff_step(self._context_and_x_for_diff_step(context, x_t), timesteps[:, -1]) - pixels = prediction if self.model_output_type == "x0" else self._euler_step(x_t, prediction, timesteps[:, -1], torch.zeros_like(timesteps[:, -1])) - return pixels[:, :, :frames, h_before * 32 : h_before * 32 + height, w_before * 32 : w_before * 32 + width].contiguous() - - def decode( - self, - latent: torch.Tensor, - tiled: bool = True, - tile_size_in_pixels: int = 512, - tile_overlap_in_pixels: int = 128, - tile_size_in_frames: int = 128, - tile_overlap_in_frames: int = 24, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - del tiled, tile_size_in_pixels, tile_overlap_in_pixels, tile_size_in_frames, tile_overlap_in_frames - return self.forward(latent, generator=generator) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/ops.py b/diffsynth/models/ltx25_diffvae/model/video_vae/ops.py deleted file mode 100644 index f945dfa79..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/ops.py +++ /dev/null @@ -1,57 +0,0 @@ -import torch -from einops import rearrange -from torch import nn - - -def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: - if patch_size_hw == 1 and patch_size_t == 1: - return x - if x.dim() == 4: - x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) - elif x.dim() == 5: - x = rearrange( - x, - "b c (f p) (h q) (w r) -> b (c p r q) f h w", - p=patch_size_t, - q=patch_size_hw, - r=patch_size_hw, - ) - else: - raise ValueError(f"Invalid input shape: {x.shape}") - - return x - - -def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: - if patch_size_hw == 1 and patch_size_t == 1: - return x - - if x.dim() == 4: - x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) - elif x.dim() == 5: - x = rearrange( - x, - "b (c p r q) f h w -> b c (f p) (h q) (w r)", - p=patch_size_t, - q=patch_size_hw, - r=patch_size_hw, - ) - - return x - - -class PerChannelStatistics(nn.Module): - def __init__(self, latent_channels: int = 128): - super().__init__() - self.register_buffer("std-of-means", torch.ones(latent_channels)) - self.register_buffer("mean-of-means", torch.zeros(latent_channels)) - - def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( - 1, -1, 1, 1, 1 - ).to(x) - - def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( - 1, -1, 1, 1, 1 - ).to(x) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py deleted file mode 100644 index 1f9d3b95d..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .blocks import NABlock -from .combined.block import CombinedDiffusionNABlock -from .layers import AdaLNZero, ChannelLinear, LinearPixelShuffleUpsample diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py deleted file mode 100644 index e52d0c304..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/attention.py +++ /dev/null @@ -1,59 +0,0 @@ -import torch -from torch import nn - -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.det_attn_rope import det_qkv_rope -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.fallback_na import EagerSdpaAttention -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.qkv import QKVProjections -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.rope_math import ( - DEFAULT_ABS_ROPE_NUM_TILES, - default_rope_dim_split, - rope_inv_freqs, -) - - -class NeighborhoodAttention3D(nn.Module): - def __init__( - self, - dim: int, - kernel_size: tuple[int, int, int], - head_dim: int = 64, - rope_dim_split: tuple[int, int, int] | None = None, - rope_base: float = 10000.0, - ) -> None: - super().__init__() - if dim % head_dim: - raise ValueError(f"dim={dim} must be divisible by head_dim={head_dim}") - rope_dim_split = rope_dim_split or default_rope_dim_split(head_dim) - if sum(rope_dim_split) != head_dim: - raise ValueError(f"rope_dim_split={rope_dim_split} must sum to head_dim={head_dim}") - - self.dim = dim - self.num_heads = dim // head_dim - self.head_dim = head_dim - self.kernel_size = tuple(kernel_size) - self.scale = head_dim**-0.5 - self.rope_dim_split = rope_dim_split - self.rope_num_tiles = DEFAULT_ABS_ROPE_NUM_TILES - self.rope_compute_dtype = torch.float32 - self.attention_function = EagerSdpaAttention() - self.register_buffer("rope_inv_t", rope_inv_freqs(rope_dim_split[0], rope_base), persistent=False) - self.register_buffer("rope_inv_h", rope_inv_freqs(rope_dim_split[1], rope_base), persistent=False) - self.register_buffer("rope_inv_w", rope_inv_freqs(rope_dim_split[2], rope_base), persistent=False) - self.qkv = QKVProjections(dim) - self.proj = nn.Linear(dim, dim, bias=True) - self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) - self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) - - def project_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - batch, frames, height, width, _ = x.shape - q, k, v = self.qkv(x) - shape = (batch, frames, height, width, self.num_heads, self.head_dim) - return q.view(shape), k.view(shape), v.view(shape) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - batch, frames, height, width, _ = x.shape - if any(size < kernel for size, kernel in zip((frames, height, width), self.kernel_size, strict=True)): - raise ValueError(f"input {(frames, height, width)} is smaller than neighborhood kernel {self.kernel_size}") - q, k, v = det_qkv_rope(self, x) - output = self.attention_function(self, q.contiguous(), k.contiguous(), v.contiguous()) - return self.proj(output.reshape(batch, frames, height, width, self.dim)) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py deleted file mode 100644 index 926f124bc..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/blocks.py +++ /dev/null @@ -1,66 +0,0 @@ -from __future__ import annotations - -import torch -from torch import nn - -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.attention import NeighborhoodAttention3D -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.layers import AdaLNZero -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.swiglu import SwiGLU, plain_mlp - -__all__ = [ - "DiffusionNABlock", - "NABlock", -] - - -class NABlock(nn.Module): - def __init__( - self, - dim: int, - kernel_size: tuple[int, int, int], - head_dim: int = 64, - mlp_ratio: float = 4.0, - rope_dim_split: tuple[int, int, int] | None = None, - ) -> None: - super().__init__() - self.norm1 = nn.RMSNorm(dim, eps=1e-6) - self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) - self.norm2 = nn.RMSNorm(dim, eps=1e-6) - hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 - self.mlp = SwiGLU(dim, hidden) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = x + self.attn(self.norm1(x)) - x = plain_mlp(x, self.mlp, self.norm2) - return x - - -class DiffusionNABlock(nn.Module): - def __init__( - self, - dim: int, - kernel_size: tuple[int, int, int], - context_channels: int, - head_dim: int = 64, - mlp_ratio: float = 4.0, - rope_dim_split: tuple[int, int, int] | None = None, - ) -> None: - super().__init__() - self.context_channels = context_channels - self.context_proj = nn.Linear(context_channels, dim, bias=True) - self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim)) - - self.norm1 = nn.RMSNorm(dim, eps=1e-6) - self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) - self.norm2 = nn.RMSNorm(dim, eps=1e-6) - hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 - self.mlp = SwiGLU(dim, hidden) - self.attn.proj.reset_parameters() - - def _modulation( - self, modulation: tuple[torch.Tensor, ...] - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ - modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) - ] - return scale_msa, shift_msa, scale_mlp, shift_mlp diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py deleted file mode 100644 index a38ec13af..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/attn.py +++ /dev/null @@ -1,114 +0,0 @@ -from __future__ import annotations - -import torch -from torch import nn - -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.attention import NeighborhoodAttention3D -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.rope_math import ( - h_positions, - rot_abs_axis_impl, - t_positions, -) - - -_rot_abs_axis = torch.compiler.nested_compile_region(rot_abs_axis_impl) - - -def _apply_nested_abs_rope_slab( - x: torch.Tensor, - rope_split: tuple[int, int, int], - inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - *, - w_pos: torch.Tensor, - compute_dtype: torch.dtype, -) -> torch.Tensor: - d_t, d_h, _ = rope_split - inv_t, inv_h, inv_w = inv_freqs - t = x.shape[1] - h = x.shape[2] - xt = _rot_abs_axis(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) - xh = _rot_abs_axis( - x[..., d_t : d_t + d_h], - h_positions(h, x.device), - inv_h, - axis=2, - compute_dtype=compute_dtype, - ) - xw = _rot_abs_axis(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) - return torch.cat([xt, xh, xw], dim=-1) - - -def _apply_nested_full_volume_rope( - x: torch.Tensor, - rope_split: tuple[int, int, int], - inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - *, - num_tiles: int, - compute_dtype: torch.dtype, -) -> torch.Tensor: - slabs = torch.chunk(x, num_tiles, dim=3) - w_off = 0 - parts: list[torch.Tensor] = [] - for slab in slabs: - w_slab = slab.shape[3] - w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off - parts.append( - _apply_nested_abs_rope_slab( - slab, - rope_split, - inv_freqs, - w_pos=w_pos, - compute_dtype=compute_dtype, - ) - ) - w_off = w_off + w_slab - return torch.cat(parts, dim=3) - - -def _qkv_nested_rope(attn: NeighborhoodAttention3D, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = attn.project_qkv(x) - q = attn.q_norm(q) * attn.scale - k = attn.k_norm(k) - inv_freqs = ( - attn.rope_inv_t.to(device=x.device), - attn.rope_inv_h.to(device=x.device), - attn.rope_inv_w.to(device=x.device), - ) - q = _apply_nested_full_volume_rope( - q, - attn.rope_dim_split, - inv_freqs, - num_tiles=attn.rope_num_tiles, - compute_dtype=attn.rope_compute_dtype, - ) - k = _apply_nested_full_volume_rope( - k, - attn.rope_dim_split, - inv_freqs, - num_tiles=attn.rope_num_tiles, - compute_dtype=attn.rope_compute_dtype, - ) - return q, k, v - - -def full( - x: torch.Tensor, - attn: NeighborhoodAttention3D, - norm: nn.RMSNorm, - scale: torch.Tensor, - shift: torch.Tensor, -) -> torch.Tensor: - y = norm(x) * (1.0 + scale) + shift - batch, t, h, w, _ = y.shape - kt, kh, kw = attn.kernel_size - if t < kt or h < kh or w < kw: - raise ValueError( - f"3D neighborhood attention requires spatial dims >= kernel_size; " - f"got (T,H,W)=({t},{h},{w}) vs kernel={attn.kernel_size}" - ) - - q, k, v = _qkv_nested_rope(attn, y) - q, k, v = q.contiguous(), k.contiguous(), v.contiguous() - out = attn.attention_function(attn, q, k, v) - out = out.reshape(batch, t, h, w, attn.dim) - return x + attn.proj(out) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py deleted file mode 100644 index aeb88051d..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/block.py +++ /dev/null @@ -1,28 +0,0 @@ -from __future__ import annotations - -import torch - -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.blocks import DiffusionNABlock -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.attn import full as residual_attn -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.context import combined as inject_context -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.combined.mlp import residual_mlp - - -class CombinedDiffusionNABlock(DiffusionNABlock): - def forward_combined( - self, - context_and_x: torch.Tensor, - modulation: tuple[torch.Tensor, ...], - ) -> torch.Tensor: - scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) - x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) - x = residual_attn(x, self.attn, self.norm1, scale_msa, shift_msa) - x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp) - return x - - def forward( - self, - context_and_x: torch.Tensor, - modulation: tuple[torch.Tensor, ...], - ) -> torch.Tensor: - return self.forward_combined(context_and_x, modulation) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py deleted file mode 100644 index 783ec42b3..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/context.py +++ /dev/null @@ -1,15 +0,0 @@ -from __future__ import annotations - -import torch -import torch.nn.functional as F - - -def combined( - context_and_x: torch.Tensor, - w_proj: torch.Tensor, - b_proj: torch.Tensor | None, -) -> torch.Tensor: - context_channels = w_proj.shape[1] - latent_context = context_and_x[..., :context_channels] - x = context_and_x[..., context_channels:] - return x + F.linear(latent_context, w_proj, b_proj) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py deleted file mode 100644 index e3bb98760..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/combined/mlp.py +++ /dev/null @@ -1,14 +0,0 @@ -import torch -from torch import nn - -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.layers import modulate - - -def residual_mlp( - x: torch.Tensor, - mlp: nn.Module, - norm: nn.RMSNorm, - scale: torch.Tensor, - shift: torch.Tensor, -) -> torch.Tensor: - return x + mlp(modulate(norm(x), scale, shift)) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py deleted file mode 100644 index f8a4c2529..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/det_attn_rope.py +++ /dev/null @@ -1,156 +0,0 @@ -from __future__ import annotations - -import torch - -from diffsynth.models.ltx25_diffvae.model.video_vae.transformer.rope_math import ( - DEFAULT_ABS_ROPE_NUM_TILES, - h_positions, - rot_abs_axis_impl, - t_positions, -) - - -def _apply_opaque_rope_slab( - x: torch.Tensor, - rope_split: tuple[int, int, int], - inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - *, - w_pos: torch.Tensor, - compute_dtype: torch.dtype, -) -> torch.Tensor: - d_t, d_h, _ = rope_split - inv_t, inv_h, inv_w = inv_freqs - t = x.shape[1] - h = x.shape[2] - xt = rot_abs_axis_impl(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) - xh = rot_abs_axis_impl( - x[..., d_t : d_t + d_h], - h_positions(h, x.device), - inv_h, - axis=2, - compute_dtype=compute_dtype, - ) - xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) - return torch.cat([xt, xh, xw], dim=-1) - - -def _apply_opaque_tiled_rope( - x: torch.Tensor, - rope_split: tuple[int, int, int], - inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - *, - num_tiles: int, - compute_dtype: torch.dtype, -) -> torch.Tensor: - slabs = torch.chunk(x, num_tiles, dim=3) - w_off = 0 - parts: list[torch.Tensor] = [] - for slab in slabs: - w_slab = slab.shape[3] - w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off - parts.append( - _apply_opaque_rope_slab( - slab, - rope_split, - inv_freqs, - w_pos=w_pos, - compute_dtype=compute_dtype, - ) - ) - w_off = w_off + w_slab - return torch.cat(parts, dim=3) - - -@torch.library.custom_op("diffsynth_ltx25::abs_rope", mutates_args=()) -def _abs_rope_op( - x: torch.Tensor, - inv_t: torch.Tensor, - inv_h: torch.Tensor, - inv_w: torch.Tensor, - d_t: int, - d_h: int, - d_w: int, - num_tiles: int, - compute_dtype_is_bf16: bool, -) -> torch.Tensor: - compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 - return _apply_opaque_tiled_rope( - x, - (d_t, d_h, d_w), - (inv_t, inv_h, inv_w), - num_tiles=num_tiles, - compute_dtype=compute_dtype, - ) - - -@_abs_rope_op.register_fake -def _abs_rope_fake( - x: torch.Tensor, - inv_t: torch.Tensor, - inv_h: torch.Tensor, - inv_w: torch.Tensor, - d_t: int, - d_h: int, - d_w: int, - num_tiles: int, - compute_dtype_is_bf16: bool, -) -> torch.Tensor: - del inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 - return torch.empty(x.shape, device=x.device, dtype=x.dtype) - - -def _apply_opaque_abs_rope( - x: torch.Tensor, - rope_split: tuple[int, int, int], - inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - *, - num_tiles: int, - compute_dtype: torch.dtype, -) -> torch.Tensor: - if num_tiles < 1: - raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") - if compute_dtype not in (torch.float32, torch.bfloat16): - raise ValueError(f"compute_dtype must be float32 or bfloat16, got {compute_dtype}") - d_t, d_h, d_w = rope_split - inv_t, inv_h, inv_w = inv_freqs - return _abs_rope_op( - x, - inv_t, - inv_h, - inv_w, - d_t, - d_h, - d_w, - num_tiles, - compute_dtype == torch.bfloat16, - ) - - -def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - q, k, v = attn.project_qkv(x) - q = attn.q_norm(q) - k = attn.k_norm(k) - q = q * attn.scale - - inv_freqs = ( - attn.rope_inv_t.to(device=x.device), - attn.rope_inv_h.to(device=x.device), - attn.rope_inv_w.to(device=x.device), - ) - num_tiles = getattr(attn, "rope_num_tiles", DEFAULT_ABS_ROPE_NUM_TILES) - compute_dtype = getattr(attn, "rope_compute_dtype", torch.float32) - q = _apply_opaque_abs_rope( - q, - attn.rope_dim_split, - inv_freqs, - num_tiles=num_tiles, - compute_dtype=compute_dtype, - ) - k = _apply_opaque_abs_rope( - k, - attn.rope_dim_split, - inv_freqs, - num_tiles=num_tiles, - compute_dtype=compute_dtype, - ) - return q, k, v diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py deleted file mode 100644 index c522eb352..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .eager import na3d - - -class EagerSdpaAttention: - def __call__(self, attn, q, k, v): - if q.dtype != v.dtype or k.dtype != v.dtype: - q, k = q.to(dtype=v.dtype), k.to(dtype=v.dtype) - return na3d(q, k, v, kernel_size=attn.kernel_size, is_causal=None, scale=1.0) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py deleted file mode 100644 index bd5154c23..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/fallback_na/eager.py +++ /dev/null @@ -1,154 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import math - -import torch -from torch.nn import functional - - -NA_SCORE_BUDGET = 2**25 - -NA_KV_STACK_BUDGET = 2**28 - - -def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], list[int]]: - starts: list[int] = [] - ends: list[int] = [] - if causal: - for i in range(length): - starts.append(max(0, i - kernel + 1)) - ends.append(i + 1) - else: - kernel = min(kernel, length) - lo = length - kernel - half = kernel // 2 - for i in range(length): - start = min(max(i - half, 0), lo) - starts.append(start) - ends.append(start + kernel) - return starts, ends - - -def _pick_tiles(dims: tuple[int, int, int], kernels: list[int]) -> list[int]: - tiles = list(dims) - - def cost(ts: list[int]) -> int: - nq = math.prod(ts) - nk = math.prod(min(d, t + k - 1) for t, k, d in zip(ts, kernels, dims, strict=True)) - return nq * nk - - while cost(tiles) > NA_SCORE_BUDGET and max(tiles) > 1: - i = max(range(3), key=lambda a: tiles[a] / kernels[a]) - if tiles[i] <= 1: - break - tiles[i] = max(1, (tiles[i] + 1) // 2) - return tiles - - -def _group_mask( - rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor: - bools = [] - for starts, ends in rel_bounds: - st = torch.tensor(starts, device=device) - en = torch.tensor(ends, device=device) - kj = torch.arange(int(en.max()), device=device) - bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) - visible = ( - bools[0][:, None, None, :, None, None] - & bools[1][None, :, None, None, :, None] - & bools[2][None, None, :, None, None, :] - ) - nq = visible.shape[0] * visible.shape[1] * visible.shape[2] - nk = visible.shape[3] * visible.shape[4] * visible.shape[5] - mask = torch.zeros((nq, nk), dtype=dtype, device=device) - mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) - return mask.reshape(1, 1, nq, nk) - - -def na3d( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - kernel_size: list[int] | tuple[int, ...], - is_causal: list[bool] | None = None, - scale: float | None = None, -) -> torch.Tensor: - batch, t, h, w, nh, hd = q.shape - dims = (t, h, w) - causal = [False, False, False] if is_causal is None else list(is_causal) - kernels = [k_ if c else min(k_, d) for k_, c, d in zip(kernel_size, causal, dims, strict=True)] - if scale is None: - scale = hd**-0.5 - device = q.device - if scale != 1.0: - q = q * scale - - bounds = [_window_bounds(d, k_, c) for d, k_, c in zip(dims, kernels, causal, strict=True)] - tile_t, tile_h, tile_w = _pick_tiles(dims, [min(k_, d) for k_, d in zip(kernels, dims, strict=True)]) - - groups: dict[ - tuple[ - tuple[tuple[int, ...], tuple[int, ...]], - tuple[tuple[int, ...], tuple[int, ...]], - tuple[tuple[int, ...], tuple[int, ...]], - ], - list[tuple[tuple[slice, slice, slice], tuple[slice, slice, slice]]], - ] = {} - for t0 in range(0, t, tile_t): - t1 = min(t0 + tile_t, t) - rt0, rt1 = bounds[0][0][t0], bounds[0][1][t1 - 1] - rel_t = ( - tuple(s - rt0 for s in bounds[0][0][t0:t1]), - tuple(e - rt0 for e in bounds[0][1][t0:t1]), - ) - for h0 in range(0, h, tile_h): - h1 = min(h0 + tile_h, h) - rh0, rh1 = bounds[1][0][h0], bounds[1][1][h1 - 1] - rel_h = ( - tuple(s - rh0 for s in bounds[1][0][h0:h1]), - tuple(e - rh0 for e in bounds[1][1][h0:h1]), - ) - for w0 in range(0, w, tile_w): - w1 = min(w0 + tile_w, w) - rw0, rw1 = bounds[2][0][w0], bounds[2][1][w1 - 1] - rel_w = ( - tuple(s - rw0 for s in bounds[2][0][w0:w1]), - tuple(e - rw0 for e in bounds[2][1][w0:w1]), - ) - groups.setdefault((rel_t, rel_h, rel_w), []).append( - ( - (slice(t0, t1), slice(h0, h1), slice(w0, w1)), - (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), - ) - ) - - out = torch.empty((batch, t, h, w, nh, hd), device=device, dtype=v.dtype) - for rel, tiles in groups.items(): - mask = _group_mask(rel, q.dtype, device) - nq, nk = mask.shape[2], mask.shape[3] - g_max = max(1, NA_KV_STACK_BUDGET // max(1, batch * nh * nk * hd * 2)) if device.type == "cuda" else 1 - qs0, _ = tiles[0] - tq = qs0[0].stop - qs0[0].start - th = qs0[1].stop - qs0[1].start - tw = qs0[2].stop - qs0[2].start - for c0 in range(0, len(tiles), g_max): - chunk = tiles[c0 : c0 + g_max] - g = len(chunk) - q_s = torch.stack([q[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) - k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) - v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) - q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nq, hd) - k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) - v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) - o = functional.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) - o = o.view(g, batch, nh, tq, th, tw, hd).permute(0, 1, 3, 4, 5, 2, 6) - for i, (qs, _) in enumerate(chunk): - out[:, qs[0], qs[1], qs[2]] = o[i] - - return out diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py deleted file mode 100644 index f879e247f..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/layers.py +++ /dev/null @@ -1,66 +0,0 @@ -from __future__ import annotations - -import math - -import torch -import torch.nn.functional as F -from einops import rearrange -from torch import nn - - -class ChannelLinear(nn.Linear): - - @property - def in_channels(self) -> int: - return self.in_features - - @property - def out_channels(self) -> int: - return self.out_features - - -class LinearPixelShuffleUpsample(nn.Module): - def __init__( - self, - in_channels: int, - stride: tuple[int, int, int], - out_channels_reduction_factor: int = 1, - ) -> None: - super().__init__() - self.stride = stride - self.proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor - self.out_channels = self.proj_out_channels // math.prod(stride) - self.proj = nn.Linear(in_channels, self.proj_out_channels, bias=True) - - def forward(self, x: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: - x = self.proj(x) - x = rearrange( - x, - "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c", - p1=self.stride[0], - p2=self.stride[1], - p3=self.stride[2], - ) - if self.stride[0] == 2 and drop_leading_frame: - x = x[:, 1:, :, :, :] - return x - - -class AdaLNZero(nn.Module): - NUM_CHUNKS: int = 7 - - def __init__(self, dim: int, t_emb_dim: int) -> None: - super().__init__() - self.dim = dim - self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True) - nn.init.zeros_(self.proj.weight) - nn.init.zeros_(self.proj.bias) - - def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: - h = self.proj(F.silu(t_emb)) - chunks = h.chunk(self.NUM_CHUNKS, dim=-1) - return tuple(c[:, None, None, None, :] for c in chunks) - - -def modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: - return x * (1.0 + scale) + shift diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py deleted file mode 100644 index 6096dfad0..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/qkv.py +++ /dev/null @@ -1,16 +0,0 @@ -from __future__ import annotations - -import torch -from torch import nn - - -class QKVProjections(nn.Module): - def __init__(self, dim: int) -> None: - super().__init__() - self.dim = dim - self.to_q = nn.Linear(dim, dim, bias=True) - self.to_k = nn.Linear(dim, dim, bias=True) - self.to_v = nn.Linear(dim, dim, bias=True) - - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return self.to_q(x), self.to_k(x), self.to_v(x) diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py deleted file mode 100644 index bac4fb16f..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/rope_math.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -import numpy as np -import torch - -DEFAULT_ABS_ROPE_NUM_TILES = 4 - - -def t_positions(t: int, device: torch.device) -> torch.Tensor: - return torch.arange(t, dtype=torch.float32, device=device) - - -def h_positions(h: int, device: torch.device) -> torch.Tensor: - return torch.arange(h, dtype=torch.float32, device=device) - - -def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: - assert head_dim % 8 == 0, f"head_dim={head_dim} must be a multiple of 8 for default split" - d_t = (head_dim // 4) // 2 * 2 - d_hw = (head_dim - d_t) // 2 - if d_hw % 2 != 0: - d_t -= 2 - d_hw = (head_dim - d_t) // 2 - assert d_t > 0 - assert d_hw > 0 - return (d_t, d_hw, d_hw) - - -def rope_inv_freqs(dim: int, base: float = 10000.0) -> torch.Tensor: - assert dim % 2 == 0, f"RoPE dim must be even, got {dim}" - exponents = np.arange(0, dim, 2, dtype=np.float64) / dim - inv_freqs = 1.0 / np.power(float(base), exponents) - return torch.from_numpy(inv_freqs).to(torch.float32) - - -def rot_abs_axis_impl( - xc: torch.Tensor, - pos: torch.Tensor, - inv: torch.Tensor, - axis: int, - *, - compute_dtype: torch.dtype, -) -> torch.Tensor: - out_dtype = xc.dtype - pairs = xc.reshape(*xc.shape[:-1], xc.shape[-1] // 2, 2) - xe = pairs[..., 0].to(compute_dtype) - xo = pairs[..., 1].to(compute_dtype) - shape = [1, 1, 1, 1, 1, inv.shape[0]] - shape[axis] = pos.shape[0] - ang = (pos[:, None] * inv[None, :]).reshape(shape) - c = ang.cos().to(compute_dtype) - s = ang.sin().to(compute_dtype) - re = xe * c - xo * s - ro = xe * s + xo * c - out = torch.stack([re, ro], dim=-1).reshape(xc.shape) - return out.to(out_dtype) if out.dtype != out_dtype else out diff --git a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py b/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py deleted file mode 100644 index d9ad9e0f1..000000000 --- a/diffsynth/models/ltx25_diffvae/model/video_vae/transformer/swiglu.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - - -_DEFAULT_TOKEN_CHUNK = 16_384 - - -def swiglu(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor) -> torch.Tensor: - if x.dtype != w_gate.dtype: - x = x.to(w_gate.dtype) - leading, dim = x.shape[:-1], x.shape[-1] - flat = x.reshape(-1, dim).contiguous() - output = torch.empty_like(flat) - for start in range(0, flat.shape[0], _DEFAULT_TOKEN_CHUNK): - end = min(start + _DEFAULT_TOKEN_CHUNK, flat.shape[0]) - tokens = flat[start:end] - workspace = torch.empty((end - start, w_gate.shape[0]), dtype=x.dtype, device=x.device) - torch.mm(tokens, w_gate.t(), out=workspace) - F.silu(workspace, inplace=True) - workspace.mul_(F.linear(tokens, w_up)) - torch.mm(workspace, w_down.t(), out=output[start:end]) - return output.view(*leading, dim) - - -class SwiGLU(nn.Module): - def __init__(self, dim: int, hidden_dim: int) -> None: - super().__init__() - self.w_up = nn.Linear(dim, hidden_dim, bias=False) - self.w_gate = nn.Linear(dim, hidden_dim, bias=False) - self.w_down = nn.Linear(hidden_dim, dim, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return swiglu(x, self.w_gate.weight, self.w_up.weight, self.w_down.weight) - - -def plain_mlp(x: torch.Tensor, mlp: SwiGLU, norm: nn.RMSNorm) -> torch.Tensor: - return x + mlp(norm(x)) diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index 075aae75e..7c464c1f5 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -160,6 +160,45 @@ def __init__(self): self.config = Gemma4UnifiedConfig(**copy.deepcopy(LTX25_GEMMA_CONFIG)) self.model = Gemma4UnifiedForConditionalGeneration(self.config) + self.reset_non_persistent_buffers() + + def reset_non_persistent_buffers(self): + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + language_model = self.model.model.language_model + text_config = self.config.text_config + rotary_embedding = language_model.rotary_emb + # VRAM management replaces submodules with wrappers; buffers live on the inner module. + rotary_embedding = getattr(rotary_embedding, "module", rotary_embedding) + for layer_type in dict.fromkeys(text_config.layer_types): + rope_parameters = text_config.rope_parameters[layer_type] + if rope_parameters is None: + continue + rope_type = rope_parameters["rope_type"] + if rope_type == "default": + inv_freq, attention_scaling = rotary_embedding.compute_default_rope_parameters( + text_config, layer_type=layer_type + ) + else: + init_kwargs = {"layer_type": layer_type} + if layer_type == "full_attention" and rope_type == "proportional": + init_kwargs["head_dim_key"] = "global_head_dim" + inv_freq, attention_scaling = ROPE_INIT_FUNCTIONS[rope_type](text_config, **init_kwargs) + for buffer_name, buffer_value in ( + (f"{layer_type}_inv_freq", inv_freq), + (f"{layer_type}_original_inv_freq", inv_freq.clone()), + ): + if hasattr(rotary_embedding, buffer_name): + delattr(rotary_embedding, buffer_name) + rotary_embedding.register_buffer(buffer_name, buffer_value, persistent=False) + setattr(rotary_embedding, f"{layer_type}_attention_scaling", attention_scaling) + + embed_scale = torch.tensor(text_config.hidden_size**0.5, device="cpu") + embed_tokens = language_model.embed_tokens + embed_tokens = getattr(embed_tokens, "module", embed_tokens) + if hasattr(embed_tokens, "embed_scale"): + delattr(embed_tokens, "embed_scale") + embed_tokens.register_buffer("embed_scale", embed_scale, persistent=False) def forward(self, *args, **kwargs): return self.model(*args, **kwargs) @@ -183,13 +222,18 @@ def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Te class LTX25FeatureExtractorV2(torch.nn.Module): def __init__( self, - video_aggregate_embed: torch.nn.Linear, - embedding_dim: int, - audio_aggregate_embed: torch.nn.Linear | None = None, + embedding_dim: int = 3840, + num_layers: int = 49, + video_out_features: int = 4096, + audio_out_features: int = 2048, ): super().__init__() - self.video_aggregate_embed = video_aggregate_embed - self.audio_aggregate_embed = audio_aggregate_embed + self.video_aggregate_embed = torch.nn.Linear(embedding_dim * num_layers, video_out_features, bias=True) + self.audio_aggregate_embed = ( + torch.nn.Linear(embedding_dim * num_layers, audio_out_features, bias=True) + if audio_out_features is not None + else None + ) self.embedding_dim = embedding_dim def forward( @@ -347,11 +391,9 @@ def _right_pad_order(additive_attention_mask: torch.Tensor) -> tuple[torch.Tenso return sort_indices, additive[:, None, None, :] -class LTX25TextEncoderPostModules(torch.nn.Module): +class LTX25EmbeddingsConnectors(torch.nn.Module): def __init__( self, - embedding_dim: int = 3840, - num_layers: int = 49, video_attention_heads: int = 32, video_attention_head_dim: int = 128, audio_attention_heads: int = 32, @@ -361,19 +403,6 @@ def __init__( connector_ff_bias: bool = True, ): super().__init__() - self.feature_extractor = LTX25FeatureExtractorV2( - video_aggregate_embed=torch.nn.Linear( - embedding_dim * num_layers, - video_attention_heads * video_attention_head_dim, - bias=True, - ), - embedding_dim=embedding_dim, - audio_aggregate_embed=torch.nn.Linear( - embedding_dim * num_layers, - audio_attention_heads * audio_attention_head_dim, - bias=True, - ), - ) connector_max_positions = [4096] if connector_max_positions is None else connector_max_positions self.video_connector = LTX25Embeddings1DConnector( attention_head_dim=video_attention_head_dim, @@ -390,6 +419,55 @@ def __init__( ff_bias=connector_ff_bias, ) + +class LTX25TextEncoderPostModules(torch.nn.Module): + def __init__( + self, + embedding_dim: int = 3840, + num_layers: int = 49, + video_attention_heads: int = 32, + video_attention_head_dim: int = 128, + audio_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + num_connector_layers: int = 8, + connector_max_positions: list[int] | None = None, + connector_ff_bias: bool = True, + feature_extractor: LTX25FeatureExtractorV2 | None = None, + connectors: LTX25EmbeddingsConnectors | None = None, + ): + super().__init__() + self.feature_extractor = ( + feature_extractor + if feature_extractor is not None + else LTX25FeatureExtractorV2( + embedding_dim=embedding_dim, + num_layers=num_layers, + video_out_features=video_attention_heads * video_attention_head_dim, + audio_out_features=audio_attention_heads * audio_attention_head_dim, + ) + ) + self.connectors = ( + connectors + if connectors is not None + else LTX25EmbeddingsConnectors( + video_attention_heads=video_attention_heads, + video_attention_head_dim=video_attention_head_dim, + audio_attention_heads=audio_attention_heads, + audio_attention_head_dim=audio_attention_head_dim, + num_connector_layers=num_connector_layers, + connector_max_positions=connector_max_positions, + connector_ff_bias=connector_ff_bias, + ) + ) + + @property + def video_connector(self): + return self.connectors.video_connector + + @property + def audio_connector(self): + return self.connectors.audio_connector + def create_embeddings( self, video_features: torch.Tensor, diff --git a/diffsynth/models/ltx2_common.py b/diffsynth/models/ltx2_common.py index a658ec6ec..21270b884 100644 --- a/diffsynth/models/ltx2_common.py +++ b/diffsynth/models/ltx2_common.py @@ -266,6 +266,8 @@ class Modality: attention. ``None`` means unrestricted (full) attention between all tokens. Built incrementally by conditioning items; see :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`. + keyframes_mask: Optional per-token marker of shape ``(B, T, 1)``. Non-zero + entries receive the LTX 2.5 keyframe absolute-position embedding. """ latent: ( @@ -280,6 +282,7 @@ class Modality: enabled: bool = True context_mask: torch.Tensor | None = None attention_mask: torch.Tensor | None = None + keyframes_mask: torch.Tensor | None = None def to_denoised( diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py index 8639113db..b510f9d09 100644 --- a/diffsynth/models/ltx2_dit.py +++ b/diffsynth/models/ltx2_dit.py @@ -578,6 +578,23 @@ def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states = self.linear_2(hidden_states) return hidden_states +KeyframesEmbeddingProvider = Callable[[], torch.Tensor | None] + + +def apply_keyframes_absolute_embedding( + hidden_states: torch.Tensor, + keyframes_mask: torch.Tensor | None, + embedding_provider: KeyframesEmbeddingProvider | None, +) -> torch.Tensor: + if embedding_provider is None or keyframes_mask is None: + return hidden_states + embedding = embedding_provider() + if embedding is None: + return hidden_states + mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype) + return hidden_states + mask * embedding.to(device=hidden_states.device, dtype=hidden_states.dtype) + + @dataclass(frozen=True) class TransformerArgs: x: torch.Tensor @@ -611,6 +628,7 @@ def __init__( # noqa: PLR0913 rope_type: LTXRopeType, caption_projection: torch.nn.Module | None = None, prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None, ) -> None: self.patchify_proj = patchify_proj self.adaln = adaln @@ -624,6 +642,7 @@ def __init__( # noqa: PLR0913 self.rope_type = rope_type self.caption_projection = caption_projection self.prompt_adaln = prompt_adaln + self.keyframes_embedding_provider = keyframes_embedding_provider def _prepare_timestep( self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype @@ -717,6 +736,7 @@ def prepare( cross_modality: Modality | None = None, # noqa: ARG002 ) -> TransformerArgs: x = self.patchify_proj(modality.latent) + x = apply_keyframes_absolute_embedding(x, modality.keyframes_mask, self.keyframes_embedding_provider) batch_size = x.shape[0] timestep, embedded_timestep = self._prepare_timestep( modality.timesteps, self.adaln, batch_size, modality.latent.dtype @@ -773,6 +793,8 @@ def __init__( # noqa: PLR0913 av_ca_timestep_scale_multiplier: int, caption_projection: torch.nn.Module | None = None, prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None, + use_tokenwise_av_ca_scale_shift: bool = False, ) -> None: self.simple_preprocessor = TransformerArgsPreprocessor( patchify_proj=patchify_proj, @@ -787,12 +809,14 @@ def __init__( # noqa: PLR0913 rope_type=rope_type, caption_projection=caption_projection, prompt_adaln=prompt_adaln, + keyframes_embedding_provider=keyframes_embedding_provider, ) self.cross_scale_shift_adaln = cross_scale_shift_adaln self.cross_gate_adaln = cross_gate_adaln self.cross_pe_max_pos = cross_pe_max_pos self.audio_cross_attention_dim = audio_cross_attention_dim self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier + self.use_tokenwise_av_ca_scale_shift = use_tokenwise_av_ca_scale_shift def prepare( self, @@ -809,10 +833,6 @@ def prepare( if cross_modality.sigma.ndim != 1: raise ValueError("Cross modality sigma must be a 1D tensor") - cross_timestep = cross_modality.sigma.view( - modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:]) - ) - cross_pe = self.simple_preprocessor._prepare_positional_embeddings( positions=modality.positions[:, 0:1, :], inner_dim=self.audio_cross_attention_dim, @@ -823,7 +843,8 @@ def prepare( ) cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( - timestep=cross_timestep, + modality_timesteps=modality.timesteps, + cross_modality_sigma=cross_modality.sigma, timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, batch_size=transformer_args.x.shape[0], hidden_dtype=modality.latent.dtype, @@ -838,23 +859,24 @@ def prepare( def _prepare_cross_attention_timestep( self, - timestep: torch.Tensor | None, + modality_timesteps: torch.Tensor, + cross_modality_sigma: torch.Tensor, timestep_scale_multiplier: int, batch_size: int, hidden_dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare cross attention timestep embeddings.""" - timestep = timestep * timestep_scale_multiplier - + """Prepare A-V cross-attention AdaLN inputs.""" av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier + cross_timestep = cross_modality_sigma.view(batch_size, 1, *[1] * (modality_timesteps.ndim - 2)) + scale_shift_input = modality_timesteps if self.use_tokenwise_av_ca_scale_shift else cross_timestep scale_shift_timestep, _ = self.cross_scale_shift_adaln( - timestep.flatten(), + (scale_shift_input * timestep_scale_multiplier).flatten(), hidden_dtype=hidden_dtype, ) scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) gate_noise_timestep, _ = self.cross_gate_adaln( - timestep.flatten() * av_ca_factor, + (cross_timestep * timestep_scale_multiplier * av_ca_factor).flatten(), hidden_dtype=hidden_dtype, ) gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) @@ -1313,6 +1335,7 @@ def __init__( # noqa: PLR0913 ff_bias: bool = True, audio_ff_bias: bool = True, use_keyframes_abs_pos_embedding: bool = False, + use_tokenwise_av_ca_scale_shift: bool = False, ): super().__init__() self._enable_gradient_checkpointing = False @@ -1325,6 +1348,7 @@ def __init__( # noqa: PLR0913 self.cross_attention_adaln = cross_attention_adaln self.use_prompt_adaln_single = use_prompt_adaln_single self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding + self.use_tokenwise_av_ca_scale_shift = use_tokenwise_av_ca_scale_shift cross_pe_max_pos = None if model_type.is_video_enabled(): if positional_embedding_max_pos is None: @@ -1376,6 +1400,24 @@ def __init__( # noqa: PLR0913 def _adaln_embedding_coefficient(self) -> int: return adaln_embedding_coefficient(self.cross_attention_adaln) + def _keyframes_embedding(self) -> torch.Tensor | None: + return getattr(self, "keyframes_abs_pos_embedding", None) + + @property + def supports_keyframes_abs_pos_embedding(self) -> bool: + embedding = self._keyframes_embedding() + return embedding is not None and not embedding.is_meta + + def enable_keyframes_abs_pos_embedding(self) -> None: + if not self.model_type.is_video_enabled(): + raise ValueError("The keyframe absolute-position embedding is a video-stream parameter") + existing = self._keyframes_embedding() + if existing is not None and not existing.is_meta: + return + shape = existing.shape if existing is not None else (1, self.inner_dim) + self.use_keyframes_abs_pos_embedding = True + self.keyframes_abs_pos_embedding = torch.nn.Parameter(torch.zeros(shape, dtype=torch.bfloat16)) + def _init_video( self, in_channels: int, @@ -1485,6 +1527,8 @@ def _init_preprocessors( av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, caption_projection=getattr(self, "caption_projection", None), prompt_adaln=getattr(self, "prompt_adaln_single", None), + keyframes_embedding_provider=self._keyframes_embedding, + use_tokenwise_av_ca_scale_shift=self.use_tokenwise_av_ca_scale_shift, ) self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( patchify_proj=self.audio_patchify_proj, @@ -1504,6 +1548,7 @@ def _init_preprocessors( av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, caption_projection=getattr(self, "audio_caption_projection", None), prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), + use_tokenwise_av_ca_scale_shift=self.use_tokenwise_av_ca_scale_shift, ) elif self.model_type.is_video_enabled(): self.video_args_preprocessor = TransformerArgsPreprocessor( @@ -1519,6 +1564,7 @@ def _init_preprocessors( rope_type=self.rope_type, caption_projection=getattr(self, "caption_projection", None), prompt_adaln=getattr(self, "prompt_adaln_single", None), + keyframes_embedding_provider=self._keyframes_embedding, ) elif self.model_type.is_audio_enabled(): self.audio_args_preprocessor = TransformerArgsPreprocessor( @@ -1589,13 +1635,6 @@ def _init_transformer_blocks( ) def set_gradient_checkpointing(self, enable: bool) -> None: - """Enable or disable gradient checkpointing for transformer blocks. - Gradient checkpointing trades compute for memory by recomputing activations - during the backward pass instead of storing them. This can significantly - reduce memory usage at the cost of ~20-30% slower training. - Args: - enable: Whether to enable gradient checkpointing - """ self._enable_gradient_checkpointing = enable def _process_transformer_blocks( @@ -1691,12 +1730,58 @@ def _forward( ) return vx, ax - def forward(self, video_latents, video_positions, video_context, video_timesteps, audio_latents, audio_positions, audio_context, audio_timesteps, sigma, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False): - cross_pe_max_pos = None - if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): - cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) - self._init_preprocessors(cross_pe_max_pos) - video = Modality(video_latents, sigma, video_timesteps, video_positions, video_context) - audio = Modality(audio_latents, sigma, audio_timesteps, audio_positions, audio_context) if audio_latents is not None else None - vx, ax = self._forward(video=video, audio=audio, perturbations=None, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload) + def forward( + self, + video_latents, + video_positions, + video_context, + video_timesteps, + audio_latents, + audio_positions, + audio_context, + audio_timesteps, + sigma, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + video_context_mask=None, + audio_context_mask=None, + video_attention_mask=None, + audio_attention_mask=None, + video_keyframes_mask=None, + perturbations=None, + ): + video = ( + Modality( + video_latents, + sigma, + video_timesteps, + video_positions, + video_context, + context_mask=video_context_mask, + attention_mask=video_attention_mask, + keyframes_mask=video_keyframes_mask, + ) + if video_latents is not None + else None + ) + audio = ( + Modality( + audio_latents, + sigma, + audio_timesteps, + audio_positions, + audio_context, + context_mask=audio_context_mask, + attention_mask=audio_attention_mask, + ) + if audio_latents is not None + else None + ) + vx, ax = self._forward( + video=video, + audio=audio, + perturbations=perturbations, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + ) return vx, ax diff --git a/diffsynth/models/ltx2_video_vae.py b/diffsynth/models/ltx2_video_vae.py index a70dc80e8..80aaa60f9 100644 --- a/diffsynth/models/ltx2_video_vae.py +++ b/diffsynth/models/ltx2_video_vae.py @@ -1335,6 +1335,12 @@ def __init__( encoder_version: str = "ltx-2", ): super().__init__() + if isinstance(norm_layer, str): + norm_layer = NormLayerType(norm_layer) + if isinstance(latent_log_var, str): + latent_log_var = LogVarianceType(latent_log_var) + if isinstance(encoder_spatial_padding_mode, str): + encoder_spatial_padding_mode = PaddingModeType(encoder_spatial_padding_mode) if encoder_version == "ltx-2": encoder_blocks = [ ['res_x', {'num_layers': 4}], @@ -1794,6 +1800,10 @@ def __init__( base_channels: int = 128, ): super().__init__() + if isinstance(norm_layer, str): + norm_layer = NormLayerType(norm_layer) + if isinstance(decoder_spatial_padding_mode, str): + decoder_spatial_padding_mode = PaddingModeType(decoder_spatial_padding_mode) # Spatiotemporal downscaling between decoded video space and VAE latents. # According to the LTXV paper, the standard configuration downsamples diff --git a/diffsynth/pipelines/ltx25_audio_video.py b/diffsynth/pipelines/ltx25_audio_video.py deleted file mode 100644 index 5d549a950..000000000 --- a/diffsynth/pipelines/ltx25_audio_video.py +++ /dev/null @@ -1,120 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Union - -import torch -from tqdm import tqdm - -from ..core import ModelConfig -from ..core.device.npu_compatible_device import get_device_type -from ..models.ltx25_tokenizer import LTX25GemmaTokenizer -from .ltx2_audio_video import ( - LTX2AudioVideoPipeline, - LTX2AudioVideoUnit_PromptEmbedder, -) - - -def _seconds_to_num_frames(seconds: float, frame_rate: float, min_frames: int = 1, max_frames: int = 1024) -> int: - raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames)) - frames = ((raw_frames - 1) // 8) * 8 + 1 - if frames < min_frames: - frames = min(-(-(min_frames - 1) // 8) * 8 + 1, max_frames) - return frames - - -class LTX25AudioVideoUnit_PromptEmbedder(LTX2AudioVideoUnit_PromptEmbedder): - def _preprocess_text(self, pipe, text: str): - token_pairs = pipe.tokenizer.tokenize_with_weights(text)["gemma"] - input_ids = torch.tensor([[token_id for token_id, _ in token_pairs]], device=pipe.device) - attention_mask = torch.tensor([[weight for _, weight in token_pairs]], device=pipe.device) - - outputs = pipe.text_encoder.model.model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - return outputs.hidden_states, attention_mask - - -class LTX25AudioVideoPipeline(LTX2AudioVideoPipeline): - def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): - super().__init__(device=device, torch_dtype=torch_dtype) - self.duration_head = None - self.units[2] = LTX25AudioVideoUnit_PromptEmbedder() - - @staticmethod - def from_pretrained( - torch_dtype: torch.dtype = torch.bfloat16, - device: Union[str, torch.device] = get_device_type(), - model_configs: list[ModelConfig] = [], - gemma_path: str | Path | None = None, - vram_limit: float | None = None, - load_duration_head: bool = False, - stage2_lora_config: ModelConfig | None = None, - stage2_lora_strength: float = 1.0, - ) -> "LTX25AudioVideoPipeline": - if gemma_path is None: - raise ValueError("gemma_path is required for the packed LTX-2.5 Gemma4 tokenizer assets.") - pipe = LTX25AudioVideoPipeline(device=device, torch_dtype=torch_dtype) - model_pool = pipe.download_and_load_models(model_configs, vram_limit) - pipe.text_encoder = model_pool.fetch_model("ltx25_text_encoder") - pipe.text_encoder_post_modules = model_pool.fetch_model("ltx25_text_encoder_post_modules") - pipe.dit = model_pool.fetch_model("ltx25_dit") - pipe.video_vae_encoder = model_pool.fetch_model("ltx25_video_vae_encoder") - pipe.video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") - pipe.audio_vae_decoder = model_pool.fetch_model("ltx25_audio_vae_decoder") - pipe.audio_vocoder = model_pool.fetch_model("ltx25_audio_vocoder") - pipe.audio_vae_encoder = model_pool.fetch_model("ltx25_audio_vae_encoder") - - pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") - if load_duration_head: - pipe.duration_head = model_pool.fetch_model("ltx25_duration_head") - if stage2_lora_config is not None: - stage2_lora_config.download_if_necessary() - pipe.stage2_lora_config = stage2_lora_config - pipe.stage2_lora_strength = stage2_lora_strength - pipe.tokenizer = LTX25GemmaTokenizer(gemma_path) - pipe.vram_management_enabled = pipe.check_vram_management_state() - return pipe - - @torch.no_grad() - def predict_num_frames(self, prompt: str, frame_rate: float = 24.0) -> int: - if self.duration_head is None: - raise ValueError("Automatic duration requires from_pretrained(..., load_duration_head=True) and its ModelConfig.") - self.load_models_to_device(("text_encoder", "text_encoder_post_modules", "duration_head")) - embedder = self.units[2] - hidden_states, attention_mask = embedder._preprocess_text(self, prompt) - video_context, audio_context, _ = self.text_encoder_post_modules.process_hidden_states(hidden_states, attention_mask) - seconds = float(self.duration_head(video_context, audio_context).item()) - return _seconds_to_num_frames(seconds, frame_rate) - - @torch.no_grad() - def __call__( - self, - *args, - use_two_stage_pipeline: bool = True, - use_distilled_pipeline: bool = True, - cfg_scale: float = 1.0, - num_inference_steps: int = 8, - progress_bar_cmd=tqdm, - **kwargs, - ): - if use_distilled_pipeline and not use_two_stage_pipeline: - raise ValueError("LTX-2.5 distilled inference requires the two-stage refinement flow.") - if use_distilled_pipeline and cfg_scale != 1.0: - raise ValueError("LTX-2.5 distilled inference requires cfg_scale=1.0.") - if use_two_stage_pipeline and not use_distilled_pipeline and not hasattr(self, "stage2_lora_config"): - raise ValueError("LTX-2.5 Dev two-stage inference requires stage2_lora_config.") - if kwargs.get("num_frames") is None: - prompt = kwargs.get("prompt", args[0] if args else "") - kwargs["num_frames"] = self.predict_num_frames(prompt, kwargs.get("frame_rate", 24.0)) - return super().__call__( - *args, - use_two_stage_pipeline=use_two_stage_pipeline, - use_distilled_pipeline=use_distilled_pipeline, - cfg_scale=cfg_scale, - num_inference_steps=num_inference_steps, - progress_bar_cmd=progress_bar_cmd, - **kwargs, - ) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 9cde576a3..337dc65b2 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -1,28 +1,28 @@ -import torch, types -import numpy as np -from PIL import Image -from einops import repeat +from functools import partial +from pathlib import Path from typing import Optional, Union -from einops import rearrange + import numpy as np +import torch +from einops import repeat from PIL import Image from tqdm import tqdm -from typing import Optional from transformers import AutoImageProcessor, Gemma3Processor +from ..core import ModelConfig from ..core.device.npu_compatible_device import get_device_type from ..diffusion import FlowMatchScheduler -from ..core import ModelConfig from ..diffusion.base_pipeline import BasePipeline, PipelineUnit - -from ..models.ltx2_text_encoder import LTX2TextEncoder, LTX2TextEncoderPostModules, LTXVGemmaTokenizer +from ..models.ltx2_audio_vae import LTX2AudioDecoder, LTX2AudioEncoder, LTX2Vocoder, AudioPatchifier, AudioProcessor +from ..models.ltx2_common import AudioLatentShape, VIDEO_SCALE_FACTORS, VideoLatentShape, VideoPixelShape, get_pixel_coords from ..models.ltx2_dit import LTXModel -from ..models.ltx2_video_vae import LTX2VideoEncoder, LTX2VideoDecoder, VideoLatentPatchifier -from ..models.ltx2_audio_vae import LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder, AudioPatchifier, AudioProcessor +from ..models.ltx2_text_encoder import LTX2TextEncoder, LTX2TextEncoderPostModules, LTXVGemmaTokenizer from ..models.ltx2_upsampler import LTX2LatentUpsampler -from ..models.ltx2_common import VideoLatentShape, AudioLatentShape, VideoPixelShape, get_pixel_coords, VIDEO_SCALE_FACTORS -from ..utils.data.media_io_ltx2 import ltx2_preprocess +from ..models.ltx2_video_vae import LTX2VideoDecoder, LTX2VideoEncoder, VideoLatentPatchifier +from ..models.ltx25_text_encoder import LTX25TextEncoderPostModules +from ..models.ltx25_tokenizer import LTX25GemmaTokenizer from ..utils.data.audio import convert_to_stereo +from ..utils.data.media_io_ltx2 import ltx2_preprocess class LTX2AudioVideoPipeline(BasePipeline): @@ -44,10 +44,14 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.dit: LTXModel = None self.video_vae_encoder: LTX2VideoEncoder = None self.video_vae_decoder: LTX2VideoDecoder = None + self.diffusion_video_vae_decoder = None + self.conv_video_vae_decoder: LTX2VideoDecoder = None self.audio_vae_encoder: LTX2AudioEncoder = None self.audio_vae_decoder: LTX2AudioDecoder = None self.audio_vocoder: LTX2Vocoder = None self.upsampler: LTX2LatentUpsampler = None + self.duration_head = None + self.is_ltx25 = False self.video_patchifier: VideoLatentPatchifier = VideoLatentPatchifier(patch_size=1) self.audio_patchifier: AudioPatchifier = AudioPatchifier(patch_size=1) @@ -56,8 +60,11 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.in_iteration_models = ("dit",) self.units = [ LTX2AudioVideoUnit_PipelineChecker(), - LTX2AudioVideoUnit_ShapeChecker(), + LTX2AudioVideoUnit_VideoDecoderSelector(), LTX2AudioVideoUnit_PromptEmbedder(), + LTX2AudioVideoUnit_AutoDuration(), + LTX2AudioVideoUnit_ShapeChecker(), + LTX25AudioVideoUnit_SetScheduleStage1Ancestral(), LTX2AudioVideoUnit_NoiseInitializer(), LTX2AudioVideoUnit_VideoRetakeEmbedder(), LTX2AudioVideoUnit_AudioRetakeEmbedder(), @@ -116,33 +123,70 @@ def from_pretrained( stage2_lora_config: Optional[ModelConfig] = None, stage2_lora_strength: float = 0.8, vram_limit: float = None, + gemma_path: Union[str, Path, None] = None, + load_duration_head: bool = False, ): - # Initialize pipeline pipe = LTX2AudioVideoPipeline(device=device, torch_dtype=torch_dtype) model_pool = pipe.download_and_load_models(model_configs, vram_limit) - # Fetch models - pipe.text_encoder = model_pool.fetch_model("ltx2_text_encoder") - tokenizer_config.download_if_necessary() - pipe.tokenizer = LTXVGemmaTokenizer(tokenizer_path=tokenizer_config.path) - image_processor = AutoImageProcessor.from_pretrained(tokenizer_config.path, local_files_only=True) - pipe.processor = Gemma3Processor(image_processor=image_processor, tokenizer=pipe.tokenizer.tokenizer) - - pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") - pipe.dit = model_pool.fetch_model("ltx2_dit") - pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") - pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") - pipe.audio_vae_decoder = model_pool.fetch_model("ltx2_audio_vae_decoder") - pipe.audio_vocoder = model_pool.fetch_model("ltx2_audio_vocoder") - pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") - pipe.audio_vae_encoder = model_pool.fetch_model("ltx2_audio_vae_encoder") + ltx25_text_encoder = model_pool.fetch_model("ltx25_text_encoder") + ltx25_dit = model_pool.fetch_model("ltx25_dit") + pipe.is_ltx25 = ltx25_text_encoder is not None or ltx25_dit is not None + if pipe.is_ltx25: + if ltx25_text_encoder is None or ltx25_dit is None: + raise ValueError("LTX-2.5 requires both ltx25_text_encoder and ltx25_dit components.") + if gemma_path is None: + for model_config in model_configs: + if isinstance(model_config.path, str) and "text_encoders" in model_config.path: + gemma_path = model_config.path + break + if gemma_path is None: + raise ValueError("gemma_path is required for the packed LTX-2.5 Gemma4 tokenizer assets.") + pipe.text_encoder = ltx25_text_encoder + pipe.text_encoder.reset_non_persistent_buffers() + pipe.tokenizer = LTX25GemmaTokenizer(gemma_path) + feature_extractor = model_pool.fetch_model("ltx25_feature_extractor") + connectors = model_pool.fetch_model("ltx25_embeddings_connectors") + if feature_extractor is None or connectors is None: + raise ValueError("LTX-2.5 requires ltx25_feature_extractor and ltx25_embeddings_connectors components.") + pipe.text_encoder_post_modules = LTX25TextEncoderPostModules( + feature_extractor=feature_extractor, + connectors=connectors, + ) + # The container holds VRAM-wrapped modules but is not itself wrapped, so mark it + # for load_models_to_device to offload/onload its wrapped children. + pipe.text_encoder_post_modules.vram_management_enabled = True + pipe.dit = ltx25_dit + pipe.video_vae_encoder = model_pool.fetch_model("ltx25_video_vae_encoder") + if pipe.video_vae_encoder is None: + pipe.video_vae_encoder = model_pool.fetch_model("ltx25_conv_video_vae_encoder") + pipe.diffusion_video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") + pipe.conv_video_vae_decoder = model_pool.fetch_model("ltx25_conv_video_vae_decoder") + pipe.audio_vae_decoder = model_pool.fetch_model("ltx25_audio_vae_decoder") + pipe.audio_vocoder = model_pool.fetch_model("ltx25_audio_vocoder") + pipe.audio_vae_encoder = model_pool.fetch_model("ltx25_audio_vae_encoder") + pipe.duration_head = model_pool.fetch_model("ltx25_duration_head") + if load_duration_head and pipe.duration_head is None: + raise ValueError("load_duration_head=True requires an ltx25_duration_head ModelConfig.") + else: + pipe.text_encoder = model_pool.fetch_model("ltx2_text_encoder") + tokenizer_config.download_if_necessary() + pipe.tokenizer = LTXVGemmaTokenizer(tokenizer_path=tokenizer_config.path) + image_processor = AutoImageProcessor.from_pretrained(tokenizer_config.path, local_files_only=True) + pipe.processor = Gemma3Processor(image_processor=image_processor, tokenizer=pipe.tokenizer.tokenizer) + pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") + pipe.dit = model_pool.fetch_model("ltx2_dit") + pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") + pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") + pipe.audio_vae_decoder = model_pool.fetch_model("ltx2_audio_vae_decoder") + pipe.audio_vocoder = model_pool.fetch_model("ltx2_audio_vocoder") + pipe.audio_vae_encoder = model_pool.fetch_model("ltx2_audio_vae_encoder") - # Stage 2 + pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") if stage2_lora_config is not None: pipe.stage2_lora_config = stage2_lora_config pipe.stage2_lora_strength = stage2_lora_strength - # VRAM Management pipe.vram_management_enabled = pipe.check_vram_management_state() return pipe @@ -151,18 +195,39 @@ def denoise_stage(self, inputs_shared, inputs_posi, inputs_nega, units, cfg_scal return inputs_shared, inputs_posi, inputs_nega for unit in units: inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) + cfg_scale = inputs_shared.get("cfg_scale", cfg_scale) self.load_models_to_device(self.in_iteration_models) models = {name: getattr(self, name) for name in self.in_iteration_models} + timestep_dtype = torch.float32 if self.is_ltx25 else self.torch_dtype for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): - timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) + if self.is_ltx25: + timestep = self.scheduler.sigmas[progress_id] + timestep = timestep.unsqueeze(0).to(dtype=timestep_dtype, device=self.device) noise_pred_video, noise_pred_audio = self.cfg_guided_model_fn( self.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) - inputs_shared["video_latents"] = self.step(self.scheduler, inputs_shared["video_latents"], progress_id=progress_id, noise_pred=noise_pred_video, - inpaint_mask=inputs_shared.get("denoise_mask_video", None), input_latents=inputs_shared.get("input_latents_video", None), **inputs_shared) - inputs_shared["audio_latents"] = self.step(self.scheduler, inputs_shared["audio_latents"], progress_id=progress_id, noise_pred=noise_pred_audio, - inpaint_mask=inputs_shared.get("denoise_mask_audio", None), input_latents=inputs_shared.get("input_latents_audio", None), **inputs_shared) + if inputs_shared.get("video_latents") is not None and noise_pred_video is not None: + inputs_shared["video_latents"] = self.step( + self.scheduler, + inputs_shared["video_latents"], + progress_id=progress_id, + noise_pred=noise_pred_video, + inpaint_mask=inputs_shared.get("denoise_mask_video", None), + input_latents=inputs_shared.get("input_latents_video", None), + ancestral_noise_shape=inputs_shared.get("video_ancestral_noise_shape"), + ancestral_noise_transform=inputs_shared.get("video_ancestral_noise_transform"), + ) + inputs_shared["audio_latents"] = self.step( + self.scheduler, + inputs_shared["audio_latents"], + progress_id=progress_id, + noise_pred=noise_pred_audio, + inpaint_mask=inputs_shared.get("denoise_mask_audio", None), + input_latents=inputs_shared.get("input_latents_audio", None), + ancestral_noise_shape=inputs_shared.get("audio_ancestral_noise_shape"), + ancestral_noise_transform=inputs_shared.get("audio_ancestral_noise_transform"), + ) return inputs_shared, inputs_posi, inputs_nega @torch.no_grad() @@ -194,6 +259,10 @@ def __call__( width: int = 768, num_frames: int = 121, frame_rate: int = 24, + auto_duration: bool = False, + auto_duration_min_seconds: float = 1.0, + auto_duration_max_seconds: float = 20.0, + generate_video: bool = True, # Classifier-free guidance cfg_scale: float = 3.0, # Scheduler @@ -204,6 +273,7 @@ def __call__( tile_overlap_in_pixels: int = 128, tile_size_in_frames: int = 128, tile_overlap_in_frames: int = 24, + use_diffusion_vae: Optional[bool] = None, # Special Pipelines use_two_stage_pipeline: bool = False, stage2_spatial_upsample_factor: int = 2, @@ -212,15 +282,13 @@ def __call__( # progress_bar progress_bar_cmd=tqdm, ): - # Scheduler - self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, special_case="distilled_stage1" if use_distilled_pipeline else None) - # Inputs - inputs_posi = { - "prompt": prompt, - } - inputs_nega = { - "negative_prompt": negative_prompt, - } + self.scheduler.set_timesteps( + num_inference_steps, + denoising_strength=denoising_strength, + special_case="distilled_stage1" if use_distilled_pipeline else None, + ) + inputs_posi = {"prompt": prompt} + inputs_nega = {"negative_prompt": negative_prompt} inputs_shared = { "input_images": input_images, "input_images_indexes": input_images_indexes, "input_images_strength": input_images_strength, "retake_video": retake_video, "retake_video_regions": retake_video_regions, @@ -228,21 +296,38 @@ def __call__( "in_context_videos": in_context_videos, "in_context_downsample_factor": in_context_downsample_factor, "seed": seed, "rand_device": rand_device, "height": height, "width": width, "num_frames": num_frames, "frame_rate": frame_rate, + "auto_duration": auto_duration, + "auto_duration_min_seconds": auto_duration_min_seconds, + "auto_duration_max_seconds": auto_duration_max_seconds, + "generate_video": generate_video, "cfg_scale": cfg_scale, "tiled": tiled, "tile_size_in_pixels": tile_size_in_pixels, "tile_overlap_in_pixels": tile_overlap_in_pixels, "tile_size_in_frames": tile_size_in_frames, "tile_overlap_in_frames": tile_overlap_in_frames, + "use_diffusion_vae": use_diffusion_vae, "use_two_stage_pipeline": use_two_stage_pipeline, "use_distilled_pipeline": use_distilled_pipeline, "clear_lora_before_state_two": clear_lora_before_state_two, "stage2_spatial_upsample_factor": stage2_spatial_upsample_factor, "video_patchifier": self.video_patchifier, "audio_patchifier": self.audio_patchifier, + "timestep_scale": 1.0 if self.is_ltx25 else 1000.0, } - # Stage 1 - inputs_shared, inputs_posi, inputs_nega = self.denoise_stage(inputs_shared, inputs_posi, inputs_nega, self.units, cfg_scale, progress_bar_cmd) - # Stage 2 - inputs_shared, inputs_posi, inputs_nega = self.denoise_stage(inputs_shared, inputs_posi, inputs_nega, self.stage2_units, 1.0, progress_bar_cmd, not inputs_shared["use_two_stage_pipeline"]) - # Decode - self.load_models_to_device(['video_vae_decoder']) - video = self.video_vae_decoder.decode(inputs_shared["video_latents"], tiled, tile_size_in_pixels, tile_overlap_in_pixels, tile_size_in_frames, tile_overlap_in_frames) - video = self.vae_output_to_video(video) - self.load_models_to_device(['audio_vae_decoder', 'audio_vocoder']) + inputs_shared, inputs_posi, inputs_nega = self.denoise_stage( + inputs_shared, inputs_posi, inputs_nega, self.units, cfg_scale, progress_bar_cmd + ) + inputs_shared, inputs_posi, inputs_nega = self.denoise_stage( + inputs_shared, + inputs_posi, + inputs_nega, + self.stage2_units, + 1.0, + progress_bar_cmd, + not inputs_shared["use_two_stage_pipeline"], + ) + video = None + if inputs_shared.get("generate_video", True): + video_decoder_name = inputs_shared["video_decoder_name"] + self.load_models_to_device([video_decoder_name]) + video_decoder = getattr(self, video_decoder_name) + video = video_decoder.decode(inputs_shared["video_latents"], **inputs_shared["video_decode_kwargs"]) + video = self.vae_output_to_video(video) + self.load_models_to_device(["audio_vae_decoder", "audio_vocoder"]) decoded_audio = self.audio_vae_decoder(inputs_shared["audio_latents"]) decoded_audio = self.audio_vocoder(decoded_audio) decoded_audio = self.output_audio_format_check(decoded_audio) @@ -251,26 +336,171 @@ def __call__( class LTX2AudioVideoUnit_PipelineChecker(PipelineUnit): def __init__(self): - super().__init__( - take_over=True, - input_params=("use_distilled_pipeline", "use_two_stage_pipeline"), - output_params=("use_two_stage_pipeline", "cfg_scale") - ) + super().__init__(take_over=True) def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inputs_nega): - if inputs_shared.get("use_distilled_pipeline", False): + use_distilled_pipeline = inputs_shared.get("use_distilled_pipeline", False) + use_two_stage_pipeline = inputs_shared.get("use_two_stage_pipeline", False) + if use_distilled_pipeline: inputs_shared["cfg_scale"] = 1.0 - print(f"Distilled pipeline requested, disable CFG by setting cfg_scale to 1.0.") - if inputs_shared.get("use_two_stage_pipeline", False): - # distill pipeline also uses two-stage, but it does not needs lora - if not inputs_shared.get("use_distilled_pipeline", False): + print("Distilled pipeline requested, disable CFG by setting cfg_scale to 1.0.") + if pipe.is_ltx25 and use_distilled_pipeline: + if not use_two_stage_pipeline: + raise ValueError("LTX-2.5 distilled inference requires use_two_stage_pipeline=True.") + if inputs_shared.get("seed") is None: + raise ValueError("LTX-2.5 distilled ancestral sampling requires an explicit seed.") + if use_two_stage_pipeline: + if not use_distilled_pipeline: if not (hasattr(pipe, "stage2_lora_config") and pipe.stage2_lora_config is not None): raise ValueError("Two-stage pipeline requested, but stage2_lora_config is not set in the pipeline.") - if not (hasattr(pipe, "upsampler") and pipe.upsampler is not None): + if pipe.upsampler is None: raise ValueError("Two-stage pipeline requested, but upsampler model is not loaded in the pipeline.") + if inputs_shared.get("auto_duration", False): + if pipe.duration_head is None: + raise ValueError( + "Automatic duration requires an ltx25_duration_head ModelConfig in from_pretrained()." + ) + min_seconds = inputs_shared["auto_duration_min_seconds"] + max_seconds = inputs_shared["auto_duration_max_seconds"] + if min_seconds <= 0 or max_seconds < min_seconds: + raise ValueError("Automatic duration requires 0 < min_seconds <= max_seconds.") return inputs_shared, inputs_posi, inputs_nega +class LTX2AudioVideoUnit_VideoDecoderSelector(PipelineUnit): + def __init__(self): + super().__init__( + input_params=( + "use_diffusion_vae", + "seed", + "rand_device", + "tiled", + "tile_size_in_pixels", + "tile_overlap_in_pixels", + "tile_size_in_frames", + "tile_overlap_in_frames", + "generate_video", + ), + output_params=("video_decoder_name", "video_decode_kwargs", "noise_generator"), + ) + + def process( + self, + pipe: LTX2AudioVideoPipeline, + use_diffusion_vae, + seed, + rand_device, + tiled, + tile_size_in_pixels, + tile_overlap_in_pixels, + tile_size_in_frames, + tile_overlap_in_frames, + generate_video=True, + ): + if not generate_video: + return { + "video_decoder_name": None, + "video_decode_kwargs": {}, + "noise_generator": None, + } + if not pipe.is_ltx25: + if use_diffusion_vae: + raise ValueError("Diffusion VAE decoding is only supported by LTX-2.5 checkpoints.") + decoder_name = "video_vae_decoder" + elif use_diffusion_vae is not False: + decoder_name = "diffusion_video_vae_decoder" + else: + decoder_name = "conv_video_vae_decoder" + + if getattr(pipe, decoder_name) is None: + requested = "DiffusionVAE" if decoder_name == "diffusion_video_vae_decoder" else "ConvVAE" + raise ValueError(f"{requested} decoder was requested but its model component is not loaded.") + + decode_kwargs = { + "tiled": tiled, + "tile_size_in_pixels": tile_size_in_pixels, + "tile_overlap_in_pixels": tile_overlap_in_pixels, + "tile_size_in_frames": tile_size_in_frames, + "tile_overlap_in_frames": tile_overlap_in_frames, + } + if decoder_name == "diffusion_video_vae_decoder": + conv_defaults = (512, 128, 128, 24) + current_values = ( + tile_size_in_pixels, + tile_overlap_in_pixels, + tile_size_in_frames, + tile_overlap_in_frames, + ) + if any(value is None for value in current_values) or current_values == conv_defaults: + decode_kwargs.update(dict.fromkeys(( + "tile_size_in_pixels", + "tile_overlap_in_pixels", + "tile_size_in_frames", + "tile_overlap_in_frames", + ))) + noise_generator = None + if pipe.is_ltx25 and seed is not None: + noise_generator = torch.Generator(device=rand_device).manual_seed(seed) + if decoder_name == "diffusion_video_vae_decoder": + decode_kwargs["generator"] = noise_generator + return { + "video_decoder_name": decoder_name, + "video_decode_kwargs": decode_kwargs, + "noise_generator": noise_generator, + } + + +class LTX2AudioVideoUnit_AutoDuration(PipelineUnit): + def __init__(self): + super().__init__(take_over=True) + + @staticmethod + def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds): + min_frames = round(min_seconds * frame_rate) + max_frames = round(max_seconds * frame_rate) + raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames)) + frames = ((raw_frames - 1) // 8) * 8 + 1 + if frames < min_frames: + frames = min(-(-(min_frames - 1) // 8) * 8 + 1, max_frames) + return frames + + def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inputs_nega): + if not inputs_shared.get("auto_duration", False): + return inputs_shared, inputs_posi, inputs_nega + pipe.load_models_to_device(("duration_head",)) + seconds = float( + pipe.duration_head(inputs_posi["video_context"], inputs_posi["audio_context"]).item() + ) + inputs_shared["num_frames"] = self.seconds_to_num_frames( + seconds, + inputs_shared["frame_rate"], + inputs_shared["auto_duration_min_seconds"], + inputs_shared["auto_duration_max_seconds"], + ) + return inputs_shared, inputs_posi, inputs_nega + + +class LTX25AudioVideoUnit_SetScheduleStage1Ancestral(PipelineUnit): + def __init__(self): + super().__init__(input_params=("use_distilled_pipeline", "seed")) + + def process(self, pipe: LTX2AudioVideoPipeline, use_distilled_pipeline, seed): + if not pipe.is_ltx25: + return {} + if use_distilled_pipeline: + pipe.scheduler.set_step_mode( + "euler_ancestral", + eta=1.0, + s_noise=1.0, + noise_seed=seed + 10000, + device=pipe.device, + roundtrip_denoised=True, + ) + else: + pipe.scheduler.set_step_mode("euler", roundtrip_denoised=True) + return {} + + class LTX2AudioVideoUnit_ShapeChecker(PipelineUnit): """ For two-stage pipelines, the resolution must be divisible by 64. @@ -310,9 +540,20 @@ def _preprocess_text( text: str, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: token_pairs = pipe.tokenizer.tokenize_with_weights(text)["gemma"] - input_ids = torch.tensor([[t[0] for t in token_pairs]], device=pipe.device) - attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=pipe.device) - outputs = pipe.text_encoder(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) + input_ids = torch.tensor([[token_id for token_id, _ in token_pairs]], device=pipe.device) + attention_mask = torch.tensor([[weight for _, weight in token_pairs]], device=pipe.device) + if pipe.is_ltx25: + outputs = pipe.text_encoder.model.model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + else: + outputs = pipe.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) return outputs.hidden_states, attention_mask def encode_prompt(self, pipe, text, padding_side="left"): hidden_states, attention_mask = self._preprocess_text(pipe, text) @@ -329,34 +570,179 @@ def process(self, pipe: LTX2AudioVideoPipeline, prompt: str): class LTX2AudioVideoUnit_NoiseInitializer(PipelineUnit): def __init__(self): super().__init__( - input_params=("height", "width", "num_frames", "seed", "rand_device", "frame_rate"), - output_params=("video_noise", "audio_noise", "video_positions", "audio_positions", "video_latent_shape", "audio_latent_shape") + input_params=( + "height", + "width", + "num_frames", + "seed", + "rand_device", + "frame_rate", + "noise_generator", + "generate_video", + ), + output_params=( + "video_noise", + "audio_noise", + "video_positions", + "audio_positions", + "video_latent_shape", + "audio_latent_shape", + "video_keyframes_mask", + "video_ancestral_noise_shape", + "video_ancestral_noise_transform", + "audio_ancestral_noise_shape", + "audio_ancestral_noise_transform", + "noise_generator", + ), + ) + + @staticmethod + def unpatchify_video_noise(noise, patchifier, latent_shape): + return patchifier.unpatchify_video( + noise, + latent_shape.frames, + latent_shape.height, + latent_shape.width, ) - def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0): + @staticmethod + def unpatchify_audio_noise(noise, patchifier, latent_shape): + return patchifier.unpatchify_audio(noise, latent_shape.channels, latent_shape.mel_bins) + + def process_stage( + self, + pipe: LTX2AudioVideoPipeline, + height, + width, + num_frames, + seed, + rand_device, + frame_rate=24.0, + noise_generator=None, + generate_video=True, + ): video_pixel_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) video_latent_shape = VideoLatentShape.from_pixel_shape(shape=video_pixel_shape, latent_channels=128) - video_noise = pipe.generate_noise(video_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) - - latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=video_latent_shape, device=pipe.device) - video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, True).float() - video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate - video_positions = video_positions.to(pipe.torch_dtype) + noise_dtype = pipe.torch_dtype if pipe.is_ltx25 else torch.float32 + video_noise = None + video_positions = None + video_keyframes_mask = None + video_ancestral_noise_shape = None + video_ancestral_noise_transform = None + if generate_video: + video_noise_shape = video_latent_shape.to_torch_shape() + if pipe.is_ltx25: + video_noise_shape = ( + video_latent_shape.batch, + video_latent_shape.frames * video_latent_shape.height * video_latent_shape.width, + video_latent_shape.channels, + ) + video_noise = pipe.generate_noise( + video_noise_shape, + seed=seed, + rand_device=rand_device, + rand_torch_dtype=noise_dtype, + generator=noise_generator, + ) + if pipe.is_ltx25: + video_noise = pipe.video_patchifier.unpatchify_video( + video_noise, + video_latent_shape.frames, + video_latent_shape.height, + video_latent_shape.width, + ) + + latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=video_latent_shape, device=pipe.device) + video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, True).float() + video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) audio_latent_shape = AudioLatentShape.from_video_pixel_shape(video_pixel_shape) - audio_noise = pipe.generate_noise(audio_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) + audio_noise_shape = audio_latent_shape.to_torch_shape() + if pipe.is_ltx25: + audio_noise_shape = ( + audio_latent_shape.batch, + audio_latent_shape.frames, + audio_latent_shape.channels * audio_latent_shape.mel_bins, + ) + audio_noise = pipe.generate_noise( + audio_noise_shape, + seed=seed, + rand_device=rand_device, + rand_torch_dtype=noise_dtype, + generator=noise_generator, + ) + if pipe.is_ltx25: + audio_noise = pipe.audio_patchifier.unpatchify_audio( + audio_noise, + audio_latent_shape.channels, + audio_latent_shape.mel_bins, + ) audio_positions = pipe.audio_patchifier.get_patch_grid_bounds(audio_latent_shape, device=pipe.device) + audio_ancestral_noise_shape = None + audio_ancestral_noise_transform = None + if pipe.is_ltx25 and generate_video: + video_keyframes_mask = torch.zeros( + video_latent_shape.batch, + 1, + video_latent_shape.frames, + video_latent_shape.height, + video_latent_shape.width, + dtype=torch.float32, + device=pipe.device, + ) + video_keyframes_mask[:, :, 0] = 1.0 + video_ancestral_noise_shape = video_noise_shape + video_ancestral_noise_transform = partial( + self.unpatchify_video_noise, + patchifier=pipe.video_patchifier, + latent_shape=video_latent_shape, + ) + audio_ancestral_noise_shape = audio_noise_shape + audio_ancestral_noise_transform = partial( + self.unpatchify_audio_noise, + patchifier=pipe.audio_patchifier, + latent_shape=audio_latent_shape, + ) return { "video_noise": video_noise, "audio_noise": audio_noise, "video_positions": video_positions, "audio_positions": audio_positions, "video_latent_shape": video_latent_shape, - "audio_latent_shape": audio_latent_shape + "audio_latent_shape": audio_latent_shape, + "video_keyframes_mask": video_keyframes_mask, + "video_ancestral_noise_shape": video_ancestral_noise_shape, + "video_ancestral_noise_transform": video_ancestral_noise_transform, + "audio_ancestral_noise_shape": audio_ancestral_noise_shape, + "audio_ancestral_noise_transform": audio_ancestral_noise_transform, + "noise_generator": noise_generator, } - def process(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0): - return self.process_stage(pipe, height, width, num_frames, seed, rand_device, frame_rate) + def process( + self, + pipe: LTX2AudioVideoPipeline, + height, + width, + num_frames, + seed, + rand_device, + frame_rate=24.0, + noise_generator=None, + generate_video=True, + ): + return self.process_stage( + pipe, + height, + width, + num_frames, + seed, + rand_device, + frame_rate, + noise_generator, + generate_video, + ) class LTX2AudioVideoUnit_InputVideoEmbedder(PipelineUnit): @@ -620,7 +1006,9 @@ def __init__(self): def process(self, pipe: LTX2AudioVideoPipeline, video_latents, video_noise, audio_latents, audio_noise): pipe.scheduler.set_timesteps(special_case="stage2") - video_latents = pipe.scheduler.add_noise(video_latents, video_noise, pipe.scheduler.timesteps[0]) + pipe.scheduler.set_step_mode("euler", roundtrip_denoised=pipe.is_ltx25) + if video_latents is not None and video_noise is not None: + video_latents = pipe.scheduler.add_noise(video_latents, video_noise, pipe.scheduler.timesteps[0]) audio_latents = pipe.scheduler.add_noise(audio_latents, audio_noise, pipe.scheduler.timesteps[0]) return {"video_latents": video_latents, "audio_latents": audio_latents} @@ -667,18 +1055,28 @@ def model_fn_ltx2( # Audio Inputs input_latents_audio=None, denoise_mask_audio=None, + # LTX-2.5 keyframe class embedding + video_keyframes_mask=None, + timestep_scale=1000.0, # Gradient Checkpointing use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, **kwargs, ): - timestep = timestep.float() / 1000. - - # patchify - b, c_v, f, h, w = video_latents.shape - video_latents = video_patchifier.patchify(video_latents) - seq_len_video = video_latents.shape[1] - video_timesteps = timestep.repeat(1, video_latents.shape[1], 1) + timestep = timestep.float() / timestep_scale + + video_timesteps = None + if video_latents is not None: + # patchify + b, c_v, f, h, w = video_latents.shape + video_latents = video_patchifier.patchify(video_latents) + if video_keyframes_mask is not None: + # Target LTX-2.5 keeps patchified video tokens as a channel-first view. + # Preserve that layout because BF16 GEMM reduction order depends on strides. + video_latents = video_latents.transpose(1, 2).contiguous().transpose(1, 2) + video_keyframes_mask = video_patchifier.patchify(video_keyframes_mask) + seq_len_video = video_latents.shape[1] + video_timesteps = timestep.repeat(1, video_latents.shape[1], 1) # Frist frame conditioning by replacing the video latents if input_latents_video is not None: denoise_mask_video = video_patchifier.patchify(denoise_mask_video) @@ -697,6 +1095,16 @@ def model_fn_ltx2( video_latents = torch.cat([video_latents, ref_frames_latent], dim=1) video_positions = torch.cat([video_positions, ref_frames_position], dim=2) video_timesteps = torch.cat([video_timesteps, ref_frames_timestep], dim=1) + if video_keyframes_mask is not None: + # Target marks appended single-frame guiding latents as keyframe tokens too. + ref_keyframes_mask = torch.ones( + ref_frames_latent.shape[0], + ref_frames_latent.shape[1], + 1, + dtype=video_keyframes_mask.dtype, + device=video_keyframes_mask.device, + ) + video_keyframes_mask = torch.cat([video_keyframes_mask, ref_keyframes_mask], dim=1) if audio_latents is not None: _, c_a, _, mel_bins = audio_latents.shape @@ -719,12 +1127,14 @@ def model_fn_ltx2( audio_context=audio_context, audio_timesteps=audio_timesteps, sigma=timestep, + video_keyframes_mask=video_keyframes_mask, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, ) - vx = vx[:, :seq_len_video, ...] - # unpatchify - vx = video_patchifier.unpatchify_video(vx, f, h, w) + if vx is not None: + vx = vx[:, :seq_len_video, ...] + # unpatchify + vx = video_patchifier.unpatchify_video(vx, f, h, w) ax = audio_patchifier.unpatchify_audio(ax, c_a, mel_bins) if ax is not None else None return vx, ax diff --git a/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py b/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py index 385690982..e6d84247b 100644 --- a/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py +++ b/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py @@ -1,4 +1,9 @@ def LTX25DiffusionVideoDecoderStateDictConverter(state_dict): + """Select DiffVAE weights and rename only legacy timestep-MLP keys. + + The flat decoder keeps fused ``qkv`` parameters, so this converter is safe for + lazy ``DiskMap`` inputs and never reads tensor metadata or values. + """ converted = {} for source_name in state_dict: if source_name.startswith("decoder."): @@ -8,20 +13,9 @@ def LTX25DiffusionVideoDecoderStateDictConverter(state_dict): else: continue - if name == "type_emb" or name.startswith("coarse_") or name.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")): + if name.startswith("coarse_") or name.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")): continue name = name.replace("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1.") name = name.replace("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2.") - value = state_dict[source_name] - if name.endswith(".attn.qkv.weight") or name.endswith(".attn.qkv.bias"): - if value.shape[0] % 3 != 0: - raise ValueError(f"Fused QKV tensor has invalid leading dimension: {source_name} {tuple(value.shape)}") - leaf = "weight" if name.endswith(".weight") else "bias" - prefix = name[: -len(leaf)] - q, k, v = value.chunk(3, dim=0) - converted[f"{prefix}to_q.{leaf}"] = q - converted[f"{prefix}to_k.{leaf}"] = k - converted[f"{prefix}to_v.{leaf}"] = v - else: - converted[name] = value + converted[name] = state_dict[source_name] return converted diff --git a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py index d86afe87f..fada522c9 100644 --- a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py +++ b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py @@ -16,12 +16,18 @@ def LTX25TextEncoderStateDictConverter(state_dict): return state_dict_ -def LTX25TextEncoderPostModulesStateDictConverter(state_dict): +def LTX25FeatureExtractorStateDictConverter(state_dict): state_dict_ = {} for name in state_dict: if name.startswith("text_embedding_projection."): - new_name = "feature_extractor." + name.removeprefix("text_embedding_projection.") - elif name.startswith("model.diffusion_model.video_embeddings_connector."): + state_dict_[name.removeprefix("text_embedding_projection.")] = state_dict[name] + return state_dict_ + + +def LTX25EmbeddingsConnectorsStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("model.diffusion_model.video_embeddings_connector."): new_name = "video_connector." + name.removeprefix("model.diffusion_model.video_embeddings_connector.") elif name.startswith("model.diffusion_model.audio_embeddings_connector."): new_name = "audio_connector." + name.removeprefix("model.diffusion_model.audio_embeddings_connector.") diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py new file mode 100644 index 000000000..1aebfc997 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -0,0 +1,56 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example audio comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +duration = num_frames / frame_rate +audio, audio_sample_rate = read_audio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_a2v.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py new file mode 100644 index 000000000..c51a3a74f --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py @@ -0,0 +1,87 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from modelscope import dataset_snapshot_download +from PIL import Image + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example frames come from the shared sample dataset, so reuse their paired prompt. +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +last_frame = Image.open("data/example_video_dataset/ltx2/last_frame.png").convert("RGB").resize((width, height)) + +# Single-image I2AV +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) + +# Multi-keyframe interpolation: any frame indexes within num_frames are supported. +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + input_images=[first_frame, last_frame], + input_images_indexes=[0, num_frames - 1], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_i2av_keyframes.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py new file mode 100644 index 000000000..0e3d49a05 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -0,0 +1,65 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data import VideoData +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], +) +pipe.load_lora( + pipe.dit, + ModelConfig( + model_id="Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", + origin_file_pattern="ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors", + ), +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The reference video comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +reference_video = VideoData("data/example_video_dataset/ltx2/video2.mp4", height=height // 4, width=width // 4).raw_data() +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + in_context_videos=[reference_video], + in_context_downsample_factor=2, + tiled=True, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + clear_lora_before_state_two=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_pixel_spatial_upscale.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2A.py b/examples/ltx2/model_inference/LTX-2.5-T2A.py new file mode 100644 index 000000000..889962e05 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2A.py @@ -0,0 +1,38 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.audio import save_audio + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + load_duration_head=True, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = "noise" +_, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + num_frames=121, + frame_rate=24, + num_inference_steps=30, + generate_video=False, +) +save_audio(audio, pipe.audio_vocoder.output_sampling_rate, "ltx2.5_t2a.wav") diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py new file mode 100644 index 000000000..c1cf54f4c --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -0,0 +1,55 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + load_duration_head=True, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width = 512 * 2, 768 * 2 +# Automatic duration: one pipe call predicts the clip length from the prompt and generates it. +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + frame_rate=24, + auto_duration=True, + auto_duration_min_seconds=1.0, + auto_duration_max_seconds=20.0, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py new file mode 100644 index 000000000..f5c98f6df --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py @@ -0,0 +1,50 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_int8_convrot_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py new file mode 100644 index 000000000..58772fc72 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -0,0 +1,65 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from diffsynth.utils.data import VideoData +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example video comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] + +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +path = "data/example_video_dataset/ltx2/video2.mp4" +video = VideoData(path, height=height, width=width).raw_data()[:num_frames] +assert len(video) == num_frames, f"Input video has {len(video)} frames, but expected {num_frames} frames based on the specified num_frames argument." +audio, audio_sample_rate = read_audio(path) + +# Regenerate the video within time regions. Retake regions are in seconds. +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_video=video, + retake_video_regions=[(1, 2), (3, 4)], + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + retake_audio_regions=[(0, 1), (4, 5)], + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_retake.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py new file mode 100644 index 000000000..911817ad2 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -0,0 +1,57 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example audio comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +duration = num_frames / frame_rate +audio, audio_sample_rate = read_audio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_a2v.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py new file mode 100644 index 000000000..2ca5577db --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py @@ -0,0 +1,88 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from modelscope import dataset_snapshot_download +from PIL import Image + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example frames come from the shared sample dataset, so reuse their paired prompt. +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +last_frame = Image.open("data/example_video_dataset/ltx2/last_frame.png").convert("RGB").resize((width, height)) + +# Single-image I2AV +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) + +# Multi-keyframe interpolation: any frame indexes within num_frames are supported. +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + input_images=[first_frame, last_frame], + input_images_indexes=[0, num_frames - 1], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_i2av_keyframes.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 3fde37ccb..0158873eb 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -1,54 +1,49 @@ -import os - import torch - -from diffsynth.core import ModelConfig -from diffsynth.pipelines.ltx25_audio_video import LTX25AudioVideoPipeline +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig from diffsynth.utils.data import VideoData from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 - -MODEL_ROOT = "models/Lightricks/LTX-2.5" -PIXEL_LORA = "models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler/ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors" -INPUT_VIDEO = "input.mp4" -GEMMA = f"{MODEL_ROOT}/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" -TRANSFORMER = f"{MODEL_ROOT}/diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors" -VRAM_LIMIT_GB = float(os.environ.get("LTX25_VRAM_LIMIT_GB", "16")) +from modelscope import dataset_snapshot_download vram_config = { "offload_dtype": torch.float8_e5m2, "offload_device": "cpu", "onload_dtype": torch.float8_e5m2, "onload_device": "cpu", - "preparing_dtype": torch.float8_e5m2, + "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, "computation_device": "cuda", } - -pipe = LTX25AudioVideoPipeline.from_pretrained( +pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ - ModelConfig(path=GEMMA, **vram_config), - ModelConfig(path=[GEMMA, TRANSFORMER], **vram_config), - ModelConfig(path=TRANSFORMER, **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), - ModelConfig( - path=f"{MODEL_ROOT}/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", - **vram_config, - ), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], - gemma_path=GEMMA, - vram_limit=VRAM_LIMIT_GB, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +pipe.load_lora( + pipe.dit, + ModelConfig( + model_id="Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", + origin_file_pattern="ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors", + ), ) -pipe.load_lora(pipe.dit, PIXEL_LORA) -height, width, num_frames = 576, 960, 121 -reference_video = VideoData(INPUT_VIDEO, height=height // 4, width=width // 4).raw_data() +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The reference video comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +reference_video = VideoData("data/example_video_dataset/ltx2/video2.mp4", height=height // 4, width=width // 4).raw_data() video, audio = pipe( - prompt="A colorful sailboat crosses a calm lake at sunrise. Gentle water sounds and distant birds.", - seed=42, + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, height=height, width=width, num_frames=num_frames, @@ -56,6 +51,8 @@ in_context_videos=[reference_video], in_context_downsample_factor=2, tiled=True, + cfg_scale=1.0, + num_inference_steps=8, use_distilled_pipeline=True, use_two_stage_pipeline=True, clear_lora_before_state_two=True, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py new file mode 100644 index 000000000..73e010f39 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py @@ -0,0 +1,39 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.audio import save_audio + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + load_duration_head=True, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = "noise" +_, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + num_frames=121, + frame_rate=24, + num_inference_steps=30, + generate_video=False, +) +save_audio(audio, pipe.audio_vocoder.output_sampling_rate, "ltx2.5_t2a.wav") diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index c9e7e1157..def7f6f52 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -1,48 +1,50 @@ -import os - import torch - -from diffsynth.core import ModelConfig -from diffsynth.pipelines.ltx25_audio_video import LTX25AudioVideoPipeline +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 -MODEL_ROOT = "models/Lightricks/LTX-2.5" -GEMMA = f"{MODEL_ROOT}/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" -TRANSFORMER = f"{MODEL_ROOT}/diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors" -VRAM_LIMIT_GB = float(os.environ.get("LTX25_VRAM_LIMIT_GB", "16")) - vram_config = { "offload_dtype": torch.float8_e5m2, "offload_device": "cpu", "onload_dtype": torch.float8_e5m2, "onload_device": "cpu", - "preparing_dtype": torch.float8_e5m2, + "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, "computation_device": "cuda", } - -pipe = LTX25AudioVideoPipeline.from_pretrained( +pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ - ModelConfig(path=GEMMA, **vram_config), - ModelConfig(path=[GEMMA, TRANSFORMER], **vram_config), - ModelConfig(path=TRANSFORMER, **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - gemma_path=GEMMA, - vram_limit=VRAM_LIMIT_GB, + load_duration_head=True, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width = 512 * 2, 768 * 2 +# Automatic duration: one pipe call predicts the clip length from the prompt and generates it. video, audio = pipe( - prompt="A gentle ocean wave rolls toward a quiet sunrise beach. Natural ambient surf audio.", - seed=42, - height=576, - width=960, - num_frames=121, + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + frame_rate=24, + auto_duration=True, + auto_duration_min_seconds=1.0, + auto_duration_max_seconds=20.0, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, tiled=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py new file mode 100644 index 000000000..dc75e2c78 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py @@ -0,0 +1,51 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_int8_convrot_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py new file mode 100644 index 000000000..c68d8164d --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -0,0 +1,66 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from diffsynth.utils.data import VideoData +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example video comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] + +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +path = "data/example_video_dataset/ltx2/video2.mp4" +video = VideoData(path, height=height, width=width).raw_data()[:num_frames] +assert len(video) == num_frames, f"Input video has {len(video)} frames, but expected {num_frames} frames based on the specified num_frames argument." +audio, audio_sample_rate = read_audio(path) + +# Regenerate the video within time regions. Retake regions are in seconds. +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_video=video, + retake_video_regions=[(1, 2), (3, 4)], + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + retake_audio_regions=[(0, 1), (4, 5)], + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_retake.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) From 7899c2b4e65e903a794be06295071f516d69485a Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Mon, 7 Sep 2026 23:44:34 +0800 Subject: [PATCH 03/31] Fix LTX-2.5 VRAM management, add training scripts and docs - Resolve aliased/preprocessor module references through the owner in ltx2_dit so VRAM-wrapped modules are used under CPU offload - Move STFT/mel buffers and DiffVAE raw parameters (scale_shift_table, fused QKV) to the input device/dtype at the use site - Read functional projection weights through the VRAM wrapper computation path; reduce the DiffVAE module map to leaf-level entries since the decoder calls block methods directly - Treat missing generate_video input param as True so training caches video positions (T2A keeps generate_video=False) - Add LTX-2.5 T2AV split training scripts (LoRA/full/debug), validate scripts and a series-local zero3 accelerate config - Rewrite LTX-2.5 docs for the unified pipeline, training and unsupported features; add README news and model table rows - Drop the PR leftover LTX-2.5-Keyframe-Interpolation low-VRAM script (keyframes are covered by I2AV) --- README.md | 10 + README_zh.md | 9 + .../configs/vram_management_module_maps.py | 7 - diffsynth/models/ltx25_diffusion_video_vae.py | 40 +++- diffsynth/models/ltx2_audio_vae.py | 5 +- diffsynth/models/ltx2_dit.py | 37 ++++ diffsynth/pipelines/ltx2_audio_video.py | 4 +- docs/en/Model_Details/LTX-2.5.md | 201 +++++++++++------- docs/zh/Model_Details/LTX-2.5.md | 195 ++++++++++------- .../LTX-2.5-Keyframe-Interpolation.py | 70 ------ .../full/LTX-2.5-T2AV-splited.sh | 39 ++++ .../full/accelerate_config_zero3.yaml | 23 ++ .../lora/LTX-2.5-T2AV-splited-test.sh | 42 ++++ .../lora/LTX-2.5-T2AV-splited.sh | 42 ++++ .../validate_full/LTX-2.5-T2AV.py | 44 ++++ .../validate_lora/LTX-2.5-T2AV.py | 45 ++++ 16 files changed, 566 insertions(+), 247 deletions(-) delete mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py create mode 100644 examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh create mode 100644 examples/ltx2/model_training/full/accelerate_config_zero3.yaml create mode 100644 examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh create mode 100644 examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh create mode 100644 examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py create mode 100644 examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py diff --git a/README.md b/README.md index 79147c95f..7d5445692 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,8 @@ See also: > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand. +- **September 7, 2026** We have integrated [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5), the latest audio-video generation model from Lightricks. The features include text-to-audio/video with automatic duration prediction, image-to-audio/video with keyframe interpolation, audio-to-video, audio-video retake, IC-LoRA pixel spatial upscaling, text-to-audio, INT8 quantized inference, low VRAM inference, and training. For details, please refer to the [documentation](/docs/en/Model_Details/LTX-2.5.md) and [code](/examples/ltx2/). + - **September 1, 2026** We have integrated [SenseNova-U1.5](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT), SenseTime's unified multimodal model, for which we provide text-to-image generation, image editing, low VRAM inference, and training support. For details, please refer to the [documentation](/docs/en/Model_Details/SenseNova-U1.md) and [example code](/examples/sensenova_u1/). - **August 31, 2026** We have integrated [Qwen-Video-Edit](https://modelscope.cn/models/yunpeng1998/Qwen-Video-Edit), a video editing model developed by open-source community contributor [yunpeng1998](https://github.com/yunpeng1998) based on the image editing model Qwen-Image-Edit. This serves as an excellent example of exploring and expanding model capabilities. @@ -327,6 +329,7 @@ Model overview: - Video generation - MiniMax-H3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-H3.html), [Example code](/examples/minimax_h3/) - LingBot-Video: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LingBot-Video.html), [Example code](/examples/lingbot_video/) + - LTX-2.5: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.5.html), [Example code](/examples/ltx2/) - LTX-2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.html), [Example code](/examples/ltx2/) - Wan: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html), [Example code](/examples/wanvideo/) - Audio generation @@ -635,6 +638,13 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44 | JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py) | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler) | [code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2A.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py) | - | - | - | - | | LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py) | - | - | - | - | diff --git a/README_zh.md b/README_zh.md index 97117f37d..fb4ea949f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -40,6 +40,8 @@ > 目前本项目的开发人员有限，大部分工作由 [Artiprocher](https://github.com/Artiprocher) 和 [mi804](https://github.com/mi804) 负责，因此新功能的开发进展会比较缓慢，issue 的回复和解决速度有限，我们对此感到非常抱歉，请各位开发者理解。 +- **2026年9月7日** 我们接入了 [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5)，这是 Lightricks 最新的音视频联合生成模型。支持的功能包括自动时长预测的文生音视频、关键帧插值的图生音视频、音频驱动视频、音视频区域重生成、IC-LoRA 像素空间上采样、文生音频、INT8 量化推理、低显存推理以及模型训练。详情请参考[文档](/docs/zh/Model_Details/LTX-2.5.md)和[示例代码](/examples/ltx2/)。 + - **2026年9月1日** 我们接入了 [SenseNova-U1.5](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT)，这是商汤科技开源的统一多模态模型，我们为其提供了文生图、图像编辑、低显存推理和训练支持。详情请参考[文档](/docs/zh/Model_Details/SenseNova-U1.md)和[示例代码](/examples/sensenova_u1/)。 - **2026年8月31日** 我们接入了 [Qwen-Video-Edit](https://modelscope.cn/models/yunpeng1998/Qwen-Video-Edit)，这是开源社区用户 [yunpeng1998](https://github.com/yunpeng1998) 基于图像编辑模型 Qwen-Image-Edit 训练的视频编辑模型，是探索模型能力拓展的优秀案例。 @@ -596,6 +598,13 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44 | JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py) | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler) | [code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2A.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py) | - | - | - | - | | LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py) | - | - | - | - | diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index b05a306d6..f31749a3c 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -293,11 +293,6 @@ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", }, "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder": { - "diffsynth.models.ltx25_diffusion_video_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_diffusion_video_vae.NABlock": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_diffusion_video_vae.CombinedDiffusionNABlock": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_diffusion_video_vae.SwiGLU": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_diffusion_video_vae.ChannelLinear": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, @@ -326,7 +321,6 @@ "diffsynth.models.ltx2_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.SnakeBeta": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.Activation1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.MelSTFT": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.DownSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", @@ -337,7 +331,6 @@ "diffsynth.models.ltx2_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.SnakeBeta": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.Activation1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.MelSTFT": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.ltx2_audio_vae.DownSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index 245608a90..f6473dca2 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -3215,6 +3215,15 @@ def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Ten +def vram_ready_linear(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: + # This decoder calls several projections functionally, bypassing the VRAM wrappers' + # forward, so ask the wrapper for computation-ready weights instead of reading them raw. + computation = getattr(module, "computation", None) + if computation is not None: + return computation() + return module.weight, module.bias + + class QKVProjections(nn.Module): """Checkpoint-fused QKV weights executed as the target's three projections.""" @@ -3225,8 +3234,10 @@ def __init__(self, dim: int) -> None: self.bias = linear.bias def forward(self, x): - weights = self.weight.chunk(3, dim=0) - biases = self.bias.chunk(3, dim=0) + weight = self.weight.to(device=x.device, dtype=x.dtype) + bias = self.bias.to(device=x.device, dtype=x.dtype) + weights = weight.chunk(3, dim=0) + biases = bias.chunk(3, dim=0) return tuple(F.linear(x, weight, bias) for weight, bias in zip(weights, biases, strict=True)) @@ -3277,11 +3288,19 @@ def swiglu_tiled(x, w_gate, w_up, w_down, tile, *, use_triton=None): return output.reshape(*leading, output.shape[-1]) +def swiglu_weights(mlp) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + vram_ready_linear(mlp.w_gate)[0], + vram_ready_linear(mlp.w_up)[0], + vram_ready_linear(mlp.w_down)[0], + ) + + def plain_mlp(x, mlp, norm, tile): y = norm(x) if y.numel() == 0: return x - return x + swiglu_tiled(y, mlp.w_gate.weight, mlp.w_up.weight, mlp.w_down.weight, tile) + return x + swiglu_tiled(y, *swiglu_weights(mlp), tile) class SwiGLU(nn.Module): @@ -3293,7 +3312,7 @@ def __init__(self, dim: int, hidden_dim: int, tile: SwiGLUTileSpec = DEFAULT_SWI self.tile = tile def forward(self, x): - return swiglu_tiled(x, self.w_gate.weight, self.w_up.weight, self.w_down.weight, self.tile) + return swiglu_tiled(x, *swiglu_weights(self), self.tile) def configure_swiglu_tile(module_root, *, num_tiles=None, tile_size=None): @@ -4425,8 +4444,9 @@ def __init__( def _modulation( self, modulation: tuple[torch.Tensor, ...] ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + scale_shift_table = self.scale_shift_table.to(dtype=modulation[0].dtype, device=modulation[0].device) scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ - modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) + modulation[i] + scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) ] return scale_msa, shift_msa, scale_mlp, shift_mlp @@ -4658,7 +4678,7 @@ def residual_mlp( y = modulate(norm(x), scale, shift) if y.numel() == 0: return x - return x + swiglu_tiled(y, mlp.w_gate.weight, mlp.w_up.weight, mlp.w_down.weight, tile) + return x + swiglu_tiled(y, *swiglu_weights(mlp), tile) """CombinedDiffusionNABlock: context_and_x inject + full-volume attn + residual MLP.""" @@ -4683,8 +4703,9 @@ def forward_combined_with_keyframes( masking. Both asymmetries are upstream's. """ scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) - x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) - keyframe_x = inject_context(keyframe_context_and_x, self.context_proj.weight, self.context_proj.bias) + w_proj, b_proj = vram_ready_linear(self.context_proj) + x = inject_context(context_and_x, w_proj, b_proj) + keyframe_x = inject_context(keyframe_context_and_x, w_proj, b_proj) x, keyframe_x = residual_attn_with_keyframes( x, keyframe_x, @@ -4705,7 +4726,8 @@ def forward_combined( modulation: tuple[torch.Tensor, ...], ) -> torch.Tensor: scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) - x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) + w_proj, b_proj = vram_ready_linear(self.context_proj) + x = inject_context(context_and_x, w_proj, b_proj) x = residual_attn(x, self.attn, self.norm1, scale_msa, shift_msa) x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) return x diff --git a/diffsynth/models/ltx2_audio_vae.py b/diffsynth/models/ltx2_audio_vae.py index 8a58f9724..d99f9382a 100644 --- a/diffsynth/models/ltx2_audio_vae.py +++ b/diffsynth/models/ltx2_audio_vae.py @@ -1718,7 +1718,8 @@ def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: y = y.unsqueeze(1) # (B, 1, T) left_pad = max(0, self.win_length - self.hop_length) # causal: left-only y = F.pad(y, (left_pad, 0)) - spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) + forward_basis = self.forward_basis.to(device=y.device, dtype=y.dtype) + spec = F.conv1d(y, forward_basis, stride=self.hop_length, padding=0) n_freqs = spec.shape[1] // 2 real, imag = spec[:, :n_freqs], spec[:, n_freqs:] magnitude = torch.sqrt(real**2 + imag**2) @@ -1761,7 +1762,7 @@ def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, """ magnitude, phase = self.stft_fn(y) energy = torch.norm(magnitude, dim=1) - mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) + mel = torch.matmul(self.mel_basis.to(device=magnitude.device, dtype=magnitude.dtype), magnitude) log_mel = torch.log(torch.clamp(mel, min=1e-5)) return log_mel, magnitude, phase, energy diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py index b510f9d09..ec50642f8 100644 --- a/diffsynth/models/ltx2_dit.py +++ b/diffsynth/models/ltx2_dit.py @@ -613,6 +613,25 @@ class TransformerArgs: ) +class OwnerModuleProxy: + # Preprocessors are plain objects holding references to modules owned by the model. + # VRAM management replaces the owned modules with wrappers, so resolve them lazily + # through the owner to avoid calling stale unwrapped modules. + def __init__(self, owner: torch.nn.Module, name: str): + self.owner = owner + self.name = name + + @property + def module(self): + return getattr(self.owner, self.name) + + def __call__(self, *args, **kwargs): + return self.module(*args, **kwargs) + + def __getattr__(self, item): + return getattr(self.module, item) + + class TransformerArgsPreprocessor: def __init__( # noqa: PLR0913 self, @@ -1581,6 +1600,24 @@ def _init_preprocessors( caption_projection=getattr(self, "audio_caption_projection", None), prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), ) + self._bind_preprocessor_modules() + + def _bind_preprocessor_modules(self) -> None: + owned_names = {id(module): name for name, module in self.named_children()} + pending = [ + getattr(self, "video_args_preprocessor", None), + getattr(self, "audio_args_preprocessor", None), + ] + while pending: + preprocessor = pending.pop() + if preprocessor is None: + continue + for attr_name, value in list(vars(preprocessor).items()): + if isinstance(value, torch.nn.Module): + if id(value) in owned_names: + setattr(preprocessor, attr_name, OwnerModuleProxy(self, owned_names[id(value)])) + elif isinstance(value, (TransformerArgsPreprocessor, MultiModalTransformerArgsPreprocessor)): + pending.append(value) def _init_transformer_blocks( self, diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 337dc65b2..fbbaadaf2 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -397,7 +397,7 @@ def process( tile_overlap_in_frames, generate_video=True, ): - if not generate_video: + if generate_video is False: return { "video_decoder_name": None, "video_decode_kwargs": {}, @@ -621,6 +621,8 @@ def process_stage( noise_generator=None, generate_video=True, ): + # The unit runner passes None for params missing from inputs_shared (e.g. in training). + generate_video = generate_video is not False video_pixel_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) video_latent_shape = VideoLatentShape.from_pixel_shape(shape=video_pixel_shape, latent_channels=128) noise_dtype = pipe.torch_dtype if pipe.is_ltx25 else torch.float32 diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md index cbc125a08..4c0e2de20 100644 --- a/docs/en/Model_Details/LTX-2.5.md +++ b/docs/en/Model_Details/LTX-2.5.md @@ -1,88 +1,125 @@ # LTX-2.5 -DiffSynth-Studio provides portable LTX-2.5 joint audio-video inference through -`LTX25AudioVideoPipeline`. It loads the official split BF16 checkpoints locally -and does not require `ltx_core`, NATTEN, Triton, or ltx-kernels at runtime. - -> LTX-2.5 weights are gated. Obtain access from Lightricks and place the -> checkpoint files in a local model directory before running an example. - -## Implemented components - -- LTX-2.5 22B Distilled and Dev DiT checkpoints -- Fine-tuned Gemma4 12B encoder, packed tokenizer assets, and dual AV connectors -- Duration head and automatic causal-grid frame-count prediction -- DiffVAE video encoder and pure-PyTorch eager diffusion decoder -- Audio VAE, 48 kHz BWE vocoder, spatial x2 latent upsampler, and temporal x2 - upsampler registration -- Dev stage-2 distilled-LoRA loading - -The eager DiffVAE decoder uses tiled scaled-dot-product attention as a portable -neighborhood-attention fallback. Its fixed-input output matches the upstream -eager implementation for deterministic decoder stages and an x0 diffusion step. - -## Inference modes - -The public `LTX25AudioVideoPipeline` API follows the existing LTX-2.3 API. - -| Mode | Pipeline parameters | Status | -|---|---|---| -| Distilled two-stage T2AV | `use_distilled_pipeline=True`, `use_two_stage_pipeline=True` | Supported | -| Distilled two-stage I2AV | `input_images`, `input_images_indexes` | Supported | -| Dev one-stage T2AV | `use_distilled_pipeline=False`, `use_two_stage_pipeline=False` | Supported | -| Dev one-stage I2AV | `input_images`, `use_two_stage_pipeline=False` | Supported | -| Dev two-stage T2AV/I2AV | `stage2_lora_config`, `use_two_stage_pipeline=True` | Supported | -| Audio-to-video | `retake_audio`, `audio_sample_rate`, optional `retake_audio_regions` | Supported | -| Video/audio retake | `retake_video`, `retake_video_regions`, `retake_audio_regions` | Supported | -| Keyframe interpolation | multiple `input_images` and `input_images_indexes` | Supported | -| Pixel Spatial Upscaler IC-LoRA | `in_context_videos`, `in_context_downsample_factor=2` | Supported | - -For a two-stage Dev run, supply the released distilled stage-2 LoRA through -`stage2_lora_config`. The two-stage path is required for distilled inference; -Dev also supports a one-stage path without a stage-2 LoRA. - -The Pixel Spatial Upscaler requires the official LTX-2.5 Pixel IC-LoRA. -Load it with `pipe.load_lora(pipe.dit, ModelConfig(path=...))`, pass the -reference video through `in_context_videos`, and set -`clear_lora_before_state_two=True`. Its reference resolution is one quarter of -the final height and width: stage 1 is half resolution and the adapter's -`reference_downscale_factor` is 2. Do not load an LTX-2.3 adapter into an -LTX-2.5 DiT. - -## Geometry and memory requirements - -- `num_frames % 8 == 1` -- One-stage height and width must be divisible by 32. -- Two-stage height and width must be divisible by 64. -- The low-VRAM examples use BF16 compute, FP8 CPU weight offload, and - fine-grained management for the DiT, Gemma4 encoder, text connectors, and - DiffVAE decoder. -- `LTX25_VRAM_LIMIT_GB` controls the GPU budget used to retain prepared model - layers. It defaults to 16; it is not a hard end-to-end VRAM limit. -- A 960×576×121 distilled T2AV run with `LTX25_VRAM_LIMIT_GB=16` measured a - 31.8 GiB PyTorch allocation peak and 44.4 GiB reserved peak. Treat 48 GiB - as a measured lower bound for this resolution, not a validated 48 GiB target, - and leave headroom for the driver. -- `tiled=True` does not yet split the portable DiffVAE decode volume. The - decoder determines the current full-resolution peak; lower-card support - requires a tiled decoder implementation. - -Run the low-VRAM examples with: - -```bash -python examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py -python examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py -python examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +LTX-2.5 is the joint audio-video generation model released by Lightricks. DiffSynth-Studio supports its inference and training through `LTX2AudioVideoPipeline` (the same pipeline class used by LTX-2 / LTX-2.3). Compared with LTX-2.3, LTX-2.5 introduces a fine-tuned Gemma4 12B text encoder (with packed tokenizer assets and dual audio/video connectors), the DiffVAE diffusion video decoder, automatic duration prediction via a Duration Head, and INT8 quantized checkpoints. + +## Installation + +Before using this project for inference or training, please install DiffSynth-Studio. + +```shell +git clone https://github.com/modelscope/DiffSynth-Studio.git +cd DiffSynth-Studio +pip install -e . +``` + +For more information, see [Setup](../Pipeline_Usage/Setup.md). The LTX-2.5 Gemma4 text encoder requires `transformers>=5.8,<5.15`. + +## Quickstart + +The code below loads [Lightricks/LTX-2.5](https://www.modelscope.cn/models/Lightricks/LTX-2.5) and runs inference. With `auto_duration` enabled, the pipeline first predicts the clip length from the prompt with the Duration Head and then generates audio and video, all in a single `pipe(...)` call. + +```python +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + load_duration_head=True, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=1024, width=1536, frame_rate=24, + auto_duration=True, + cfg_scale=1.0, num_inference_steps=8, + use_distilled_pipeline=True, use_two_stage_pipeline=True, + tiled=True, +) +write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate) +``` + +## Models + +|Model ID|Extra parameters|Inference|Low-VRAM inference|Full training|Validate full|LoRA training|Validate LoRA| +|-|-|-|-|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| + +## Inference + +Models are loaded with `LTX2AudioVideoPipeline.from_pretrained`; see [Load Models](../Pipeline_Usage/Model_Inference.md#load-models). LTX-2.5 shares `LTX2AudioVideoPipeline` with LTX-2.3, and the framework detects the model version from the loaded weights. + +LTX-2.5 related `from_pretrained` arguments: + +* `load_duration_head`: require the Duration Head (needed for automatic duration prediction). Defaults to `False`. +* `gemma_path`: path to the Gemma4 checkpoint, used to load the packed tokenizer assets. When omitted, it is derived from the text encoder entry in `model_configs`. +* `stage2_lora_config`: the stage-2 distilled LoRA used for two-stage inference with the Dev weights. + +For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#inference). The new or LTX-2.5 specific `LTX2AudioVideoPipeline` arguments are: + +* `auto_duration`: predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. +* `auto_duration_min_seconds` / `auto_duration_max_seconds`: lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. +* `generate_video`: whether to generate video. Defaults to `True`. Set it to `False` to generate audio only (T2A); the video VAE and latent upsampler are then not required. +* `use_diffusion_vae`: video decoder selection. `None` (default) selects by model version (LTX-2.5 uses the DiffVAE diffusion decoder); `False` uses the ConvVAE convolutional decoder (requires `ltx-2.5-video-vae-conv-bf16.safetensors`). +* `input_images` / `input_images_indexes`: keyframe images and their frame indexes. A single first frame gives image-to-video; first and last (or more) frames give keyframe interpolation. +* `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: audio-to-video (A2V) and audio region retake. + +Geometry constraints: `num_frames % 8 == 1`; height and width must be multiples of 32 for the one-stage pipeline and multiples of 64 for the two-stage pipeline. + +If you are short of VRAM, enable [VRAM management](../Pipeline_Usage/VRAM_management.md). Each example script ships a recommended low-VRAM configuration (FP8 CPU weight offload plus fine-grained VRAM management); see the table above. + +## Training + +LTX-2.5 shares the training script [`examples/ltx2/model_training/train.py`](/examples/ltx2/model_training/train.py) with LTX-2 / LTX-2.3. The general training arguments are documented in the [LTX-2 documentation](LTX-2.md#training). + +The 22B DiT and the 12B Gemma4 encoder do not fit on a single GPU together, so the LTX-2.5 training scripts use the two-stage (splited) scheme: + +1. `--task "sft:data_process"`: run text encoding and VAE encoding and cache the results to disk. +2. `--task "sft:train"`: read the cached results and train the DiT only. + +Both stages keep the full `--model_id_with_origin_paths` list and use `--fp8_models` to declare the models that are not forwarded in that stage (stage two loads the text encoder and the VAEs in FP8). The training dataset columns are `video,prompt,input_audio,frame_rate`, which map to `--data_file_keys "video,input_audio"` and `--extra_inputs "input_audio"`. + +A sample dataset is available for testing: + +```shell +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset ``` -The split checkpoints are expected under `models/Lightricks/LTX-2.5`. The -Pixel example additionally expects the separately gated adapter under -`models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler`. Set a larger -prepared-layer budget with `LTX25_VRAM_LIMIT_GB=24` when GPU memory is -available; it can improve speed but does not lower the decoder peak. Adjust -paths in an example if the local model directory is different. +After training, use `examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py` (LoRA) or `examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py` (full) to run inference with the trained checkpoint. For more details on writing training scripts, see [Model Training](../Pipeline_Usage/Model_Training.md). + +## Unsupported features -## Scope +The following official LTX-2.5 capabilities are not integrated yet: -This integration is inference-only. Training, LTX-2.5-specific DFR, Dub-It, -and HDR/EXR pipelines are not included in this scope. +* DFR (Diffusion Frame Rate) +* Native HDR / EXR output +* HDR IC-LoRA +* Dub-It dubbing diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md index f697e7880..3d0f67a21 100644 --- a/docs/zh/Model_Details/LTX-2.5.md +++ b/docs/zh/Model_Details/LTX-2.5.md @@ -1,82 +1,125 @@ # LTX-2.5 -DiffSynth-Studio 通过 `LTX25AudioVideoPipeline` 提供可移植的 LTX-2.5 -音视频联合推理。该实现从本地加载官方 BF16 分组件权重，运行时不依赖 -`ltx_core`、NATTEN、Triton 或 ltx-kernels。 - -> LTX-2.5 权重受门控保护。运行示例前，请先在 Lightricks 页面申请访问权限， -> 并将权重放到本地模型目录。 - -## 已实现的组件 - -- LTX-2.5 22B Distilled 和 Dev DiT -- LTX 微调 Gemma4 12B 编码器、内嵌 tokenizer 资产和音视频双 connector -- Duration Head 与因果帧网格的自动帧数预测 -- DiffVAE 视频编码器和纯 PyTorch eager diffusion 解码器 -- 音频 VAE、48 kHz BWE vocoder、空间 x2 latent upsampler 及时域 x2 - upsampler 注册 -- Dev 第二阶段 distilled-LoRA 加载 - -Eager DiffVAE 解码器使用 tiled scaled-dot-product attention 实现可移植的 -邻域注意力后备路径。固定输入下，其确定性解码阶段和一个 x0 diffusion step -与上游 eager 实现数值一致。 - -## 推理模式 - -公开的 `LTX25AudioVideoPipeline` API 与现有 LTX-2.3 API 对齐。 - -| 模式 | Pipeline 参数 | 状态 | -|---|---|---| -| Distilled 两阶段 T2AV | `use_distilled_pipeline=True`，`use_two_stage_pipeline=True` | 支持 | -| Distilled 两阶段 I2AV | `input_images`，`input_images_indexes` | 支持 | -| Dev 单阶段 T2AV | `use_distilled_pipeline=False`，`use_two_stage_pipeline=False` | 支持 | -| Dev 单阶段 I2AV | `input_images`，`use_two_stage_pipeline=False` | 支持 | -| Dev 两阶段 T2AV/I2AV | `stage2_lora_config`，`use_two_stage_pipeline=True` | 支持 | -| A2V | `retake_audio`，`audio_sample_rate`，可选 `retake_audio_regions` | 支持 | -| Video/audio Retake | `retake_video`，`retake_video_regions`，`retake_audio_regions` | 支持 | -| Keyframe interpolation | 多个 `input_images` 和 `input_images_indexes` | 支持 | -| Pixel Spatial Upscaler IC-LoRA | `in_context_videos`，`in_context_downsample_factor=2` | 支持 | - -Dev 两阶段推理需要通过 `stage2_lora_config` 提供发布的第二阶段 distilled-LoRA。 -Distilled 推理必须使用两阶段；Dev 同时支持不加载第二阶段 LoRA 的单阶段路径。 - -Pixel Spatial Upscaler 需要官方的 LTX-2.5 Pixel IC-LoRA。通过 -`pipe.load_lora(pipe.dit, ModelConfig(path=...))` 加载它，将参考视频传给 -`in_context_videos`，并设置 `clear_lora_before_state_two=True`。参考视频的 -宽高应为最终输出的四分之一：第一阶段为半分辨率，adapter 的 -`reference_downscale_factor` 为 2。不要将 LTX-2.3 adapter 加载到 -LTX-2.5 DiT 中。 - -## 几何与显存要求 - -- `num_frames % 8 == 1` -- 单阶段的高度、宽度必须是 32 的倍数。 -- 两阶段的高度、宽度必须是 64 的倍数。 -- 低显存示例使用 BF16 计算、FP8 CPU 权重卸载，并为 DiT、Gemma4 编码器、 - text connector 和 DiffVAE decoder 启用细粒度显存管理。 -- `LTX25_VRAM_LIMIT_GB` 控制保留在 GPU 上的预加载模型层预算，默认值为 16； - 它不是端到端显存硬上限。 -- 在 `LTX25_VRAM_LIMIT_GB=16` 下，960×576×121 distilled T2AV 实测 PyTorch - allocated 峰值为 31.8 GiB、reserved 峰值为 44.4 GiB。48 GiB 仅是该分辨率 - 的实测下界，不代表已在 48 GiB 显存卡上验证，并应为驱动预留余量。 -- 当前 `tiled=True` 不会切分 portable DiffVAE 的完整 decode volume。decoder - 决定了目前的全分辨率峰值；要支持更低显存卡，需要实现 tiled decoder。 - -运行低显存示例： - -```bash -python examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py -python examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py -python examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +LTX-2.5 是 Lightricks 发布的音视频联合生成模型。DiffSynth-Studio 通过 `LTX2AudioVideoPipeline` 提供其推理与训练支持（与 LTX-2 / LTX-2.3 共用同一个 Pipeline 类）。相比 LTX-2.3，LTX-2.5 引入了微调版 Gemma4 12B 文本编码器（内嵌 tokenizer 资产与音视频双 connector）、DiffVAE 扩散视频解码器、Duration Head 自动时长预测，以及 INT8 量化权重。 + +## 安装 + +在使用本项目进行模型推理和训练前，请先安装 DiffSynth-Studio。 + +```shell +git clone https://github.com/modelscope/DiffSynth-Studio.git +cd DiffSynth-Studio +pip install -e . +``` + +更多关于安装的信息，请参考[安装依赖](../Pipeline_Usage/Setup.md)。LTX-2.5 的 Gemma4 文本编码器需要 `transformers>=5.8,<5.15`。 + +## 快速开始 + +运行以下代码可以快速加载 [Lightricks/LTX-2.5](https://www.modelscope.cn/models/Lightricks/LTX-2.5) 模型并进行推理。开启 `auto_duration` 后，Pipeline 会先用 Duration Head 从提示词预测视频时长，再完成音视频生成，整个过程只需一次 `pipe(...)` 调用。 + +```python +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + load_duration_head=True, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=1024, width=1536, frame_rate=24, + auto_duration=True, + cfg_scale=1.0, num_inference_steps=8, + use_distilled_pipeline=True, use_two_stage_pipeline=True, + tiled=True, +) +write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate) +``` + +## 模型总览 + +|模型 ID|额外参数|推理|低显存推理|全量训练|全量训练后验证|LoRA 训练|LoRA 训练后验证| +|-|-|-|-|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| + +## 模型推理 + +模型通过 `LTX2AudioVideoPipeline.from_pretrained` 加载，详见[加载模型](../Pipeline_Usage/Model_Inference.md#加载模型)。LTX-2.5 与 LTX-2.3 共用 `LTX2AudioVideoPipeline`，框架根据加载到的权重自动识别模型版本。 + +`from_pretrained` 的 LTX-2.5 相关参数： + +* `load_duration_head`: 是否要求加载 Duration Head（自动时长预测所需），默认为 `False`。 +* `gemma_path`: Gemma4 权重路径，用于加载内嵌的 tokenizer 资产。留空时自动从 `model_configs` 中的 text encoder 路径推导。 +* `stage2_lora_config`: Dev 权重两阶段推理时使用的第二阶段 distilled-LoRA。 + +`LTX2AudioVideoPipeline` 的通用推理参数见 [LTX-2 文档](LTX-2.md#模型推理)，LTX-2.5 新增或特有的参数为： + +* `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 +* `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 +* `generate_video`: 是否生成视频，默认为 `True`。设置为 `False` 时只生成音频（T2A），此时无需加载视频 VAE 与 latent upsampler。 +* `use_diffusion_vae`: 视频解码器选择。`None`（默认）表示按模型版本自动选择（LTX-2.5 使用 DiffVAE 扩散解码器），`False` 表示使用 ConvVAE 卷积解码器（需加载 `ltx-2.5-video-vae-conv-bf16.safetensors`）。 +* `input_images` / `input_images_indexes`: 关键帧图像及其帧索引。传入首帧即为图生视频，传入首尾（或多帧）即为关键帧插值。 +* `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: 音频驱动视频（A2V）与音频区域重生成。 + +几何约束：`num_frames % 8 == 1`；单阶段的宽高为 32 的倍数，两阶段的宽高为 64 的倍数。 + +如果显存不足，请开启[显存管理](../Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置（FP8 CPU 权重卸载 + 细粒度显存管理），详见前文"模型总览"中的表格。 + +## 模型训练 + +LTX-2.5 与 LTX-2 / LTX-2.3 共用训练脚本 [`examples/ltx2/model_training/train.py`](/examples/ltx2/model_training/train.py)，通用训练参数的说明见 [LTX-2 文档](LTX-2.md#模型训练)。 + +由于 22B DiT 与 12B Gemma4 编码器无法同时放入单卡，LTX-2.5 的训练脚本采用双阶段（splited）方案： + +1. `--task "sft:data_process"`：运行文本编码与 VAE 编码，把结果缓存到硬盘。 +2. `--task "sft:train"`：从缓存读取前处理结果，只训练 DiT。 + +两个阶段都保留完整的 `--model_id_with_origin_paths`，并用 `--fp8_models` 声明该阶段不需要前向的模型（阶段二的 TextEncoder 与 VAE 使用 FP8 加载）。训练数据集的字段为 `video,prompt,input_audio,frame_rate`，对应 `--data_file_keys "video,input_audio"` 与 `--extra_inputs "input_audio"`。 + +我们构建了一个样例视频数据集，以方便您进行测试，通过以下命令可以下载这个数据集： + +```shell +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset ``` -分组件权重默认位于 `models/Lightricks/LTX-2.5`。Pixel 示例还需要单独门控的 -adapter，默认位于 -`models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler`。显存充足时可以 -通过 `LTX25_VRAM_LIMIT_GB=24` 增大预加载层预算以改善速度，但它不会降低 decoder -峰值。如本地模型目录不同，请修改示例中的路径。 +训练完成后，可以使用 `examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py`（LoRA）或 `examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py`（全量）加载训练产物进行推理验证。关于如何编写模型训练脚本，请参考[模型训练](../Pipeline_Usage/Model_Training.md)。 + +## 暂不支持的功能 -## 范围 +以下 LTX-2.5 官方能力暂未接入： -该接入仅覆盖推理。训练、LTX-2.5 专有 DFR、Dub-It 和 HDR/EXR -pipeline 不在本次范围内。 +* DFR（Diffusion Frame Rate） +* Native HDR / EXR 输出 +* HDR IC-LoRA +* Dub-It 配音 diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py deleted file mode 100644 index 4d306e77d..000000000 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-Keyframe-Interpolation.py +++ /dev/null @@ -1,70 +0,0 @@ -import os - -import torch -from PIL import Image - -from diffsynth.core import ModelConfig -from diffsynth.pipelines.ltx25_audio_video import LTX25AudioVideoPipeline -from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 - -MODEL_ROOT = "models/Lightricks/LTX-2.5" -START_IMAGE = "start.png" -MIDDLE_IMAGE = "middle.png" -END_IMAGE = "end.png" -GEMMA = f"{MODEL_ROOT}/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors" -TRANSFORMER = f"{MODEL_ROOT}/diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors" -STAGE2_LORA = f"{MODEL_ROOT}/loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors" -VRAM_LIMIT_GB = float(os.environ.get("LTX25_VRAM_LIMIT_GB", "16")) - -vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", - "preparing_dtype": torch.float8_e5m2, - "preparing_device": "cuda", - "computation_dtype": torch.bfloat16, - "computation_device": "cuda", -} - -pipe = LTX25AudioVideoPipeline.from_pretrained( - torch_dtype=torch.bfloat16, - device="cuda", - model_configs=[ - ModelConfig(path=GEMMA, **vram_config), - ModelConfig(path=[GEMMA, TRANSFORMER], **vram_config), - ModelConfig(path=TRANSFORMER, **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), - ModelConfig(path=f"{MODEL_ROOT}/vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), - ModelConfig( - path=f"{MODEL_ROOT}/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", - **vram_config, - ), - ], - gemma_path=GEMMA, - stage2_lora_config=ModelConfig(path=STAGE2_LORA), - vram_limit=VRAM_LIMIT_GB, -) - -video, audio = pipe( - prompt="A colorful sailboat crosses a calm lake at sunrise. Gentle water sounds and distant birds.", - seed=42, - height=576, - width=960, - num_frames=121, - frame_rate=24, - input_images=[Image.open(path).convert("RGB") for path in (START_IMAGE, MIDDLE_IMAGE, END_IMAGE)], - input_images_indexes=[0, 60, 120], - tiled=True, - use_distilled_pipeline=False, - use_two_stage_pipeline=True, - cfg_scale=3.0, - num_inference_steps=8, -) -write_video_audio_ltx2( - video=video, - audio=audio, - output_path="ltx2.5_keyframe_interpolation.mp4", - fps=24, - audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, -) diff --git a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh new file mode 100644 index 000000000..ba8aed93c --- /dev/null +++ b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV-full-splited-cache" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch --config_file examples/ltx2/model_training/full/accelerate_config_zero3.yaml examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-T2AV-full-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV-full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --task "sft:train" diff --git a/examples/ltx2/model_training/full/accelerate_config_zero3.yaml b/examples/ltx2/model_training/full/accelerate_config_zero3.yaml new file mode 100644 index 000000000..e6a8d2733 --- /dev/null +++ b/examples/ltx2/model_training/full/accelerate_config_zero3.yaml @@ -0,0 +1,23 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 1 + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: true + zero3_save_16bit_model: true + zero_stage: 3 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +enable_cpu_affinity: false +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh new file mode 100644 index 000000000..54e3a69d2 --- /dev/null +++ b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh @@ -0,0 +1,42 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training (debug: 1 epoch, dataset_repeat 1) +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV_lora-splited-cache-test" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-T2AV_lora-splited-cache-test \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV_lora-test" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:train" diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh new file mode 100644 index 000000000..9888d5867 --- /dev/null +++ b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh @@ -0,0 +1,42 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV_lora-splited-cache" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-T2AV_lora-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:train" diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py new file mode 100644 index 000000000..c07fdcd83 --- /dev/null +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py @@ -0,0 +1,44 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(path="./models/train/LTX2.5-T2AV-full/epoch-1.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512, 768, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=4.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_t2av_full.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py new file mode 100644 index 000000000..1f104b6be --- /dev/null +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py @@ -0,0 +1,45 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +pipe.load_lora(pipe.dit, "models/train/LTX2.5-T2AV_lora/epoch-4.safetensors") +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +height, width, num_frames = 512, 768, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=4.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_t2av_lora.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) From 5e77bf97e4fc2864053aa217310224dbd0b09a0d Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 00:35:33 +0800 Subject: [PATCH 04/31] Use ZeRO-2 with CPU offload for LTX-2.5 full training ZeRO-3 is incompatible with the Gemma4 text encoder: DiffSynth constructs models inside deepspeed.zero.Init, where transformers' _init_weights indexes weight[padding_idx] on an empty sharded embedding and raises IndexError. Replace the series-local zero3 config with the ZeRO-2 + CPU optimizer/param offload config used by the LTX-2.3 full scripts. --- .../ltx2/model_training/full/LTX-2.5-T2AV-splited.sh | 2 +- ...ig_zero3.yaml => accelerate_config_zero2offload.yaml} | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) rename examples/ltx2/model_training/full/{accelerate_config_zero3.yaml => accelerate_config_zero2offload.yaml} (74%) diff --git a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh index ba8aed93c..45ecf8a39 100644 --- a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh +++ b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh @@ -19,7 +19,7 @@ accelerate launch examples/ltx2/model_training/train.py \ --use_gradient_checkpointing \ --task "sft:data_process" -accelerate launch --config_file examples/ltx2/model_training/full/accelerate_config_zero3.yaml examples/ltx2/model_training/train.py \ +accelerate launch --config_file examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml examples/ltx2/model_training/train.py \ --dataset_base_path ./models/train/LTX2.5-T2AV-full-splited-cache \ --data_file_keys "video,input_audio" \ --extra_inputs "input_audio" \ diff --git a/examples/ltx2/model_training/full/accelerate_config_zero3.yaml b/examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml similarity index 74% rename from examples/ltx2/model_training/full/accelerate_config_zero3.yaml rename to examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml index e6a8d2733..8a75f3d91 100644 --- a/examples/ltx2/model_training/full/accelerate_config_zero3.yaml +++ b/examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml @@ -2,11 +2,10 @@ compute_environment: LOCAL_MACHINE debug: false deepspeed_config: gradient_accumulation_steps: 1 - offload_optimizer_device: none - offload_param_device: none - zero3_init_flag: true - zero3_save_16bit_model: true - zero_stage: 3 + offload_optimizer_device: 'cpu' + offload_param_device: 'cpu' + zero3_init_flag: false + zero_stage: 2 distributed_type: DEEPSPEED downcast_bf16: 'no' enable_cpu_affinity: false From 9cf43135efc916e769b45300ca469f868d0a45f0 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 01:04:57 +0800 Subject: [PATCH 05/31] Remove decorative separator comment in ltx2_audio_vae --- diffsynth/models/ltx2_audio_vae.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/diffsynth/models/ltx2_audio_vae.py b/diffsynth/models/ltx2_audio_vae.py index d99f9382a..b23fefae9 100644 --- a/diffsynth/models/ltx2_audio_vae.py +++ b/diffsynth/models/ltx2_audio_vae.py @@ -1285,10 +1285,8 @@ def get_padding(kernel_size: int, dilation: int = 1) -> int: return int((kernel_size * dilation - dilation) / 2) -# --------------------------------------------------------------------------- # Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2 # Adopted from https://github.com/NVIDIA/BigVGAN -# --------------------------------------------------------------------------- def _sinc(x: torch.Tensor) -> torch.Tensor: From 2e063ecfe5016b3b27e1b797846e8da01811f2c1 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 02:49:39 +0800 Subject: [PATCH 06/31] Fix LTX-2.5 A2V fidelity and two-stage refinement - Return the original input audio (resampled and trimmed) when the audio modality is fully frozen, matching the upstream A2V pipeline which skips the VAE/vocoder round trip to preserve fidelity (mel corr vs input 0.928 -> 0.998) - Run the second stage without classifier-free guidance, as upstream uses a simple denoiser there; the stage-1 cfg scale previously leaked into stage 2 - Align the stage-2 distilled LoRA strength default with upstream (1.0 instead of 0.8), which removes residual dithering in dev two-stage outputs --- diffsynth/pipelines/ltx2_audio_video.py | 30 ++++++++++++++++++------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index fbbaadaf2..f528d930a 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -21,7 +21,7 @@ from ..models.ltx2_video_vae import LTX2VideoDecoder, LTX2VideoEncoder, VideoLatentPatchifier from ..models.ltx25_text_encoder import LTX25TextEncoderPostModules from ..models.ltx25_tokenizer import LTX25GemmaTokenizer -from ..utils.data.audio import convert_to_stereo +from ..utils.data.audio import convert_to_stereo, resample_waveform from ..utils.data.media_io_ltx2 import ltx2_preprocess @@ -121,7 +121,7 @@ def from_pretrained( model_configs: list[ModelConfig] = [], tokenizer_config: ModelConfig = ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config: Optional[ModelConfig] = None, - stage2_lora_strength: float = 0.8, + stage2_lora_strength: float = 1.0, vram_limit: float = None, gemma_path: Union[str, Path, None] = None, load_duration_head: bool = False, @@ -327,10 +327,23 @@ def __call__( video_decoder = getattr(self, video_decoder_name) video = video_decoder.decode(inputs_shared["video_latents"], **inputs_shared["video_decode_kwargs"]) video = self.vae_output_to_video(video) - self.load_models_to_device(["audio_vae_decoder", "audio_vocoder"]) - decoded_audio = self.audio_vae_decoder(inputs_shared["audio_latents"]) - decoded_audio = self.audio_vocoder(decoded_audio) - decoded_audio = self.output_audio_format_check(decoded_audio) + retake_audio = inputs_shared.get("retake_audio") + denoise_mask_audio = inputs_shared.get("denoise_mask_audio") + audio_fully_frozen = ( + retake_audio is not None + and denoise_mask_audio is not None + and float(denoise_mask_audio.abs().max()) == 0.0 + ) + if audio_fully_frozen: + waveform, waveform_sample_rate = retake_audio + decoded_audio = resample_waveform(waveform, waveform_sample_rate, self.audio_vocoder.output_sampling_rate) + num_samples = int(inputs_shared["num_frames"] / inputs_shared["frame_rate"] * self.audio_vocoder.output_sampling_rate) + decoded_audio = self.output_audio_format_check(decoded_audio[..., :num_samples]) + else: + self.load_models_to_device(["audio_vae_decoder", "audio_vocoder"]) + decoded_audio = self.audio_vae_decoder(inputs_shared["audio_latents"]) + decoded_audio = self.audio_vocoder(decoded_audio) + decoded_audio = self.output_audio_format_check(decoded_audio) return video, decoded_audio @@ -1003,7 +1016,7 @@ class LTX2AudioVideoUnit_SetScheduleStage2(PipelineUnit): def __init__(self): super().__init__( input_params=("video_latents", "video_noise", "audio_latents", "audio_noise"), - output_params=("video_latents", "audio_latents"), + output_params=("video_latents", "audio_latents", "cfg_scale"), ) def process(self, pipe: LTX2AudioVideoPipeline, video_latents, video_noise, audio_latents, audio_noise): @@ -1012,7 +1025,8 @@ def process(self, pipe: LTX2AudioVideoPipeline, video_latents, video_noise, audi if video_latents is not None and video_noise is not None: video_latents = pipe.scheduler.add_noise(video_latents, video_noise, pipe.scheduler.timesteps[0]) audio_latents = pipe.scheduler.add_noise(audio_latents, audio_noise, pipe.scheduler.timesteps[0]) - return {"video_latents": video_latents, "audio_latents": audio_latents} + # The refinement stage runs without classifier-free guidance. + return {"video_latents": video_latents, "audio_latents": audio_latents, "cfg_scale": 1.0} class LTX2AudioVideoUnit_LatentsUpsampler(PipelineUnit): From ff245b4fc81e6c546ee92e64158764352beb6ee4 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 02:54:43 +0800 Subject: [PATCH 07/31] Remove dead code from the LTX-2.5 DiffVAE module Drop unreferenced helpers, tiling constructors and the NATTEN attention fallback that the portable eager path never selects. Verified bit-identical outputs (max abs pixel diff 0) for the distilled T2AV and I2AV keyframe examples before/after. --- diffsynth/models/ltx25_diffusion_video_vae.py | 327 ------------------ 1 file changed, 327 deletions(-) diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index f6473dca2..600dd8923 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -244,21 +244,6 @@ def from_blocks(cls, blocks: list, patch_size: int) -> "SpatioTemporalScaleFacto spatial = patch_size * (2**spatial_steps) return cls(time=2**temporal_steps, height=spatial, width=spatial) - @classmethod - def from_model_config(cls, model_config: dict) -> "SpatioTemporalScaleFactors": - """Derive the video scale factors from a checkpoint's model config dict. - Reads the embedded VAE block list (see ``from_blocks``). Falls back to the - default when the config carries no VAE block list -- either no ``vae`` section - or a ``vae`` section without encoder/decoder blocks (e.g. audio-only - checkpoints), where video tools are never built. - """ - vae_config = model_config.get("vae", {}) - blocks = vae_config.get("encoder_blocks") or vae_config.get("decoder_blocks") - if not blocks: - return cls.default() - return cls.from_blocks(blocks, vae_config.get("patch_size", 4)) - - VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default() @@ -289,13 +274,6 @@ def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": width=shape[4], ) - def token_count(self) -> int: - """Number of tokens after patchification with the default patch size of 1.""" - return self.frames * self.height * self.width - - def mask_shape(self) -> "VideoLatentShape": - return self._replace(channels=1) - @staticmethod def from_pixel_shape( shape: VideoPixelShape, @@ -337,13 +315,6 @@ class AudioLatentShape(NamedTuple): def to_torch_shape(self) -> torch.Size: return torch.Size([self.batch, self.channels, self.frames, self.mel_bins]) - def token_count(self) -> int: - """Number of tokens after patchification.""" - return self.frames - - def mask_shape(self) -> "AudioLatentShape": - return self._replace(channels=1, mel_bins=1) - @staticmethod def from_torch_shape(shape: torch.Size) -> "AudioLatentShape": return AudioLatentShape( @@ -434,11 +405,6 @@ def num_keyframes(self) -> int: def num_tokens(self) -> int: return self.num_keyframes * self.tokens_per_keyframe - @property - def token_slice(self) -> slice: - return slice(self.first_token, self.first_token + self.num_tokens) - - @dataclass(frozen=True) class LatentState: """ @@ -500,12 +466,6 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1 return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) -def check_config_value(config: dict, key: str, expected: Any) -> None: # noqa: ANN401 - actual = config.get(key) - if actual != expected: - raise ValueError(f"Config value {key} is {actual}, expected {expected}") - - def to_velocity( sample: torch.Tensor, sigma: float | torch.Tensor, @@ -540,16 +500,6 @@ def to_denoised( return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) -def find_matching_file(root_path: str, pattern: str) -> Path: - """ - Recursively search for files matching a glob pattern and return the first match. - """ - matches = list(Path(root_path).rglob(pattern)) - if not matches: - raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}") - return matches[0] - - def compute_trapezoidal_mask_1d( length: int, ramp_left: int, @@ -785,28 +735,6 @@ def split(dimension_size: int) -> DimensionIntervals: return split -def split_temporal(tile_size_frames: int, overlap_frames: int) -> SplitOperation: - """Split a temporal axis in video frame space into overlapping tiles. - Args: - tile_size_frames: Tile length in frames. - overlap_frames: Overlap between consecutive tiles in frames. - Returns: - Split operation that takes frame count and returns DimensionIntervals in frame indices. - """ - non_causal_split = split_by_size(tile_size_frames, overlap_frames) - - def split(dimension_size: int) -> DimensionIntervals: - if dimension_size <= tile_size_frames: - return DEFAULT_SPLIT_OPERATION(dimension_size) - dim_intervals = non_causal_split(dimension_size) - modified_intervals = [ - replace(interval, end=interval.end + 1, right_ramp=0) for interval in dim_intervals.intervals[:-1] - ] + [replace(dim_intervals.intervals[-1], right_ramp=0)] - return DimensionIntervals(intervals=modified_intervals) - - return split - - def split_by_count_temporal_causal( num_tiles: int, overlap: int = 0, min_tile_size: int | None = None ) -> SplitOperation: @@ -1065,25 +993,6 @@ def masks_are_complementary( return True -def compute_summed_weights( - tiles: Sequence[Tile], - full_shape: Sequence[int], -) -> torch.Tensor: - """Build the dense denominator for weighted blending over ``full_shape``. - Uses separable per-axis mask broadcasts — never ``Tile.blend_mask``. - Requires concrete ``out_coords`` (``stop`` not ``None``) on every axis. - Always builds on CPU float32 so CUDA masks / a non-CPU default device cannot - place a multi-GB ``[F,H,W]`` tensor on GPU. - """ - weights = torch.zeros(*full_shape, dtype=torch.float32, device="cpu") - for tile in tiles: - masks = tuple(m.detach().float().cpu() for m in tile.masks_1d) - region_shape = tuple(s.stop - s.start for s in tile.out_coords) - region = torch.ones(region_shape, dtype=torch.float32, device="cpu") - weights[tile.out_coords] += scale_by_masks_1d(region, masks) - return weights.clamp(min=1e-8) - - def create_tiles_from_intervals_and_mappers( intervals: LatentIntervals, mappers: list[MappingOperation], @@ -1189,21 +1098,6 @@ def is_tiled(self) -> bool: """True when this axis is split into more than one tile (or has overlap).""" return self.num_tiles > 1 or self.overlap > 0 - @classmethod - def from_tile_size(cls, dim_size: int, tile_size: int, overlap: int = 0) -> DimensionTilingConfig: - """Create config by computing ``num_tiles`` from dimension size and tile size. - Args: - dim_size: Total length of the dimension. - tile_size: Desired tile size. - overlap: Overlap between consecutive tiles. - Returns: - A ``DimensionTilingConfig`` with the computed ``num_tiles``. - """ - split_op = split_by_size(tile_size, overlap) - intervals = split_op(dim_size) - return cls(num_tiles=len(intervals.intervals), overlap=overlap) - - @dataclass(frozen=True) class DimensionSizeConfig: """Tile size and overlap for a single video axis (frames / height / width). @@ -1294,12 +1188,6 @@ def axis_split(cfg: DimensionTilingConfig, axis_min: int | None, *, temporal: bo axis_split(self.width, min_w, temporal=False), ) - def video_chunks_number(self, num_frames: int) -> int: - """Number of temporal decode chunks for ``num_frames`` under this layout.""" - del num_frames - return max(1, self.frames.num_tiles) - - @dataclass(frozen=True) class TileSizeConfig: """Size-based tiling layout for a ``(F, H, W)`` video — mirror of ``TileCountConfig``. @@ -1339,49 +1227,6 @@ def default(cls) -> TileSizeConfig: width=DimensionSizeConfig(tile_size=768, overlap=64), ) - @classmethod - def from_long_side( - cls, - *, - long_side: DimensionSizeConfig, - height: int, - width: int, - scale_factors: SpatioTemporalScaleFactors, - frames: DimensionSizeConfig | None = None, - ) -> TileSizeConfig: - """Aspect-coupled construction — old single-spatial long-side behavior, explicit. - Matches main-era ``latent_tile_splitters``: scale the long-side tile in - *latent* units with ``round(size_lat * axis_lat / long_lat)``, then - multiply back by the VAE factor. Pixel-space ``round`` + ceil-snap would - bias the short axis up by almost one latent (e.g. 680 → 704 vs 672). - Both axes share ``long_side.overlap``. - """ - if height < 1 or width < 1: - raise ValueError(f"height/width must be >= 1, got {height}x{width}") - if not long_side.is_tiled(): - raise ValueError("long_side must be tiled (tile_size > 0)") - if scale_factors.height < 1 or scale_factors.width < 1: - raise ValueError(f"scale_factors height/width must be >= 1, got {scale_factors}") - span = max(height, width) - - def axis_size(axis_len: int, factor: int) -> int: - # Latent-grid round (same as main decode enable_on_axis), not pixel ceil. - axis_lat = axis_len // factor - long_lat = span // factor - size_lat = long_side.tile_size // factor - overlap_lat = long_side.overlap // factor - lower_threshold = max(2, overlap_lat + 1) - tile_lat = max(lower_threshold, round(size_lat * axis_lat / long_lat)) - tile_px = tile_lat * factor - min_legal = max(2 * factor, long_side.overlap + factor) - return max(tile_px, min_legal) - - return cls( - frames=DimensionSizeConfig() if frames is None else frames, - height=DimensionSizeConfig(tile_size=axis_size(height, scale_factors.height), overlap=long_side.overlap), - width=DimensionSizeConfig(tile_size=axis_size(width, scale_factors.width), overlap=long_side.overlap), - ) - def to_splitters( self, scale_factors: SpatioTemporalScaleFactors, @@ -1421,26 +1266,6 @@ def enable_size_axis( enable_size_axis(scale_factors.width, min_w, self.width, "width", temporal=False), ) - def video_chunks_number(self, num_frames: int, *, time_scale: int = VIDEO_SCALE_FACTORS.time) -> int: - """Number of temporal decode chunks for ``num_frames`` under this layout. - Mirrors what decode actually does: :meth:`to_splitters` converts this axis to the - latent grid and hands it to :func:`split_by_size`, so the count must be taken there - too. Doing the arithmetic in pixel units instead over-reports by one whenever the - trailing tile is absorbed -- including the common case of a tile larger than the - clip, which is a single tile but used to report two. - """ - if not self.frames.is_tiled(): - return 1 - # Same derivation as ``to_splitters.enable_size_axis``. - overlap = self.frames.overlap // time_scale - size = max(2, overlap + 1, self.frames.tile_size // time_scale) - latent_frames = (num_frames - 1) // time_scale + 1 - if latent_frames <= size: - return 1 - # Same tile count as ``split_by_size``. - return (latent_frames + size - 2 * overlap - 1) // (size - overlap) - - TilingConfig = TileSizeConfig | TileCountConfig @@ -1531,20 +1356,6 @@ def _validate_overlap( raise ValueError(f"{axis_name} overlap {cfg.overlap} {unit} is below the required {recommended} {unit}.") -def balanced_tile_split(num_tiles: int) -> tuple[int, int]: - """Factor ``num_tiles`` into ``(small, large)`` as square as possible. - ``small`` is the largest divisor not exceeding the square root, so - ``small * large == num_tiles`` and ``small <= large``. E.g. 2 -> (1, 2), - 4 -> (2, 2), 8 -> (2, 4), 16 -> (4, 4). The caller decides which tiled - dimension gets which factor. - """ - if num_tiles < 1: - raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") - small = next(d for d in range(math.isqrt(num_tiles), 0, -1) if num_tiles % d == 0) - return small, num_tiles // small - - - class DiffVAEMode(Enum): COMBINED_COMPILE = "combined_compile" CHUNKED_COMPILE = "chunked_compile" @@ -1722,24 +1533,6 @@ def validate(self, *, num_frames: int | None = None) -> None: # keep the nearest plane on each side so |dt| matches a whole-clip decode. A far plane # on a full clip is the same geometry -- joint attention ranks it by distance. - def for_frame_span(self, frame_lo: int, frame_hi: int) -> "DecodeKeyframes": - """Keep the planes a decode of pixel frames ``[lo, hi]`` needs; indices stay global. - Selection is :func:`planes_for_tile`: every plane inside the span **plus the nearest - plane on each side outside it**. Those two are not optional. DiffVAE's joint attention - picks a frame's anchors by ``|dt|``, so a window ending at frame 64 whose last inside - plane is 48 still has to carry the plane at 96 -- drop it and frames near the boundary - anchor on 48 alone, which is exactly how a split decode stops matching a whole one. - ``pixel_frame_indices`` are not rewritten. :attr:`clip_start_frame` becomes ``frame_lo`` - so the decoder subtracts ``t_s(48) - t_s(56)`` rather than treating the slice as a new - clip that starts at pixel 0. - """ - keep = planes_for_tile(self.pixel_frame_indices, frame_lo, frame_hi) - return DecodeKeyframes( - latents=self.latents[:, :, keep.to(self.latents.device)], - pixel_frame_indices=self.pixel_frame_indices[keep.to(self.pixel_frame_indices.device)], - clip_start_frame=frame_lo, - ) - def crop_spatial(self, height: slice, width: slice) -> "DecodeKeyframes": """Crop the planes to a spatial window, with the *same* latent slices the video used. For a decode that splits the latent across workers (see @@ -3052,22 +2845,6 @@ def _abs_rope_op( ) -@_abs_rope_op.register_fake -def _abs_rope_fake( - x: torch.Tensor, - inv_t: torch.Tensor, - inv_h: torch.Tensor, - inv_w: torch.Tensor, - d_t: int, - d_h: int, - d_w: int, - num_tiles: int, - compute_dtype_is_bf16: bool, -) -> torch.Tensor: - del inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 - return torch.empty(x.shape, device=x.device, dtype=x.dtype) - - def _apply_opaque_abs_rope( x: torch.Tensor, rope_split: tuple[int, int, int], @@ -3120,23 +2897,6 @@ def _abs_rope_at_t_op( ) -@_abs_rope_at_t_op.register_fake -def _abs_rope_at_t_fake( - x: torch.Tensor, - t_pos: torch.Tensor, - inv_t: torch.Tensor, - inv_h: torch.Tensor, - inv_w: torch.Tensor, - d_t: int, - d_h: int, - d_w: int, - num_tiles: int, - compute_dtype_is_bf16: bool, -) -> torch.Tensor: - del t_pos, inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 - return torch.empty(x.shape, device=x.device, dtype=x.dtype) - - def _rope_config(attn: object, x: torch.Tensor) -> tuple[tuple[torch.Tensor, ...], int, torch.dtype]: """``(inv_freqs, num_tiles, compute_dtype)`` read off an attention module.""" inv_freqs = ( @@ -3258,15 +3018,6 @@ def __post_init__(self) -> None: if self.tile_size is not None and self.tile_size < 1: raise ValueError("tile_size must be >= 1") - @classmethod - def by_count(cls, num_tiles: int = DEFAULT_SWIGLU_TILES): - return cls(num_tiles=num_tiles) - - @classmethod - def by_size(cls, tile_size: int = DEFAULT_SWIGLU_TILE_SIZE): - return cls(tile_size=tile_size) - - DEFAULT_SWIGLU_TILE_SPEC = SwiGLUTileSpec(tile_size=DEFAULT_SWIGLU_TILE_SIZE) @@ -3315,16 +3066,6 @@ def forward(self, x): return swiglu_tiled(x, *swiglu_weights(self), self.tile) -def configure_swiglu_tile(module_root, *, num_tiles=None, tile_size=None): - if num_tiles is None and tile_size is None: - return - tile = SwiGLUTileSpec(num_tiles=num_tiles, tile_size=tile_size) - for module in module_root.modules(): - if isinstance(module, SwiGLU): - module.tile = tile - - - """Limited-workspace 3D neighborhood attention (NATTEN ``na3d`` semantics) in pure torch. Vendored from comfy-kitchen ``backends/eager/na.py`` (Apache-2.0) for DiffVAE hosts without natten or Triton. Queries are tiled; tiles that share window geometry stack @@ -4134,10 +3875,6 @@ def __call__(self, attn, q, k, v, keyframe_q, keyframe_k, keyframe_v, keyframe_t _NATTEN_AVAILABLE = False -def natten_available() -> bool: - return _NATTEN_AVAILABLE - - class NAAttentionCallable(Protocol): """A windowed 3D neighborhood-attention backend. Q/K/V arrive as ``(B, T, H, W, NH, HD)``, already normed, scaled and RoPE'd; @@ -4181,38 +3918,6 @@ def __call__( ) -> tuple[torch.Tensor, torch.Tensor]: ... -class NattenAttention(NAAttentionCallable): - """``natten.na3d``, the default backend. - ``backend`` pins ``na3d``'s own kernel choice (e.g. ``"cutlass-fna"``); ``None`` - leaves NATTEN's auto-pick (hopper-fna on H100, etc.). - """ - - def __init__(self, backend: str | None = None) -> None: - self._backend = backend - - def __call__( - self, - attn: NeighborhoodAttention3D, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - ) -> torch.Tensor: - if not _NATTEN_AVAILABLE: - raise ImportError( - "natten is required for NeighborhoodAttention3D. " - "Install with: uv sync --package ltx-core --extra natten " - '(or: uv pip install "natten==0.21.7+torch2130cu132" -f https://whl.natten.org; ' - "requires torch==2.13.0+cu132)" - ) - # scale=1.0: callers already applied ``attn.scale`` to Q. - # RMSNorm under bf16 autocast can leave Q/K in float32 while V stays bf16; - # natten requires a uniform dtype (same cast pattern as flash-attn paths). - if q.dtype != v.dtype or k.dtype != v.dtype: - q = q.to(dtype=v.dtype) - k = k.to(dtype=v.dtype) - return natten.na3d(q, k, v, kernel_size=attn.kernel_size, scale=1.0, backend=self._backend) - - class NeighborhoodAttention3D(nn.Module): """3D Neighborhood Attention with absolute RoPE + pluggable NA backend. Q/K receive absolute RoPE; attention is ``attention_function`` (NATTEN by @@ -4334,19 +4039,6 @@ def forward_with_keyframes( return out, dataclasses.replace(keyframes, x=keyframe_out) -def configure_w_chunks(module_root: nn.Module, w_chunks: int = 1) -> None: - """Set W-chunking on ``NeighborhoodAttention3D`` under ``module_root``. - When ``w_chunks > 1``, also sets ``rope_num_tiles=1`` so ``chunked.attn`` - owns the W axis (RoPE W-tiling would double-split). Pass only the diffusion - residual subtree — det-stage attention must keep its default RoPE tiling. - """ - for module in module_root.modules(): - if isinstance(module, NeighborhoodAttention3D): - module.w_chunks = w_chunks - if w_chunks > 1: - module.rope_num_tiles = 1 - - """NABlock and DiffusionNABlock parameter shells for DiffVAE. Pathway subclasses live in ``chunked/`` and ``combined/``; ``apply`` installs them via ``__class__`` swap (same pattern as ``Fp8CastLinear``). The shell owns @@ -6256,15 +5948,6 @@ def forward( """Decode via ``_decode_pixels`` with ``tiling_config=None`` (single full tile).""" return next(self._decode_pixels(sample, tiling_config=None, generator=generator)) - def tiled_decode( - self, - latent: torch.Tensor, - tiling_config: TilingConfig, - generator: torch.Generator | None = None, - ) -> Iterator[torch.Tensor]: - """Tiled decode: stages 1-3 once, stages 4-5 per tile, pixel blend.""" - yield from self._decode_pixels(latent, tiling_config, generator=generator) - def decode_video( self, latent: torch.Tensor, @@ -6291,16 +5974,6 @@ def to_rgb(frames: torch.Tensor) -> torch.Tensor: for chunk in self._decode_pixels(latent, tiling_config, generator=generator, as_fhwc=True): yield to_rgb(chunk) - def decode_single_frames( - self, - latents: Sequence[torch.Tensor], - generator: torch.Generator | Sequence[torch.Generator | None] | None = None, - ) -> Iterator[torch.Tensor]: - """Decode each latent as its own one-frame clip, yielding one RGB tensor per latent.""" - yield from iter_decoded_single_frames(self, latents, generator) - - - class LTX25DiffusionVideoDecoder(DiffusionVideoDecoder): """DiffSynth-facing decoder with one full/tiled/keyframe interface.""" From db59a437b24142ac491a084c40f6cddb3f7bf692 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 03:26:46 +0800 Subject: [PATCH 08/31] Point LTX-2.3 examples at the current spatial upscaler checkpoint The Lightricks/LTX-2.3 repo no longer ships ltx-2.3-spatial-upscaler-x2-1.0.safetensors; the registered ltx2_latent_upsampler hash matches the x2-1.1 file, so update the example scripts, docs and the registration example comment accordingly. --- diffsynth/configs/model_configs.py | 2 +- docs/en/Model_Details/LTX-2.md | 2 +- docs/zh/Model_Details/LTX-2.md | 2 +- examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py | 2 +- examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py | 2 +- examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py | 2 +- examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py | 2 +- .../LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py | 2 +- .../ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py | 2 +- examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py | 2 +- examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py | 2 +- examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py | 2 +- .../model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py | 2 +- examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py | 2 +- .../model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py | 2 +- .../LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py | 2 +- .../LTX-2.3-T2AV-IC-LoRA-Union-Control.py | 2 +- .../model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py | 2 +- examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py | 2 +- 19 files changed, 19 insertions(+), 19 deletions(-) diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py index 02ee04f35..7c2ea4f1b 100644 --- a/diffsynth/configs/model_configs.py +++ b/diffsynth/configs/model_configs.py @@ -850,7 +850,7 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter", }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors") + # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors") "model_hash": "aed408774d694a2452f69936c32febb5", "model_name": "ltx2_latent_upsampler", "model_class": "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler", diff --git a/docs/en/Model_Details/LTX-2.md b/docs/en/Model_Details/LTX-2.md index 6760cdd89..29f5654ca 100644 --- a/docs/en/Model_Details/LTX-2.md +++ b/docs/en/Model_Details/LTX-2.md @@ -39,7 +39,7 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/docs/zh/Model_Details/LTX-2.md b/docs/zh/Model_Details/LTX-2.md index 6ba595be5..baa674347 100644 --- a/docs/zh/Model_Details/LTX-2.md +++ b/docs/zh/Model_Details/LTX-2.md @@ -39,7 +39,7 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py index ad0b91834..e474c2f24 100644 --- a/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py index d23460f5e..81274195c 100644 --- a/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), ) diff --git a/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py index 0dd0854c5..03c5c1dd5 100644 --- a/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py index 047a0b0d5..03e9c6d48 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), ) diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py index 5708f0d6f..9c923018d 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py index 25883f681..ecc38e757 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py index d241f69da..dc1667379 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py @@ -21,7 +21,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py index 52345894e..ca172e7f9 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py index 022160bb9..1a5437174 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py index 5585efc67..54ee01e7c 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py index d23bd619d..a7e80b2cc 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py index 8d67de931..fa3d8f2c2 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py index 094289a7a..e14c87a67 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py index b015f06ca..2f0ef7420 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py index 65a6ebfe5..40185f9d8 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py @@ -21,7 +21,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py index a954b40e8..a37ae6b01 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), From d67b3df68a59502673495b77a472632d745b1682 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 03:54:53 +0800 Subject: [PATCH 09/31] Skip video decoder validation during training caching stages The decoder selector validated the loaded decoder component at unit execution time, which breaks split training for LTX-2/2.3 repackaged checkpoints whose stage 1 loads only the VAE encoder. Caching stages never decode, so skip the validation when the scheduler is in training mode; inference behavior is unchanged. --- diffsynth/pipelines/ltx2_audio_video.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index f528d930a..704014566 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -416,6 +416,13 @@ def process( "video_decode_kwargs": {}, "noise_generator": None, } + if pipe.scheduler.training: + # Caching stages never decode, so the decoder component is not required here. + return { + "video_decoder_name": None, + "video_decode_kwargs": {}, + "noise_generator": None, + } if not pipe.is_ltx25: if use_diffusion_vae: raise ValueError("Diffusion VAE decoding is only supported by LTX-2.5 checkpoints.") From 5506c890138c14083b6746dc734fa8cba52b5448 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 09:30:23 +0800 Subject: [PATCH 10/31] Strip docstrings from the LTX-2.5 DiffVAE module Remove 192 docstring blocks (~1000 lines) from the portable DiffVAE implementation; the module keeps its inline WHY comments. Distilled T2AV output remains bit-identical (max abs pixel diff 0) to the pre-change baseline. File size 6391 -> 5036 lines. --- diffsynth/models/ltx25_diffusion_video_vae.py | 1032 +---------------- 1 file changed, 2 insertions(+), 1030 deletions(-) diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index 600dd8923..b26614cef 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -1,10 +1,3 @@ -"""LTX-2.5 diffusion video decoder with eager attention and internal tiling. - -The implementation is intentionally self-contained so DiffSynth does not depend on -``ltx-core`` or a nested source package at runtime. It preserves checkpoint names -and provides full, tiled, and keyframe-aware decode through one public interface. -""" - from __future__ import annotations import dataclasses @@ -27,13 +20,9 @@ class Disposable: - """Compatibility marker used by the target implementation.""" - - + pass class VideoDecoder: - """Structural marker for video decoders.""" - - + pass def _clip_generators(count, generator): if isinstance(generator, Sequence): if len(generator) != count: @@ -55,7 +44,6 @@ def iter_decoded_single_frames(decoder, latents, generator=None): yield torch.cat(chunks, dim=0) - def get_timestep_embedding( timesteps: torch.Tensor, embedding_dim: int, @@ -64,24 +52,6 @@ def get_timestep_embedding( scale: float = 1, max_period: int = 10000, ) -> torch.Tensor: - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - Args - timesteps (torch.Tensor): - a 1-D Tensor of N indices, one per batch element. These may be fractional. - embedding_dim (int): - the dimension of the output. - flip_sin_to_cos (bool): - Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) - downscale_freq_shift (float): - Controls the delta between frequencies between dimensions - scale (float): - Scaling factor applied to the embeddings. - max_period (int): - Controls the maximum frequency of the embeddings - Returns - torch.Tensor: an [N x dim] Tensor of positional embeddings. - """ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" half_dim = embedding_dim // 2 @@ -169,12 +139,6 @@ def forward(self, timesteps: torch.Tensor) -> torch.Tensor: class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): - """ - For PixArt-Alpha. - Reference: - https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 - """ - def __init__( self, embedding_dim: int, @@ -197,10 +161,6 @@ def forward( class VideoPixelShape(NamedTuple): - """ - Shape of the tensor representing the video pixel array. Assumes BGR channel format. - """ - batch: int frames: int height: int @@ -209,13 +169,6 @@ class VideoPixelShape(NamedTuple): class SpatioTemporalScaleFactors(NamedTuple): - """ - Describes the spatiotemporal downscaling between decoded video space and - the corresponding VAE latent grid. - Field order matches the (frame/time, height, width) axis layout used by - latent tensors and meshgrid coordinates elsewhere in the codebase. - """ - time: int height: int width: int @@ -226,14 +179,6 @@ def default(cls) -> "SpatioTemporalScaleFactors": @classmethod def from_blocks(cls, blocks: list, patch_size: int) -> "SpatioTemporalScaleFactors": - """Derive the scale factors from a VAE encoder/decoder block list. - Each ``compress_*`` block halves (encoder) or doubles (decoder) its target - axes by a stride of 2, independent of any channel ``multiplier``. The initial - patchify contributes an extra ``patch_size`` of spatial compression. Deriving - the factors from the blocks keeps a single source of truth that stays correct - across VAE variants (e.g. the 32x32x8 default and the 16x16x4 variant) instead - of relying on a hardcoded constant. - """ spatial_steps = 0 temporal_steps = 0 for block_name, _ in blocks: @@ -248,13 +193,6 @@ def from_blocks(cls, blocks: list, patch_size: int) -> "SpatioTemporalScaleFacto class VideoLatentShape(NamedTuple): - """ - Shape of the tensor representing video in VAE latent space. - The latent representation is a 5D tensor with dimensions ordered as - (batch, channels, frames, height, width). Spatial and temporal dimensions - are downscaled relative to pixel space according to the VAE's scale factors. - """ - batch: int channels: int frames: int @@ -302,11 +240,6 @@ def upscale(self, scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTOR class AudioLatentShape(NamedTuple): - """ - Shape of audio in VAE latent space: (batch, channels, frames, mel_bins). - mel_bins is the number of frequency bins from the mel-spectrogram encoding. - """ - batch: int channels: int frames: int @@ -365,13 +298,6 @@ def from_video_pixel_shape( @dataclass(frozen=True) class Audio: - """ - Container for decoded audio samples and metadata. - Attributes: - waveform: Audio waveform tensor. - sampling_rate: Sampling rate (Hz) of the waveform. - """ - waveform: torch.Tensor sampling_rate: int @@ -381,18 +307,6 @@ def to(self, **kwargs: object) -> "Audio": @dataclass(frozen=True) class GeneratedKeyframeLayout: - """Where a state's generated-keyframe slot tokens live, and what they represent. - Recorded by :class:`~ltx_core.conditioning.types.keyframe_slots.VideoGeneratedKeyframeSlots` - when it appends the slots, so they can later be located and extracted *exactly* rather - than by assuming they are the trailing tokens. Conditioning items are applied in list - order and each appends to the end, so a state built with slots plus any other appending - conditioning item has no fixed trailing layout. - Attributes: - pixel_frame_indices: Target pixel-frame index of each slot, in token order. - tokens_per_keyframe: Number of tokens one slot occupies (one latent frame's worth). - first_token: Index of the first slot token in the token sequence. - """ - pixel_frame_indices: tuple[int, ...] tokens_per_keyframe: int first_token: int @@ -407,33 +321,6 @@ def num_tokens(self) -> int: @dataclass(frozen=True) class LatentState: - """ - State of latents during the diffusion denoising process. - Attributes: - latent: The current noisy latent tensor being denoised. - denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising). - positions: Positional indices for each latent element, used for positional embeddings. - clean_latent: Initial state of the latent before denoising, may include conditioning latents. - attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention, - 0 = no attention. None means full attention everywhere. Built incrementally by conditioning items. - keyframes_mask: Optional per-token marker of shape (B, T, 1) -- same layout as - ``denoise_mask`` -- non-zero on tokens whose latent encodes a *single standalone pixel - frame* rather than the usual multi-frame span. That set is the target's first latent - frame (the video encoder is causal, so its first temporal latent frame covers 1 pixel - frame while the rest cover 8) plus any generated keyframe slots. Selects the tokens - that receive the model's learned keyframe absolute-position embedding; ignored - entirely by models built without ``use_keyframes_abs_pos_embedding``. - generated_keyframe_layout: Set when generated keyframe slots were appended; locates them. - generated_keyframes: Populated by ``clear_conditioning`` when a layout is present: the - denoised slot content as an unpatchified ``(B, C, K, H, W)`` latent, one latent frame - per keyframe. Each frame must be decoded as a standalone one-frame clip, never as a - K-frame video -- a causal decode would blend slots that were never adjacent. - frozen: When True, this stream is held fixed: token denoising is disabled (``denoise_mask`` - should be all zeros; pipeline builders enforce that) and the scalar noise level used for - prompt / cross-modality AdaLN gates is forced to 0 when the state is converted for the - transformer. - """ - latent: torch.Tensor denoise_mask: torch.Tensor positions: torch.Tensor @@ -459,10 +346,6 @@ def clone(self) -> "LatentState": def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor: - """Root-mean-square (RMS) normalize `x` over its last dimension. - Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized - shape and forwards `weight` and `eps`. - """ return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) @@ -472,11 +355,6 @@ def to_velocity( denoised_sample: torch.Tensor, calc_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: - """ - Convert the sample and its denoised version to velocity. - Returns: - Velocity - """ if isinstance(sigma, torch.Tensor): sigma = sigma.to(calc_dtype).item() if sigma == 0: @@ -490,11 +368,6 @@ def to_denoised( sigma: float | torch.Tensor, calc_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: - """ - Convert the sample and its denoising velocity to denoised sample. - Returns: - Denoised sample - """ if isinstance(sigma, torch.Tensor): sigma = sigma.to(calc_dtype) return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) @@ -506,17 +379,6 @@ def compute_trapezoidal_mask_1d( ramp_right: int, left_starts_from_0: bool = False, ) -> torch.Tensor: - """ - Generate a 1D trapezoidal blending mask with linear ramps. - Args: - length: Output length of the mask. - ramp_left: Fade-in length on the left. - ramp_right: Fade-out length on the right. - left_starts_from_0: Whether the ramp starts from 0 or first non-zero value. - Useful for temporal tiles where the first tile is causal. - Returns: - A 1D tensor of shape `(length,)` with values in [0, 1]. - """ if length <= 0: raise ValueError("Mask length must be positive.") @@ -544,15 +406,6 @@ def compute_rectangular_mask_1d( left_ramp: int, right_ramp: int, ) -> torch.Tensor: - """ - Generate a 1D rectangular (pulse) mask. - Args: - length: Output length of the mask. - left_ramp: Number of elements at the start of the mask to set to 0. - right_ramp: Number of elements at the end of the mask to set to 0. - Returns: - A 1D tensor of shape `(length,)` with values 0 or 1. - """ if length <= 0: raise ValueError("Mask length must be positive.") @@ -574,23 +427,11 @@ class DimensionInterval: @dataclass(frozen=True) class DimensionIntervals: - """Intervals which a single dimension of the latent space is split into. - Each interval is defined by its start, end, left ramp, and right ramp. - The start and end are the indices of the first and last element (exclusive) in the interval. - Ramps are regions of the interval where the value of the mask tensor is - interpolated between 0 and 1 for blending with neighboring intervals. - The left ramp and right ramp values are the lengths of the left and right ramps. - """ - intervals: list[DimensionInterval] @dataclass(frozen=True) class LatentIntervals: - """Intervals which the latent tensor of given shape is split into. - Each dimension of the latent space is split into intervals based on the length along said dimension. - """ - original_shape: torch.Size dimension_intervals: tuple[DimensionIntervals, ...] @@ -609,7 +450,6 @@ def default_split_operation(length: int) -> DimensionIntervals: def untiled_mask_1d() -> torch.Tensor: - """Length-1 ones that broadcast over an untiled axis (historical ``None`` mask).""" return torch.ones(1) @@ -623,7 +463,6 @@ def default_mapping_operation( def _grow_last_tile_to_min(intervals: list[DimensionInterval], min_tile_size: int) -> list[DimensionInterval]: - """Grow a short last tile left to ``min_tile_size``; widen penultimate ``right_ramp``.""" if len(intervals) <= 1: return list(intervals) last = intervals[-1] @@ -640,7 +479,6 @@ def _grow_last_tile_to_min(intervals: list[DimensionInterval], min_tile_size: in def _validate_tile_intervals(intervals: list[DimensionInterval], *, dim_size: int, min_tile_size: int) -> None: - """Validate coverage, ramp/overlap consistency, and ``min_tile_size``.""" if not intervals or intervals[0].start != 0 or intervals[-1].end != dim_size: raise ValueError(f"tiles must cover [0, {dim_size})") for i, iv in enumerate(intervals): @@ -657,20 +495,6 @@ def _validate_tile_intervals(intervals: list[DimensionInterval], *, dim_size: in def split_by_size(size: int, overlap: int, min_tile_size: int | None = None) -> SplitOperation: - """Split a dimension into overlapping tiles of a given size. - Tiles are sized ``size`` with ``overlap`` shared elements between - consecutive tiles. The last tile may be shorter if the dimension - doesn't divide evenly. If ``min_tile_size`` is set and the last tile is - shorter, it is grown leftward (penultimate ``right_ramp`` widens); the - result is validated and invalid layouts raise ``ValueError``. - Args: - size: Target tile size (in axis units). - overlap: Overlap between consecutive tiles. - min_tile_size: Optional minimum tile length. ``None`` keeps legacy - short-last-tile behavior. - Returns: - A split operation that divides a dimension into tiles. - """ if size <= 0: raise ValueError(f"size must be > 0, got {size}") if overlap < 0 or overlap >= size: @@ -708,16 +532,6 @@ def split(dimension_size: int) -> DimensionIntervals: def split_temporal_causal(size: int, overlap: int, min_tile_size: int | None = None) -> SplitOperation: - """Split a temporal axis into overlapping tiles with causal handling. - Each tile after the first is shifted back by 1 and its left ramp is - increased by 1, ensuring causal continuity through the blend ramps. - Args: - size: Tile size in axis units. - overlap: Overlap between tiles in the same units. - min_tile_size: Optional floor forwarded to :func:`split_by_size`. - Returns: - Split operation that divides temporal dimension with causal handling. - """ non_causal_split = split_by_size(size, overlap, min_tile_size=min_tile_size) def split(dimension_size: int) -> DimensionIntervals: @@ -738,17 +552,6 @@ def split(dimension_size: int) -> DimensionIntervals: def split_by_count_temporal_causal( num_tiles: int, overlap: int = 0, min_tile_size: int | None = None ) -> SplitOperation: - """Split a temporal dimension by count with causal handling. - Wraps :func:`split_by_count` with the same causal adjustment as - :func:`split_temporal_causal`: each tile after the first is shifted - back by 1 and its left ramp is increased by 1. - Args: - num_tiles: Number of tiles. Must be >= 1. - overlap: Overlap between adjacent tiles (default 0). - min_tile_size: Optional floor forwarded to :func:`split_by_count`. - Returns: - A split operation that divides a temporal dimension into tiles. - """ non_causal_split = split_by_count(num_tiles, overlap, min_tile_size=min_tile_size) def split(dimension_size: int) -> DimensionIntervals: @@ -765,24 +568,6 @@ def split(dimension_size: int) -> DimensionIntervals: def split_at_seams(boundaries: Sequence[int], num_tiles: int, overlap: int = 0) -> SplitOperation: - """Split a dimension on boundary cells whose content is already known, dropping the overlap. - ``boundaries`` are the ``K + 1`` segment edges in grid cells, starting at 0 and ending at the - last cell of the dimension. The ``K`` segments are dealt largest-first so leftover segments go to the leading tiles; - ``num_tiles`` larger than ``K`` is clamped. Each tile but the first starts ``overlap`` cells - before the boundary it resumes after. That run-up is context only: it lands in the interval's - ``left_ramp``, which :func:`identity_mapping_operation` with ``rectangular=True`` masks to zero, - so the earlier tile keeps the boundary cell and this one contributes strictly after it. - The point of cutting here is that nothing needs blending. A ramp is what a pair of tiles needs - when neither of them knows the truth at the seam; on a boundary cell both reproduce the same - known frame, so averaging them only smears it. - Args: - boundaries: Segment edges in grid cells, strictly increasing, starting at 0. - num_tiles: Number of tiles. Must be >= 1. Extra tiles beyond the segment count are dropped. - overlap: Context cells each non-first tile denoises before the cell it resumes at, in grid - units. Clamped at the start of the dimension. - Returns: - A split operation that divides a dimension on ``boundaries``. - """ boundaries = tuple(boundaries) if num_tiles < 1: raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") @@ -820,23 +605,6 @@ def split(dim_size: int) -> DimensionIntervals: def split_by_count(num_tiles: int, overlap: int = 0, min_tile_size: int | None = None) -> SplitOperation: - """Split a dimension into a given number of tiles with overlap. - Computes the tile size as - ``(dim_size + overlap * (num_tiles - 1)) // num_tiles`` so that - ``num_tiles`` tiles of that size with ``overlap`` shared elements - cover the dimension evenly. Delegates to :func:`split_by_size` for - the actual interval construction. - When the total ``dim_size + overlap * (num_tiles - 1)`` is not evenly - divisible by ``num_tiles``, the first ``remainder`` tiles each absorb - one extra unit. - Args: - num_tiles: Number of tiles. Must be >= 1. - overlap: Overlap between adjacent tiles (default 0). Must be >= 0 - and less than the computed tile size. - min_tile_size: Optional floor forwarded to last-tile growth / validation. - Returns: - A split operation that divides a dimension into tiles. - """ if num_tiles < 1: raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") if overlap < 0: @@ -884,13 +652,6 @@ def identity_mapping_operation( *, rectangular: bool = False, ) -> tuple[list[slice], list[torch.Tensor]]: - """Map each DimensionInterval to an output region at the same position. - For every interval the output start/end matches the input start/end and a 1-D mask is built - from the interval's left_ramp and right_ramp. The default mask is trapezoidal (blend on the - ramps). ``rectangular=True`` drops the ramps outright: the overlap is context the tile denoised - but does not contribute. Pair that with a split whose ramps are one-sided, such as - :func:`split_at_seams` -- ramps on both sides of an interval would leave a hole between tiles. - """ mask_1d = compute_rectangular_mask_1d if rectangular else compute_trapezoidal_mask_1d out_slices: list[slice] = [] masks: list[torch.Tensor] = [] @@ -901,22 +662,6 @@ def identity_mapping_operation( class Tile(NamedTuple): - """ - Represents a single tile. - Attributes: - in_coords: - Tuple of slices specifying where to cut the tile from the INPUT tensor. - out_coords: - Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor. - masks_1d: - Per-dimension masks in OUTPUT units. - Untiled axes use a length-1 ones tensor (broadcasts). These are used - for separable blending (and for the dense ``blend_mask`` property). - Methods: - blend_mask: - Create a single N-D mask from the per-dimension masks. - """ - in_coords: tuple[slice, ...] out_coords: tuple[slice, ...] masks_1d: tuple[torch.Tensor, ...] @@ -942,10 +687,6 @@ def blend_mask(self) -> torch.Tensor: def scale_by_masks_1d(x: torch.Tensor, masks_1d: Sequence[torch.Tensor]) -> torch.Tensor: - """Multiply ``x`` by separable 1d masks with broadcasting. - ``len(masks_1d)`` must equal ``x.ndim``. Prefer float32 masks so bf16/fp16 ``x`` promotes. - Length-1 masks (untiled axes) broadcast over that dimension. - """ if len(masks_1d) != x.ndim: raise ValueError(f"masks_1d length {len(masks_1d)} != x.ndim {x.ndim}") out = x @@ -962,11 +703,6 @@ def masks_are_complementary( *, atol: float = 1e-5, ) -> bool: - """Return whether per-axis 1d blend masks partition unity (sum to 1). - Checks each axis independently over the unique out-slices on that axis - (cartesian tile products would otherwise multi-count the same 1d interval). - When True, weighted accumulation needs no denominator. - """ if not tiles: return True ndim = len(full_shape) @@ -1046,11 +782,6 @@ def create_tiles( def group_tiles_by_temporal_slice(tiles: list[Tile]) -> list[list[Tile]]: - """Group consecutive tiles that share the same temporal ``out_coords`` slice. - Assumes ``tiles`` is ordered with the temporal axis varying slowest (true - for every tile list this codebase builds via ``itertools.product`` with - the temporal axis first), so equal temporal slices are always contiguous. - """ if not tiles: return [] @@ -1075,16 +806,6 @@ def group_tiles_by_temporal_slice(tiles: list[Tile]) -> list[list[Tile]]: @dataclass(frozen=True) class DimensionTilingConfig: - """Tiling parameters for a single dimension of the patchified grid. - Attributes: - num_tiles: Number of tiles along this dimension. ``1`` with ``overlap=0`` - means the axis is not tiled. - overlap: Overlap between adjacent tiles, in latent grid units. - Adjacent tiles share ``overlap`` grid cells at their - boundary, producing an overlap zone blended with - trapezoidal masks. - """ - num_tiles: int = 1 overlap: int = 0 @@ -1095,18 +816,10 @@ def __post_init__(self) -> None: raise ValueError(f"overlap must be >= 0, got {self.overlap}") def is_tiled(self) -> bool: - """True when this axis is split into more than one tile (or has overlap).""" return self.num_tiles > 1 or self.overlap > 0 @dataclass(frozen=True) class DimensionSizeConfig: - """Tile size and overlap for a single video axis (frames / height / width). - Mirrors :class:`DimensionTilingConfig`, but specifies tile *size* rather than - tile *count*. ``tile_size=0`` means the axis is not tiled (covers the whole - length). Axis-specific VAE pixel constraints (divisibility / minimums) are - enforced by :meth:`TileSizeConfig.validate` for tiled axes only. - """ - tile_size: int = 0 overlap: int = 0 @@ -1123,29 +836,16 @@ def __post_init__(self) -> None: raise ValueError(f"Overlap must be less than tile size, got {self.overlap} and {self.tile_size}") def is_tiled(self) -> bool: - """True when this axis has a positive tile size (caller intends to split it).""" return self.tile_size > 0 @dataclass(frozen=True) class TileCountConfig: - """Tiling layout for a ``(F, H, W)`` grid by tile *counts*. - Overlaps are in latent-grid units. Mirror of :class:`TileSizeConfig`. - Attributes: - frames: Tiling along the temporal (frames) dimension. - height: Tiling along the latent height dimension. - width: Tiling along the latent width dimension. - """ - frames: DimensionTilingConfig = DimensionTilingConfig() height: DimensionTilingConfig = DimensionTilingConfig() width: DimensionTilingConfig = DimensionTilingConfig() def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: - """Raise if this count layout cannot tile ``video_shape`` under ``scale_factors``. - Counts/overlaps are in latent-grid units. ``video_shape.frames <= 0`` skips the - temporal axis (duration not yet known). Spatial axes always checked. - """ check_temporal = _assert_video_on_vae_grid(scale_factors, video_shape) latent_h = video_shape.height // scale_factors.height latent_w = video_shape.width // scale_factors.width @@ -1162,14 +862,6 @@ def to_splitters( *, causal_temporal: bool = True, ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: - """Build ``(T, H, W)`` latent-grid split operations for this count layout. - ``scale_factors`` is accepted for signature parity with - :meth:`TileSizeConfig.to_splitters` and ignored — counts are already in - grid units. When ``causal_temporal`` is True (VAE encode/decode), the - frames axis uses :func:`split_by_count_temporal_causal`; otherwise plain - :func:`split_by_count`. ``min_tile_size`` is a per-axis floor in the same - units as the split. - """ del scale_factors min_t = min_h = min_w = None if min_tile_size is not None: @@ -1190,29 +882,11 @@ def axis_split(cfg: DimensionTilingConfig, axis_min: int | None, *, temporal: bo @dataclass(frozen=True) class TileSizeConfig: - """Size-based tiling layout for a ``(F, H, W)`` video — mirror of ``TileCountConfig``. - Each axis is a non-optional :class:`DimensionSizeConfig`; ``tile_size=0`` means - untiled on that axis (:meth:`DimensionSizeConfig.is_tiled`). Sizes and overlaps - are in pixel / frame units. Conversion to a split grid is an explicit - ``scale_factors`` argument to :meth:`to_splitters` (not stored on the config). - Legality vs a VAE grid is checked by :meth:`validate` (same factors decode will - pass to :meth:`to_splitters`), not at construction. - Attributes: - frames: Temporal tile size/overlap in frames. - height: Spatial height tile size/overlap in pixels. - width: Spatial width tile size/overlap in pixels. - """ - frames: DimensionSizeConfig = DimensionSizeConfig() height: DimensionSizeConfig = DimensionSizeConfig() width: DimensionSizeConfig = DimensionSizeConfig() def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: - """Raise if this size layout is illegal for ``video_shape`` under ``scale_factors``. - Checks tile/overlap divisibility and minimums against the VAE grid, and that - the video extents are compatible with that grid. ``video_shape.frames <= 0`` - skips the temporal axis (duration not yet known); height/width always checked. - """ check_temporal = _assert_video_on_vae_grid(scale_factors, video_shape) _validate_size_axis(self.height, scale_factors.height, "height") _validate_size_axis(self.width, scale_factors.width, "width") @@ -1234,9 +908,6 @@ def to_splitters( *, causal_temporal: bool = True, ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: - """Build ``(T, H, W)`` grid split ops from pixel/frame sizes via ``scale_factors``. - When ``causal_temporal`` is True, frames use :func:`split_temporal_causal`. - """ min_t = min_h = min_w = None if min_tile_size is not None: min_t, min_h, min_w = min_tile_size @@ -1270,10 +941,6 @@ def enable_size_axis( class AutoTiling: - """Sentinel: pipeline should recommend decode tiling (DiffVAE-aware / Conv default). - Distinct from ``None``, which means untiled decode. - """ - __slots__ = () def __repr__(self) -> str: @@ -1290,10 +957,6 @@ def _assert_video_on_vae_grid( scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape, ) -> bool: - """Raise if ``video_shape`` is incompatible with the VAE ``scale_factors`` grid. - Returns whether the temporal axis is known (``video_shape.frames > 0``). When - False, callers skip frames-axis checks (duration not yet resolved). - """ if scale_factors.time < 1 or scale_factors.height < 1 or scale_factors.width < 1: raise ValueError(f"scale_factors must be >= 1 on each axis, got {scale_factors}") if video_shape.height < 1 or video_shape.width < 1: @@ -1310,7 +973,6 @@ def _assert_video_on_vae_grid( def _validate_size_axis(cfg: DimensionSizeConfig, factor: int, axis_name: str) -> None: - """Pixel/frame size-axis legality vs VAE ``factor``.""" if not cfg.is_tiled(): return min_size = 2 * factor @@ -1323,7 +985,6 @@ def _validate_size_axis(cfg: DimensionSizeConfig, factor: int, axis_name: str) - def _validate_count_axis(cfg: DimensionTilingConfig, latent_extent: int, axis_name: str) -> None: - """Latent count-axis legality vs latent ``extent``.""" if not cfg.is_tiled(): return if cfg.num_tiles > latent_extent: @@ -1343,7 +1004,6 @@ def _validate_overlap( min_overlap_frames: int, min_overlap_pixels: int, ) -> None: - """Raise if any tiled ``TileSizeConfig`` axis overlap is below the given floors.""" if not isinstance(tiling_config, TileSizeConfig): return @@ -1387,18 +1047,7 @@ def frames_per_yuv_gemm(height: int, width: int) -> int: return 2**31 - 1 - def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: - """ - Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks - and moves pixels from each block into separate channels (space-to-depth). - Args: - x: Input tensor (4D or 5D) - patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks. - patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching). - For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw) - Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128) - """ if patch_size_hw == 1 and patch_size_t == 1: return x if x.dim() == 4: @@ -1418,16 +1067,6 @@ def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torc def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: - """ - Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from - channels back into patch_size x patch_size blocks (depth-to-space). - Args: - x: Input tensor (4D or 5D) - patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x. - patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion). - For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw) - Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512) - """ if patch_size_hw == 1 and patch_size_t == 1: return x @@ -1446,13 +1085,6 @@ def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> to class PerChannelStatistics(nn.Module): - """ - Per-channel statistics for normalizing and denormalizing the latent representation. - This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict. - Defaults are identity (std=1, mean=0) so models constructed without a checkpoint - do not inherit allocator garbage / NaNs from ``torch.empty``. - """ - def __init__(self, latent_channels: int = 128): super().__init__() self.register_buffer("std-of-means", torch.ones(latent_channels)) @@ -1488,27 +1120,11 @@ def normalize(self, x: torch.Tensor) -> torch.Tensor: @dataclass(frozen=True) class DecodeKeyframes: - """Caller-facing keyframe input to a DiffVAE decode. - Attributes: - latents: ``(B, C, P, H, W)`` per-channel-normalized latents, exactly one latent - frame per keyframe. Each plane must have been encoded as a standalone - one-pixel-frame clip -- the VAE is causal, so a ``P``-frame encode would - blend planes that were never temporally adjacent. - pixel_frame_indices: ``(P,)`` int64 **global** pixel frame index of each plane. - Never rebased onto a tile. A Dist slice whose first pixel is 56 still carries - a plane at 48 as ``48``; :attr:`clip_start_frame` is how DiffVAE learns the - 8-frame gap. - clip_start_frame: first global pixel frame of the video latent in this decode - call. ``0`` for a full-clip decode. Dist sets it to the tile origin so stage - times are ``t_s(index) - t_s(clip_start)``. - """ - latents: torch.Tensor pixel_frame_indices: torch.Tensor clip_start_frame: int = 0 def validate(self, *, num_frames: int | None = None) -> None: - """Raise if shapes/indices are inconsistent (optionally against a frame count).""" if self.latents.ndim != 5: raise ValueError(f"keyframe latents must be (B, C, P, H, W), got {tuple(self.latents.shape)}") if self.pixel_frame_indices.ndim != 1: @@ -1534,15 +1150,6 @@ def validate(self, *, num_frames: int | None = None) -> None: # on a full clip is the same geometry -- joint attention ranks it by distance. def crop_spatial(self, height: slice, width: slice) -> "DecodeKeyframes": - """Crop the planes to a spatial window, with the *same* latent slices the video used. - For a decode that splits the latent across workers (see - :class:`~ltx_core.multigpu.vae.distributed_decoder.DistributedVideoDecoder`): each worker - holds a crop of the video latent, so it must hold the matching crop of every keyframe - plane. Cropping one and not the other offsets every plane from the video by the - difference, which reads as ghosting rather than as an obvious failure. - Plane count, ``pixel_frame_indices``, and :attr:`clip_start_frame` are untouched -- a - spatial split leaves every worker the full frame range. - """ return DecodeKeyframes( latents=self.latents[:, :, :, height, width], pixel_frame_indices=self.pixel_frame_indices, @@ -1556,40 +1163,19 @@ def num_planes(self) -> int: @dataclass(frozen=True) class KeyframeStream: - """The keyframe half of the dual stream at one decoder stage. - Attributes: - x: ``(B, P, H, W, C)`` channels-last activations. ``H``/``W`` always match the - video stream at the same stage; ``P`` is invariant across the whole decode. - times: ``(P,)`` float32 plane position in *this stage's* temporal units and - *local to the current tile* -- the same origin the video stream's RoPE uses. - Both streams must share one origin or the joint softmax sees wrong offsets. - valid: ``(P,)`` bool. Invalid planes are masked out of every softmax and their - activations are re-zeroed after each upsample. - """ - x: torch.Tensor times: torch.Tensor valid: torch.Tensor def masked(self) -> KeyframeStream: - """Re-zero invalid planes' activations (channels-last plane axis).""" return KeyframeStream(x=self.x * self.valid[None, :, None, None, None], times=self.times, valid=self.valid) def select_planes(self, keep: torch.Tensor) -> KeyframeStream: - """Subset the plane axis, keeping ``x``/``times``/``valid`` in step. - ``keep`` is a ``(P,)`` bool mask. Used by tiled decode, which gives each tile only the - planes near it -- see :func:`planes_for_tile`. - """ if keep.shape != (self.num_planes,): raise ValueError(f"keep must be ({self.num_planes},) bool, got {tuple(keep.shape)}") return KeyframeStream(x=self.x[:, keep], times=self.times[keep], valid=self.valid[keep]) def crop_spatial(self, height: slice, width: slice) -> KeyframeStream: - """Crop H/W with the *same* slices the video stream's tile used. - Cropping only one stream offsets every plane from the video by the difference, which - reads as ghosting rather than as an obvious failure -- the same hazard as the spatial - padding rule. - """ return KeyframeStream(x=self.x[:, :, height, width, :], times=self.times, valid=self.valid) @property @@ -1598,15 +1184,6 @@ def num_planes(self) -> int: def keyframe_stage_times(pixel_frame_indices: torch.Tensor, remaining_time_stride: int) -> torch.Tensor: - """Chunk-center position of each keyframe in a stage's temporal units. - A stage whose remaining temporal upsampling is ``r`` has cells covering ``r`` pixel - frames each, except cell 0 which covers only pixel frame 0 (the causal first frame). - So ``t_s(0) = 0`` and ``t_s(f) = (f + (r - 1) / 2) / r`` -- the center of the chunk - holding ``f``. At stage 5 ``r == 1``, making the times the raw pixel indices. - Args: - pixel_frame_indices: ``(P,)`` global pixel frame index per plane. - remaining_time_stride: product of the temporal strides *still to come*. - """ if remaining_time_stride < 1: raise ValueError(f"remaining_time_stride must be positive, got {remaining_time_stride}") frames = pixel_frame_indices.to(torch.float32) @@ -1621,15 +1198,6 @@ def keyframe_clip_times( clip_start_frame: int, extra_origin: float = 0.0, ) -> torch.Tensor: - """Stage times relative to a decode whose first pixel frame is ``clip_start_frame``. - ``t_s(global) - t_s(clip_start)`` is the gap joint attention should see. Rebasing the - indices onto the tile (``48 -> -8``) and then calling :func:`keyframe_stage_times` is not - the same: ``t_s`` is not linear through a fake clip start, so stages with ``r > 1`` get - the wrong ``|dt|``. - Single-GPU tiled decode uses ``clip_start_frame=0`` and passes the in-volume tile origin - as ``extra_origin``. A Dist slice whose first pixel is global 56 uses - ``clip_start_frame=56`` and ``extra_origin=0``. - """ times = keyframe_stage_times(pixel_frame_indices, remaining_time_stride) origin = keyframe_stage_times( torch.as_tensor([clip_start_frame], dtype=torch.int64, device=pixel_frame_indices.device), @@ -1645,22 +1213,6 @@ def planes_for_tile( *, clip_start_frame: int = 0, ) -> torch.Tensor: - """``(P,)`` bool: which planes a tile spanning pixel frames ``[lo, hi]`` should carry. - Every plane inside the span, **plus the nearest plane on each side outside it**. Those two - boundary planes are the point: without them a video frame at a tile edge ranks only - in-tile planes and attends to the wrong one, which is what made tiled keyframe decode - non-invariant. They arrive with negative / past-the-end tile-local times, and the - ``(|dt|, index)`` slot ranking already handles those, so nothing downstream changes. - Selection is by *value*, not position: ``pixel_frame_indices`` is not required to be - sorted. - Args: - pixel_frame_indices: ``(P,)`` global pixel frame index per plane. - frame_lo: first pixel frame in the tile, relative to ``clip_start_frame``. - frame_hi: last pixel frame in the tile (inclusive), relative to ``clip_start_frame``. - clip_start_frame: first global pixel of this latent. A full-clip decode leaves it 0. - Dist tiles keep global indices and pass the slice origin so a local ``[0, 72)`` - still selects global ``[56, 127]``. - """ frame_lo = frame_lo + clip_start_frame frame_hi = frame_hi + clip_start_frame indices = pixel_frame_indices.to(torch.int64) @@ -1678,10 +1230,6 @@ def planes_for_tile( def remaining_time_strides(upsamples: Sequence[torch.nn.Module]) -> tuple[int, ...]: - """Remaining temporal upsampling at each stage input, plus 1 for stage 5. - For the production ladder (temporal strides ``1, 2, 2, 2``) this is - ``(8, 8, 4, 2, 1)``: stage ``i``'s blocks see the product of strides ``i..end``. - """ strides = [int(up.stride[0]) for up in upsamples] remaining: list[int] = [] for index in range(len(strides)): @@ -1694,17 +1242,6 @@ def remaining_time_strides(upsamples: Sequence[torch.nn.Module]) -> tuple[int, . def upsample_keyframe_planes(upsample: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: - """Spatially upsample keyframe planes, keeping the plane count invariant. - Each plane is folded into the batch as its own ``T=1`` clip and pushed through the - *same* upsample module as video, always with ``drop_leading_frame=True``: a temporal - stride of 2 expands ``T=1`` to 2 and the leading-frame drop takes it back to 1, - keeping phase 1. So temporal strides collapse and only ``H``/``W`` grow. - Passing ``drop_leading_frame=False`` here (as tiled video does for non-origin tiles) - would invent a second temporal plane per keyframe and is always wrong. - Args: - upsample: the video stream's ``LinearPixelShuffleUpsample`` for this stage. - x: ``(B, P, H, W, C)`` keyframe activations. - """ planes = x.shape[1] flat = rearrange(x, "b p h w c -> (b p) 1 h w c") upsampled = upsample(flat, drop_leading_frame=True) @@ -1722,10 +1259,6 @@ def _nearest_slots( candidate_valid: torch.Tensor | None, num_slots: int, ) -> torch.Tensor: - """``(Q, num_slots)`` candidate indices ranked by ``(|dt|, index)``, ``-1`` when empty. - A stable argsort on ``|dt|`` breaks ties by ascending candidate index, which is - exactly upstream's ``distances + arange * 1e-6`` tie-break. - """ distances = (query_times[:, None] - candidate_times[None, :]).abs().to(torch.float32) if candidate_valid is not None: distances = distances.masked_fill(~candidate_valid[None, :], float("inf")) @@ -1747,10 +1280,6 @@ def video_keyframe_slots( video_length: int, num_slots: int = KEYFRAME_CONTEXT_SLOTS, ) -> torch.Tensor: - """``(T, num_slots)`` keyframe plane index per video frame, ``-1`` for an empty slot. - Ranked by ``(|t_s(plane) - t|, plane)``. Deliberately independent of the temporal - kernel: the nearest planes are visible even when they lie outside ``K_t``. - """ query = torch.arange(video_length, dtype=torch.float32, device=keyframe_times.device) return _nearest_slots(query, keyframe_times.to(torch.float32), keyframe_valid, num_slots) @@ -1761,9 +1290,6 @@ def keyframe_video_slots( video_length: int, num_slots: int = KEYFRAME_CONTEXT_SLOTS, ) -> torch.Tensor: - """``(P, num_slots)`` video frame index per keyframe plane, ``-1`` for an empty slot. - Ranked by ``(|t' - t_s(plane)|, t')``. Rows of invalid planes are all ``-1``. - """ candidates = torch.arange(video_length, dtype=torch.float32, device=keyframe_times.device) slots = _nearest_slots(keyframe_times.to(torch.float32), candidates, None, num_slots) return torch.where(keyframe_valid[:, None], slots, torch.full_like(slots, -1)) @@ -1783,8 +1309,6 @@ def keyframe_video_slots( @dataclass(frozen=True, slots=True) class _StageFiveBudget: - """One mode's stage-5 multiplicity and withheld reserve, with and without keyframes.""" - coef: float coef_keyframes: float reserve_bytes: int @@ -1814,20 +1338,11 @@ class _StageFiveBudget: def _falls_back_to_eager_na(mode: DiffVAEMode) -> bool: - """True when this host has no natten and the mode's NA remaps to Triton/eager. - For chunked modes that remap also switches ``compile_blocks`` off, so the peak is the eager - one; :func:`resolve_attention_for_host` is the single owner of that decision. - """ resolved = resolve_attention_for_host(mode.resolve()) return resolved.attention in (NAttentionKind.TRITON, NAttentionKind.EAGER_SDPA) and not resolved.compile_blocks def stage5_mem_coef(mode: DiffVAEMode, *, keyframes: bool = False) -> float: - """Stage-5 working-set multiplicity for auto tiling, after host NA resolve. - Args: - mode: the decode preset. - keyframes: whether this decode carries a keyframe stream, which runs eager blocks. - """ try: budget = _BUDGET_BY_MODE[mode] except KeyError as exc: @@ -1849,14 +1364,6 @@ def stage5_mem_coef(mode: DiffVAEMode, *, keyframes: bool = False) -> float: def max_emitted_frames(*, num_frames: int, tile_frames: int, overlap_frames: int) -> int: - """Longest chunk the tiled decode yields, in pixel frames. - ``_decode_groups_with_keyframes`` yields only a group's *exclusive* span -- one - stride -- and keeps the trailing overlap as a stub for the next group. Only the - final group yields its whole buffer. Charging ``tile_frames`` for every chunk - therefore roughly doubles the estimate on long clips, which blocks layouts that - would have fit. ``+1`` covers the causal shift, which moves each group after the - first back by one frame. - """ if tile_frames >= num_frames: return num_frames stride = tile_frames - overlap_frames @@ -1875,19 +1382,6 @@ def emit_convert_bytes( out_channels: int, element_size: int, ) -> int: - """Downstream bytes a yielded chunk costs while the encoder consumes it. - The decode budget alone is not enough to size a tile: whatever ``_emit`` yields is - handed straight to the video encoder, and that peak overlaps the decode because the - decoder is a generator -- it stays suspended holding stage-4 features and the - accumulator while the consumer converts the chunk it just yielded. A layout that - decodes comfortably can therefore still die in the encoder, and because the - recommender spends spare VRAM on *larger* temporal tiles, more free memory used to - make that failure more likely rather than less. - ``out_channels`` is the width of the intermediate **YUV** tensor, not a second copy - of the emitted RGB: the chunk itself is charged by the accumulator term. On the - write-back path YUV lands in that same RGB storage, so it is not charged here at all - and only the GEMM temporary is. - """ frames = int(tile_frames) gemm_frames = min(frames, frames_per_yuv_gemm(height, width)) # Write-back reuses the RGB storage, so no full YUV tensor survives into the pack. @@ -1909,14 +1403,6 @@ def budget_safety_bytes( keyframes: bool = False, joint_sdpa_materializes: bool = False, ) -> int: - """Extra bytes withheld from the recommend budget. - Args: - mode: the decode preset. - keyframes: whether this decode carries a keyframe stream. - joint_sdpa_materializes: whether the joint attention will run on torch's MATH SDPA kernel - (see ``fallback_na.joint_eager.sdpa_materializes_scores``). Ignored without keyframes, - and irrelevant when a fused joint kernel serves the decode. - """ try: budget = _BUDGET_BY_MODE[mode] except KeyError as exc: @@ -1929,9 +1415,6 @@ def budget_safety_bytes( def accumulator_element_size(feature_dtype: torch.dtype) -> int: - """Bytes per accumulator element; mirrors ``_decode_temporal_group_isolated``. - ``accum_dtype = float16 if feat_s4.dtype == bfloat16 else feat_s4.dtype``. - """ if feature_dtype is torch.bfloat16: return 2 # stored as fp16 return int(torch.tensor([], dtype=feature_dtype).element_size()) @@ -1947,10 +1430,6 @@ def stage4_feature_bytes( element_size: int = _DEFAULT_ELEMENT_SIZE, natten_trailing_pad_latent_frames: int = 0, ) -> int: - """Resident stages-1-3 output size (full volume tiled into stage 4). - Matches ``DiffusionVideoDecoder.forward_stages_1_to_3`` after optional NATTEN - trailing latent pad: channels-last ``(B, T, H, W, C)`` at stage-4 input resolution. - """ if stage4_channels < 1: raise ValueError(f"stage4_channels must be >= 1, got {stage4_channels}") if element_size < 1: @@ -1998,37 +1477,6 @@ def recommended_decode_tiling_config( # noqa: PLR0913 keyframes: bool = False, joint_sdpa_materializes: bool = False, ) -> TileSizeConfig: - """Pick DiffVAE decode tiling from stage-4/5 halos and free VRAM. - Always enables both spatial and temporal tiling (temporal-only full-frame slabs - are unsafe on Hopper / some natten builds). - Selection (size-grid, accumulator-aware): - 1. Enumerate legal tile **sizes** on the LCM of DiffVAE ``pixel_scale`` and - :data:`~ltx_core.types.VIDEO_SCALE_FACTORS` (so configs also pass - :class:`~ltx_core.tiling.TileSizeConfig` construction); derive tile - counts from :func:`~ltx_core.tiling.split_by_size` (same as decode). - 2. Drop triples whose peak-bytes estimate exceeds ``usable`` bytes - (``free - max(model, 1 GiB) - safety - stage4_feature``; safety is - 1 GiB eager / 2 GiB compiled, see :func:`budget_safety_bytes`). - Stage-4 input features stay resident for the whole tiled decode. - 3. Among feasible triples, pick minimal :func:`volumetric_overlap_waste`. - Peak-bytes estimate:: - stage4_feature_bytes(...) # hard, full volume - + H * W * (2 * tile_t) * out_channels * element_size - + stage5_tokens * stage5_channels * element_size * coef - Accumulator is full output HxW (not spatially tiled) with temporal extent - ``2 * tile_t``: current group buffer plus the still-live previous exclusive - emit / overlap stub during handoff (not merely ``tile_t + overlap_t``). - RGBx``element_size`` by default. ``element_size`` is the activation width: - production bf16 features use fp16 accumulators (2), matching - :func:`accumulator_element_size`. Stage-5 uses the same element size x - ``stage5_channels`` x ``coef``, which :func:`stage5_mem_coef` reads off the - mode and ``keyframes`` (11 / 7 / 5 / 2.5 by mode; 15 / 5 / 5 / 2.5 with a - keyframe stream, which runs eager blocks). - Args beyond the geometry: - keyframes: this decode carries a keyframe stream (joint attention, eager blocks). - joint_sdpa_materializes: the joint attention will run on torch's MATH SDPA kernel; - costs one more GiB of reserve. See :func:`budget_safety_bytes`. - """ if height < 1 or width < 1 or num_frames < 1: raise ValueError(f"height/width/num_frames must be >= 1, got {height}x{width}x{num_frames}") if patch_size < 1: @@ -2146,14 +1594,6 @@ def prepare_tile_schedule( min_tile_size: Tuple[int, int, int], tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], ) -> List[Tile]: - """Build pixel-blend tiles whose ``in_coords`` land on the stage-4 input grid. - DiffVAE temporal tiling deliberately skips ConvVAE causal split/mask tricks - (``split_temporal_causal``, ``left_starts_from_0``): pixel overlap already - covers blend+halo, and interval propagation follows - :class:`~ltx_core.model.video_vae.transformer.layers.LinearPixelShuffleUpsample` - (``drop_leading_frame`` only on the origin tile) with *symmetric* trapezoid - ramps so masks stay complementary without a weight buffer. - """ pixel_scale = stage4_to_pixel_scale_factors(upsample3_stride, patch_size) if tiling_config is None: return [ @@ -2230,7 +1670,6 @@ def slice_stage4_tile( *, content_frames: int, ) -> tuple[torch.Tensor, bool, bool, tuple[int, int, int]]: - """Slice a stage-4 feature tile, extending trailing tiles to include ghost frames.""" is_origin = tile.in_coords[1].start in (0, None) _, stop, _ = tile.in_coords[1].indices(content_frames) pad_trailing = stop == content_frames @@ -2247,8 +1686,6 @@ def slice_stage4_tile( @dataclass(frozen=True) class AxisPad: - """How many elements were added (pad) or removed (crop) on each side of one axis.""" - before: int after: int @@ -2260,15 +1697,6 @@ def resize_axis( *, mode: ResizeAxisMode, ) -> tuple[torch.Tensor, AxisPad]: - """Pad or crop axis ``dim`` so its length becomes ``size``. - Pad (``len < size``): - ``repeat_last`` - append copies of the last slice. - ``symmetric`` - edge-replicate first/last; leftover goes to the end - (``before = need // 2``, ``after = need - before``). - Crop (``len > size``): - ``repeat_last`` - drop from the end. - ``symmetric`` - drop from both ends with the same split rule as pad. - """ if size < 1: raise ValueError(f"resize_axis target size must be >= 1, got {size}") if dim < 0: @@ -2318,7 +1746,6 @@ def ensure_min_latent_shape( latent: torch.Tensor, min_tile_sizes: Tuple[int, int, int], ) -> tuple[torch.Tensor, tuple[AxisPad, AxisPad, AxisPad]]: - """Pad latent ``(B, C, T, H, W)`` up to ``min_tile_sizes`` if needed.""" min_t, min_h, min_w = min_tile_sizes t_pad = AxisPad(0, 0) h_pad = AxisPad(0, 0) @@ -2334,7 +1761,6 @@ def ensure_min_latent_shape( def scale_axis_pad(pad: AxisPad, scale: int) -> AxisPad: - """Scale a latent-grid ``AxisPad`` into pixel (or other) units.""" return AxisPad(pad.before * scale, pad.after * scale) @@ -2348,12 +1774,6 @@ def crop_pixels_to_content( w_pad: AxisPad | None = None, spatial_scale: Tuple[int, int] = (1, 1), ) -> torch.Tensor: - """Crop padded decode output ``(B, C, F, H, W)`` back to the content shape. - Temporal pad is always trailing (``repeat_last``), so T is cropped from the - end. Spatial size-floor pads must pass the recorded ``h_pad`` / ``w_pad`` - (latent units) plus ``spatial_scale`` ``(H, W)`` so odd leftovers are not - re-split by a center-crop after upscaling. - """ x, _ = resize_axis(pixels, 2, frames, mode="repeat_last") scale_h, scale_w = spatial_scale if h_pad is not None: @@ -2384,7 +1804,6 @@ def stage5_pixel_shape_from_stage4( drop_leading_frame: bool, pad_trailing: bool, ) -> tuple[int, int, int]: - """Pixel ``(F, H, W)`` for a stage-4-input extent (one remaining NA hop + patch).""" st, sh, sw = upsample_stride frames = stage4_t * st - 1 if drop_leading_frame and st == 2 else stage4_t * st if pad_trailing: @@ -2393,7 +1812,6 @@ def stage5_pixel_shape_from_stage4( def pad_trailing_latent_for_natten_border(latent: torch.Tensor, n_frames: int) -> torch.Tensor: - """Replicate the last latent frame ``n_frames`` times for NATTEN last-frame border.""" if n_frames <= 0: return latent padded, _ = resize_axis(latent, 2, latent.shape[2] + n_frames, mode="repeat_last") @@ -2407,7 +1825,6 @@ def crop_trailing_context_natten_pad( time_scale: int, stage5_kernel_t: int, ) -> torch.Tensor: - """Crop ghosting appendix before stage 5, leaving at least ``stage5_kernel_t``.""" if n_latent_frames <= 0: return context ghost = n_latent_frames * time_scale @@ -2418,7 +1835,6 @@ def crop_trailing_context_natten_pad( def _weight_floor(dtype: torch.dtype) -> float: - """Smallest divisor that safely guards ``buffer / weights`` in ``dtype``.""" return max(1e-8, torch.finfo(dtype).tiny) @@ -2430,7 +1846,6 @@ def stage4_thw_from_latent( *, drop_leading_frame: bool = True, ) -> Tuple[int, int, int]: - """Stage-4 input ``(T, H, W)`` after the first three upsample hops.""" t, h, w = latent_t, latent_h, latent_w for st, sh, sw in upsample_strides[:3]: t, h, w = t * st, h * sh, w * sw @@ -2443,7 +1858,6 @@ def stage4_to_pixel_scale_factors( upsample_stride: Tuple[int, int, int], patch_size: int, ) -> SpatioTemporalScaleFactors: - """Pixel/frame units per stage-4-input cell (last NA hop + unpatchify).""" st, sh, sw = upsample_stride return SpatioTemporalScaleFactors(time=st, height=sh * patch_size, width=sw * patch_size) @@ -2453,7 +1867,6 @@ def compute_tile_min_size( stage5_kernel: Tuple[int, int, int], upsample3_stride: Tuple[int, int, int], ) -> Tuple[int, int, int]: - """Min stage-4-input ``(T, H, W)`` so stages 4 and 5 each see ``>= kernel``.""" return tuple(max(stage4_kernel[a], -(-stage5_kernel[a] // upsample3_stride[a])) for a in range(3)) @@ -2464,7 +1877,6 @@ def compute_tile_halos( stage5_depth: int, upsample3_stride: Tuple[int, int, int], ) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]: - """One-sided halos in stage-4-input units for stages 4 and 5.""" halo4 = tuple(stage4_depth * (stage4_kernel[a] // 2) for a in range(3)) halo5 = tuple(-(-(stage5_depth * (stage5_kernel[a] // 2)) // upsample3_stride[a]) for a in range(3)) return halo4, halo5 # type: ignore[return-value] @@ -2473,7 +1885,6 @@ def compute_tile_halos( def _cumulative_upsample_strides( upsamples: Sequence[Tuple[Tuple[int, int, int], int]], ) -> List[Tuple[int, int, int]]: - """Per-axis product of hop strides for ``upsamples[:i]`` (``cumulative[0] = (1,1,1)``).""" cumulative = [(1, 1, 1)] t, h, w = 1, 1, 1 for stride, _ in upsamples: @@ -2487,7 +1898,6 @@ def all_stages_min_tile_size( upsamples: Sequence[Tuple[Tuple[int, int, int], int]], stage5_kernel: Tuple[int, int, int], ) -> Tuple[int, int, int]: - """Per-axis latent-grid floor so every stage's NA sees dims ``>= kernel_size``.""" cumulative = _cumulative_upsample_strides(upsamples) mins = [1, 1, 1] for stage_i in range(len(upsamples)): @@ -2516,11 +1926,6 @@ def recommended_pixel_overlaps( tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], pixel_scale: SpatioTemporalScaleFactors, ) -> Tuple[int, int]: - """Stage-4/5-safe ``(temporal_overlap_frames, spatial_overlap_pixels)``. - Shared by :func:`recommended_decode_tiling_config` (to *set* overlaps) and - :func:`~ltx_core.tiling._validate_overlap` (to reject undersized configs). - """ - def dominant(axis: int) -> int: return max(tile_halos[i][axis] for i in range(len(tile_halos))) @@ -2537,14 +1942,12 @@ def stage5_tokens_for_pixel_tile( *, patch_size: int, ) -> int: - """Pre-unpatchify stage-5 token count for a pixel-space tile (NATTEN volume).""" h5 = max(1, tile_height // patch_size) w5 = max(1, tile_width // patch_size) return tile_frames * h5 * w5 def _axis_candidates(length: int, overlap: int, min_size: int, multiple: int) -> list[tuple[int, int]]: - """``(tile_size, num_tiles)`` for every legal size on ``multiple``'s grid.""" out: list[tuple[int, int]] = [] max_size = max(_round_up(length, multiple), min_size) for size in range(min_size, max_size + multiple, multiple): @@ -2567,7 +1970,6 @@ def volumetric_overlap_waste( n_h: int, n_w: int, ) -> float: - """``processed_volume / unique_volume`` (>= 1). Lower means less overlap recompute.""" processed = n_t * n_h * n_w * tile_frames * tile_height * tile_width unique = max(1, num_frames * height * width) return processed / unique @@ -2578,16 +1980,6 @@ def _propagate_interval_through_upsample_hops( strides: Sequence[int], causal: bool, ) -> DimensionInterval: - """Forward-propagate one interval through a sequence of upsample hops on one axis. - Mirrors :class:`~ltx_core.model.video_vae.transformer.layers.LinearPixelShuffleUpsample`: - multiply by ``stride``, and for the causal temporal axis when ``stride == 2`` apply - the duplicate-frame drop (``end -= 1``; non-origin also ``start -= 1``). - This is *not* :func:`~ltx_core.model.video_vae.video_vae.map_temporal_slice` (ConvVAE). - DiffVAE non-origin tiles run with ``drop_leading_frame=False`` and must keep length - ``tile_t * stride``; the ConvVAE ``1+(L-1)*stride`` mapping is one frame short and - shifts non-origin ``out_coords``, which breaks tiled↔untiled temporal blend even - when masks are complementary. - """ x = interval for stride in strides: if stride < 1: @@ -2608,8 +2000,6 @@ def _propagate_interval_through_upsample_hops( class ChannelLinear(nn.Linear): - """``nn.Linear`` exposing ``in_channels``/``out_channels`` for config introspection.""" - @property def in_channels(self) -> int: return self.in_features @@ -2620,8 +2010,6 @@ def out_channels(self) -> int: class LinearPixelShuffleUpsample(nn.Module): - """Decoder-side resampler: Linear channel-expand, then channels-last PixelShuffle.""" - def __init__( self, in_channels: int, @@ -2635,15 +2023,6 @@ def __init__( self.proj = nn.Linear(in_channels, self.proj_out_channels, bias=True) def forward(self, x: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: - """Upsample; when ``stride[0] == 2`` the pixel-shuffle produces a duplicate - leading frame that must be dropped to preserve the causal 1:2 (then - composed 1:8) frame mapping. ``drop_leading_frame`` gates that drop: it - must be ``True`` only for the chunk that contains the tensor's true - temporal origin (t=0). Tiled callers processing a later chunk in - isolation must pass ``False`` -- that chunk has no duplicate leading - frame of its own to drop, since the one duplicate frame in the full - (untiled) tensor belongs solely to the origin chunk. - """ x = self.proj(x) x = rearrange( x, @@ -2658,11 +2037,6 @@ def forward(self, x: torch.Tensor, drop_leading_frame: bool = True) -> torch.Ten class AdaLNZero(nn.Module): - """Per-block AdaLN-Zero modulation: ``t_emb`` -> 7 (scale/shift/gate) chunks. - Zero-init output projection so the block is an identity at every timestep - until the modulation pathway opens up during training. - """ - NUM_CHUNKS: int = 7 # scale_msa, shift_msa, gate_msa, scale_mlp, shift_mlp, gate_mlp, gate_ctx def __init__(self, dim: int, t_emb_dim: int) -> None: @@ -2679,7 +2053,6 @@ def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: def modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: - """Apply AdaLN-style scale + shift modulation to a channels-last tensor.""" return x * (1.0 + scale) + shift @@ -2698,7 +2071,6 @@ def h_positions(h: int, device: torch.device) -> torch.Tensor: def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: - """Default split of head_dim across (T, H, W) RoPE chunks.""" assert head_dim % 8 == 0, f"head_dim={head_dim} must be a multiple of 8 for default split" d_t = (head_dim // 4) // 2 * 2 d_hw = (head_dim - d_t) // 2 @@ -2711,7 +2083,6 @@ def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: def rope_inv_freqs(dim: int, base: float = 10000.0) -> torch.Tensor: - """Inverse RoPE frequencies: ``1 / base**(i/dim)`` for ``i`` in ``[0, dim, 2)``.""" assert dim % 2 == 0, f"RoPE dim must be even, got {dim}" exponents = np.arange(0, dim, 2, dtype=np.float64) / dim inv_freqs = 1.0 / np.power(float(base), exponents) @@ -2726,7 +2097,6 @@ def rot_abs_axis_impl( *, compute_dtype: torch.dtype, ) -> torch.Tensor: - """Absolute RoPE on one axis chunk ``xc[..., D]`` (D even) → new tensor.""" out_dtype = xc.dtype pairs = xc.reshape(*xc.shape[:-1], xc.shape[-1] // 2, 2) xe = pairs[..., 0].to(compute_dtype) @@ -2767,10 +2137,6 @@ def _apply_opaque_rope_slab( compute_dtype: torch.dtype, t_pos: torch.Tensor | None = None, ) -> torch.Tensor: - """Rotate one W-extent with raw abs-RoPE (runs inside the opaque op). - ``t_pos`` overrides the default integer ``arange`` on the first axis; the keyframe - stream passes its (possibly fractional) plane times there. - """ d_t, d_h, _ = rope_split inv_t, inv_h, inv_w = inv_freqs t = x.shape[1] @@ -2797,11 +2163,6 @@ def _apply_opaque_tiled_rope( compute_dtype: torch.dtype, t_pos: torch.Tensor | None = None, ) -> torch.Tensor: - """Fixed-``num_tiles`` W split + per-slab rotation (body of the opaque op). - The keyframe stream shares the video stream's W extent at every stage, so passing the - same ``num_tiles`` yields identical slab boundaries and identical ``w_pos`` -- which - is what keeps the two streams' W phases comparable inside the joint softmax. - """ slabs = torch.chunk(x, num_tiles, dim=3) w_off = 0 parts: list[torch.Tensor] = [] @@ -2834,7 +2195,6 @@ def _abs_rope_op( num_tiles: int, compute_dtype_is_bf16: bool, ) -> torch.Tensor: - """Opaque out-of-place abs-RoPE: Dynamo sees one node.""" compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 return _apply_opaque_tiled_rope( x, @@ -2885,7 +2245,6 @@ def _abs_rope_at_t_op( num_tiles: int, compute_dtype_is_bf16: bool, ) -> torch.Tensor: - """Opaque abs-RoPE with caller-supplied (possibly fractional) first-axis positions.""" compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 return _apply_opaque_tiled_rope( x, @@ -2898,7 +2257,6 @@ def _abs_rope_at_t_op( def _rope_config(attn: object, x: torch.Tensor) -> tuple[tuple[torch.Tensor, ...], int, torch.dtype]: - """``(inv_freqs, num_tiles, compute_dtype)`` read off an attention module.""" inv_freqs = ( attn.rope_inv_t.to(device=x.device), # type: ignore[attr-defined] attn.rope_inv_h.to(device=x.device), # type: ignore[attr-defined] @@ -2910,7 +2268,6 @@ def _rope_config(attn: object, x: torch.Tensor) -> tuple[tuple[torch.Tensor, ... def _det_project_qkv(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Shared Q/K/V proj + norm + scale for :func:`det_qkv_rope` and ``_at_times``.""" q, k, v = attn.project_qkv(x) # type: ignore[attr-defined] q = attn.q_norm(q) # type: ignore[attr-defined] k = attn.k_norm(k) # type: ignore[attr-defined] @@ -2923,13 +2280,6 @@ def det_qkv_rope_at_times( x: torch.Tensor, t_pos: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Q/K/V proj + norm/scale + abs-RoPE with explicit first-axis positions. - The keyframe-stream counterpart of :func:`det_qkv_rope`. ``t_pos`` is ``(P,)`` float - stage times *in the same origin* the video stream's RoPE uses, so the joint softmax - sees true relative offsets: absolute-vs-local RoPE only cancels as a global phase - when every token in the softmax shares one origin, which no longer holds once - keyframe tokens join a video window. - """ if t_pos.ndim != 1 or t_pos.shape[0] != x.shape[1]: raise ValueError(f"t_pos must be ({x.shape[1]},) to match the plane axis, got {tuple(t_pos.shape)}") q, k, v = _det_project_qkv(attn, x) @@ -2954,7 +2304,6 @@ def det_qkv_rope_at_times( def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Q/K/V proj + norm/scale + opaque full-volume abs-RoPE.""" q, k, v = _det_project_qkv(attn, x) inv_freqs, num_tiles, compute_dtype = _rope_config(attn, x) q = _apply_opaque_abs_rope( @@ -2974,7 +2323,6 @@ def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Ten return q, k, v - def vram_ready_linear(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: # This decoder calls several projections functionally, bypassing the VRAM wrappers' # forward, so ask the wrapper for computation-ready weights instead of reading them raw. @@ -2985,8 +2333,6 @@ def vram_ready_linear(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor | N class QKVProjections(nn.Module): - """Checkpoint-fused QKV weights executed as the target's three projections.""" - def __init__(self, dim: int) -> None: super().__init__() linear = nn.Linear(dim, dim * 3, bias=True) @@ -3080,7 +2426,6 @@ def forward(self, x): def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], list[int]]: - """Per-index (start, end) of the attended window along one axis.""" starts: list[int] = [] ends: list[int] = [] if causal: @@ -3099,7 +2444,6 @@ def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], l def _pick_tiles(dims: tuple[int, int, int], kernels: list[int]) -> list[int]: - """Per-axis query-tile lengths keeping one tile's [Nq, Nk] under budget.""" tiles = list(dims) def cost(ts: list[int]) -> int: @@ -3120,7 +2464,6 @@ def _group_mask( dtype: torch.dtype, device: torch.device, ) -> torch.Tensor: - """Additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group.""" bools = [] for starts, ends in rel_bounds: st = torch.tensor(starts, device=device) @@ -3147,9 +2490,6 @@ def na3d( is_causal: list[bool] | None = None, scale: float | None = None, ) -> torch.Tensor: - """3D neighborhood attention over ``(B, T, H, W, NH, HD)`` tensors. - ``scale`` defaults to ``head_dim**-0.5``. Pass ``scale=1.0`` when Q is already scaled. - """ batch, t, h, w, nh, hd = q.shape dims = (t, h, w) causal = [False, False, False] if is_causal is None else list(is_causal) @@ -3294,26 +2634,18 @@ def na3d( def sdpa_materializes_scores(device: torch.device) -> bool: - """Whether torch's SDPA will materialize this backend's score block on ``device``. - CUDA takes the memory-efficient (or cuDNN) kernel for this backend's broadcast mask and - aligned head dim. Everywhere else the math path runs and the ``(G, NH, Nq, Nk)`` scores land - in memory -- notably MPS. Auto tiling reads this to size its reserve. - """ return device.type != "cuda" def staging_factor(device: torch.device) -> float: - """The peak-to-staging multiplier for this host, from the table above.""" return _STAGING_FACTOR_MATERIALIZED if sdpa_materializes_scores(device) else _STAGING_FACTOR_FUSED def _key_channels(head_dim: int) -> int: - """Key/query channel count: ``head_dim``, the bias channel, and alignment padding.""" return -(-(head_dim + 1) // _HEAD_DIM_ALIGN) * _HEAD_DIM_ALIGN def _window(kernel: int) -> tuple[int, int]: - """``(lo, hi)`` halo for one axis: the offsets ``range(-k // 2, k - k // 2)`` reach.""" lo = kernel // 2 return lo, kernel - lo - 1 @@ -3325,17 +2657,11 @@ def pick_brick( target: int = DEFAULT_BRICK_QUERIES, depth: int = DEFAULT_BRICK_DEPTH, ) -> tuple[int, int, int]: - """``(bt, bh, bw)``: ``depth`` frames deep with the squarest ~``target``-query face. - Square minimizes the key slab ``(bh + Kh - 1) * (bw + Kw - 1)`` at a fixed query count, which - is exactly the wasted-work term. - """ side = max(1, round(math.sqrt(target))) return min(depth, time), min(side, height), min(side, width) class _Geometry: - """Brick decomposition of a volume, plus the padding and slab extents it implies.""" - def __init__( self, height: int, @@ -3363,16 +2689,10 @@ def __init__( self.padded_width = width + sum(self.pad_w) def row_extent(self, rows: int) -> int: - """Padded ``H`` extent a group of ``rows`` brick rows needs from the staged volume.""" return (rows - 1) * self.brick[1] + self.span[0] class _Schedule: - """How the nested loops are cut so transient memory stays inside the budget. - ``group_axis`` counts *bricks* along the volume's leading axis (frames for the video pass, - planes for the keyframe pass); ``stage_axis`` counts them per staging pass. - """ - def __init__( self, geometry: _Geometry, @@ -3404,18 +2724,12 @@ def __init__( def _banded(queries: int, span: int, kernel: int, device: torch.device) -> torch.Tensor: - """``(queries, span)`` bool: key ``i`` is visible to query ``j`` iff ``j <= i < j + kernel``.""" key = torch.arange(span, device=device)[None, :] query = torch.arange(queries, device=device)[:, None] return (key >= query) & (key < query + kernel) def _joint_mask(geometry: _Geometry, num_slots: int, device: torch.device) -> torch.Tensor: - """``(1, 1, Nq, Nk)`` visibility, shared by every brick. - Query order is ``(jt, jh, jw)``; key order is the video slab ``(it, p, r)`` followed by the - keyframe slab ``(slot, p, r)``. Keyframe keys carry no temporal condition -- a plane is visible - to every frame in the brick, which is what makes the run grouping legal. - """ brick_t, brick_h, brick_w = geometry.brick kernel_t, kernel_h, kernel_w = geometry.kernel spatial = ( @@ -3441,10 +2755,6 @@ def _stage( *, with_bias_channel: bool, ) -> torch.Tensor: - """``(B, A, H, W, NH, HD)`` -> padded head-major ``(B, NH, A + pad, Hp, Wp, C)``. - Head-major so a brick slab's innermost ``(ew, C)`` block is contiguous in both source and - destination; channels-last staging makes the same gather markedly slower. - """ batch, axis, height, width, heads, head_dim = x.shape channels = _key_channels(head_dim) if with_bias_channel else head_dim out = x.new_zeros((batch, heads, axis + sum(pad_t), geometry.padded_height, geometry.padded_width, channels)) @@ -3472,12 +2782,6 @@ def _slabs( *, group_stride: int, ) -> torch.Tensor: - """Overlapping brick slabs as a *view*: ``(B, bricks, rows, Gw, NH, blocks, eh, ew, C)``. - ``staged`` is head-major ``(B, NH, A, Hp, Wp, C)``, already sliced to this group's first brick - and brick row, so the view inherits its storage offset. ``group_stride`` is how far consecutive - bricks advance along ``A``: the brick depth for the sliding video window, and **zero** for the - keyframe planes, which every brick in a run shares. - """ batch, heads = staged.shape[0], staged.shape[1] stride_b, stride_nh, stride_a, stride_h, stride_w, _ = staged.stride() return staged.as_strided( @@ -3497,11 +2801,6 @@ def _slabs( def _query_bricks(x: torch.Tensor, geometry: _Geometry, bricks: int, rows: int) -> torch.Tensor: - """``(B, A, h, W, NH, HD)`` -> ``(B * bricks * rows * Gw, NH, Nq, C)``, unit channel set. - ``x`` is this group's slice, so ``A`` may be short of ``bricks * bt`` and ``h`` short of - ``rows * bh`` at a volume edge; the shortfall is zero-padded here and cropped by - :func:`_unbrick`. - """ batch, axis, height, width, heads, head_dim = x.shape brick_t, brick_h, brick_w = geometry.brick pad_t, pad_h, pad_w = bricks * brick_t - axis, rows * brick_h - height, geometry.grid[1] * brick_w - width @@ -3526,7 +2825,6 @@ def _unbrick( rows: int, extent: tuple[int, int], ) -> torch.Tensor: - """Inverse of :func:`_query_bricks`, cropping to ``extent`` frames/rows and the real width.""" brick_t, brick_h, brick_w = geometry.brick heads, head_dim = attended.shape[1], attended.shape[3] plane = ( @@ -3538,12 +2836,10 @@ def _unbrick( def _with_null(slots: torch.Tensor, null_index: int) -> torch.Tensor: - """Map empty slots (``-1``) onto the appended null row, which biases itself out.""" return torch.where(slots < 0, torch.full_like(slots, null_index), slots) def _append_null(keys: torch.Tensor, values: torch.Tensor, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]: - """Append one all-dead key plane (and a zero value plane) along the staged plane axis.""" shape = (keys.shape[0], keys.shape[1], 1, *keys.shape[3:]) null_key = keys.new_zeros(shape) null_key[..., head_dim] = _DEAD @@ -3552,12 +2848,6 @@ def _append_null(keys: torch.Tensor, values: torch.Tensor, head_dim: int) -> tup def _slot_runs(slots: torch.Tensor) -> list[tuple[int, int]]: - """Maximal ``[start, stop)`` runs of leading-axis positions whose slot row is identical. - A brick spanning several frames shares one keyframe key slab, so it may not straddle a change - of slot row. At production keyframe spacing these runs are ~16 frames long, so the constraint - costs little; carrying every frame's slots in the slab instead would more than give back what - brick depth wins. - """ rows = slots.tolist() runs: list[tuple[int, int]] = [] start = 0 @@ -3577,10 +2867,6 @@ def _attend_group( shape: tuple[int, int], mask: torch.Tensor, ) -> torch.Tensor: - """Gather one ``(bricks x brick rows)`` block's keys, attend, and un-brick the result. - ``key_views`` / ``value_views`` are the strided slab views to concatenate along the key axis, - in order. ``shape`` is ``(bricks, rows)``. - """ bricks, rows = shape batch = query_slice.shape[0] heads, head_dim = query_slice.shape[4], query_slice.shape[5] @@ -3606,7 +2892,6 @@ def _attend_group( def _row_groups(geometry: _Geometry, schedule: _Schedule) -> list[tuple[int, int, slice, slice]]: - """``(row, rows, staged H slice, output H slice)`` per brick-row group.""" brick_h = geometry.brick[1] groups = [] for row in range(0, geometry.grid[0], schedule.group_rows): @@ -3633,7 +2918,6 @@ def _video_query_pass( workspace_bytes: int, factor: float, ) -> torch.Tensor: - """Video queries: the local ``Kt x Kh x Kw`` window plus the nearest keyframe planes.""" time, heads, head_dim = q.shape[1], q.shape[4], q.shape[5] brick_t = geometry.brick[0] lo_t, hi_t = geometry.pad_t @@ -3724,10 +3008,6 @@ def _keyframe_query_pass( workspace_bytes: int, factor: float, ) -> torch.Tensor: - """Keyframe queries: own plane only (``d_t == 0``) plus the nearest video frames. - Runs one plane per brick: planes have no temporal window, so depth would buy nothing, and each - plane's video slots differ anyway. - """ planes_total, heads, head_dim = keyframe_q.shape[1], keyframe_q.shape[4], keyframe_q.shape[5] num_slots = slots.shape[1] blocks = 1 + num_slots @@ -3801,21 +3081,6 @@ def joint_na3d( # noqa: PLR0913 workspace_bytes: int = DEFAULT_WORKSPACE_BYTES, factor: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Joint neighborhood attention over a video volume and a keyframe plane stack. - Args: - q, k, v: ``(B, T, H, W, NH, HD)`` video stream. ``q`` must arrive pre-scaled by - ``head_dim ** -0.5``, as the shared attention module does. - keyframe_q, keyframe_k, keyframe_v: ``(B, P, H, W, NH, HD)`` keyframe stream. - keyframe_times: ``(P,)`` float32 plane times, same origin as the video RoPE. - keyframe_valid: ``(P,)`` bool. - kernel_size: ``(Kt, Kh, Kw)``. - num_slots: cross-stream slots per query. - brick: query brick ``(bt, bh, bw)``; defaults to :func:`pick_brick`. - workspace_bytes: transient budget bounding the staged window and key/value block. - factor: peak-to-staging multiplier; :func:`staging_factor` supplies it when omitted. - Returns: - ``(video_out, keyframe_out)``, each shaped like its stream's ``q``. - """ time, height, width = q.shape[1], q.shape[2], q.shape[3] video_slots = video_keyframe_slots(keyframe_times, keyframe_valid, time, num_slots) keyframe_slots = keyframe_video_slots(keyframe_times, keyframe_valid, time, num_slots) @@ -3839,7 +3104,6 @@ def joint_na3d( # noqa: PLR0913 ) - class EagerNAAttention: def __call__(self, attn, q, k, v): return na3d(q, k, v, kernel_size=attn.kernel_size, scale=1.0) @@ -3853,7 +3117,6 @@ def __call__(self, attn, q, k, v, keyframe_q, keyframe_k, keyframe_v, keyframe_t ) - """3D Neighborhood Attention via NATTEN + absolute RoPE prelude. Parameter shell shared by det ``NABlock`` and both diff-attn roles. Diffusion AdaLN residuals live in pathway packages (each owns its RoPE); @@ -3876,13 +3139,6 @@ def __call__(self, attn, q, k, v, keyframe_q, keyframe_k, keyframe_v, keyframe_t class NAAttentionCallable(Protocol): - """A windowed 3D neighborhood-attention backend. - Q/K/V arrive as ``(B, T, H, W, NH, HD)``, already normed, scaled and RoPE'd; - the return is ``(B, T, H, W, NH*HD)`` or anything reshapeable to it. The owning - module is passed so a backend can read configuration (``kernel_size``, softmax - bound, …). Backend-specific settings (NATTEN's kernel pin) live on the callable. - """ - def __call__( self, attn: NeighborhoodAttention3D, @@ -3893,17 +3149,6 @@ def __call__( class JointNAAttentionCallable(Protocol): - """A windowed 3D NA backend that also carries a keyframe plane stack. - Same conventions as :class:`NAAttentionCallable` -- Q/K/V already normed, scaled and - absolutely RoPE'd -- with a second stream shaped ``(B, P, H, W, NH, HD)`` whose plane - axis sits in video's temporal slot. Both streams' RoPE must share one origin, so - ``keyframe_times`` are tile-local. Returns one output per stream. - NATTEN and the CuTe DSL kernel cannot express a joint window, so this is a separate - slot from ``attention_function`` rather than a widening of it: it keeps the shipping - keyframe-less hot path untouched, and it is immune to the install-order hazard that - ``configure_natten_backend`` creates by overwriting ``attention_function`` wholesale. - """ - def __call__( self, attn: NeighborhoodAttention3D, @@ -3919,16 +3164,6 @@ def __call__( class NeighborhoodAttention3D(nn.Module): - """3D Neighborhood Attention with absolute RoPE + pluggable NA backend. - Q/K receive absolute RoPE; attention is ``attention_function`` (NATTEN by - default; Triton or eager SDPA when natten is missing; CuTe DSL via - DiffVAE BLACKWELL_DSL install). Relative gather-based NA is not used as a - production gather path on this branch. - NATTEN shifts its window inward at grid boundaries instead of - clamp-and-mask; interior positions match the gather reference closely, - boundary positions may differ slightly. - """ - def __init__( self, dim: int, @@ -3972,18 +3207,12 @@ def __init__( self.w_chunks = 1 # 1 = no chunking def project_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Q/K/V as owned contiguous ``(B,T,H,W,NH,HD)`` tensors.""" batch, t, h, w, _ = x.shape q, k, v = self.qkv(x) shape = (batch, t, h, w, self.num_heads, self.head_dim) return q.view(shape), k.view(shape), v.view(shape) def forward(self, x: torch.Tensor) -> torch.Tensor: - """Det-stage NA: opaque abs-RoPE via ``det_attn_rope`` + ``attention_function``. - ``x``/output: (B, T, H, W, C) — channels-last. RoPE positions are local - 0-based (see ``det_attn_rope`` module docstring for why that is - equivalent under tiled decode). - """ batch, t, h, w, _ = x.shape kt, kh, kw = self.kernel_size if t < kt or h < kh or w < kw: @@ -4004,10 +3233,6 @@ def forward_with_keyframes( x: torch.Tensor, keyframes: KeyframeStream, ) -> tuple[torch.Tensor, KeyframeStream]: - """Dual-stream det NA: one joint softmax over video and keyframe planes. - Unlike :meth:`forward` there is no ``dims >= kernel_size`` floor: the joint window - is clamp-and-mask, so an undersized axis simply masks its out-of-range offsets. - """ if self.joint_attention_function is None: raise RuntimeError( "keyframe decode needs joint_attention_function installed; build the decoder " @@ -4057,8 +3282,6 @@ def forward_with_keyframes( class NABlock(nn.Module): - """Pre-norm transformer block: NA -> SwiGLU MLP with residual adds.""" - def __init__( self, dim: int, @@ -4075,7 +3298,6 @@ def __init__( self.mlp = SwiGLU(dim, hidden) def forward(self, x: torch.Tensor) -> torch.Tensor: - """Channels-last in/out: (B, T, H, W, C).""" x = x + self.attn(self.norm1(x)) x = plain_mlp(x, self.mlp, self.norm2, self.mlp.tile) return x @@ -4085,13 +3307,6 @@ def forward_with_keyframes( x: torch.Tensor, keyframes: KeyframeStream, ) -> tuple[torch.Tensor, KeyframeStream]: - """Dual-stream block: video ``(B,T,H,W,C)`` and keyframe planes ``(B,P,H,W,C)``. - Every weight is shared with :meth:`forward`; the streams meet only inside the - joint attention softmax. Invalid planes are deliberately *not* re-zeroed here -- - the decoder re-zeroes after each upsample instead, matching upstream, so a masked - plane's hidden state may drift within a stage. It is masked out of every softmax - regardless, so this is cosmetic; reproducing it keeps us comparable. - """ attn_out, keyframe_attn = self.attn.forward_with_keyframes( self.norm1(x), dataclasses.replace(keyframes, x=self.norm1(keyframes.x)), @@ -4104,14 +3319,6 @@ def forward_with_keyframes( class DiffusionNABlock(nn.Module): - """Parameter shell for diffusion NA + SwiGLU with shared AdaLN-Zero. - Mode-specific subclasses (:class:`~ltx_core.model.video_vae.transformer.combined.block.CombinedDiffusionNABlock`, - :class:`~ltx_core.model.video_vae.transformer.chunked.block.ChunkedDiffusionNABlock`) - are installed via ModuleOps ``__class__`` swap and own the forward path. - Not a ``Protocol``: must be a concrete ``nn.Module`` so checkpoint load and - ``__class__`` swap keep one parameter identity. - """ - def __init__( self, dim: int, @@ -4151,10 +3358,6 @@ def combined( w_proj: torch.Tensor, b_proj: torch.Tensor | None, ) -> torch.Tensor: - """Split ``context_and_x`` via ``w_proj.shape[1]`` (context channels), add projected ctx. - Returns the updated ``x`` half (not the full concatenated buffer). - ``w_proj`` is ``context_proj.weight`` with shape ``(dim, context_channels)``. - """ context_channels = w_proj.shape[1] latent_context = context_and_x[..., :context_channels] x = context_and_x[..., context_channels:] @@ -4182,10 +3385,6 @@ def _apply_nested_abs_rope_slab( compute_dtype: torch.dtype, t_pos: torch.Tensor | None = None, ) -> torch.Tensor: - """Rotate one W-extent with nested per-axis abs-RoPE. - ``t_pos`` overrides the integer ``arange`` on the first axis; the keyframe stream passes - its fractional plane times there so both streams' RoPE shares one origin. - """ d_t, d_h, _ = rope_split inv_t, inv_h, inv_w = inv_freqs t = x.shape[1] @@ -4212,11 +3411,6 @@ def _apply_nested_full_volume_rope( compute_dtype: torch.dtype, t_pos: torch.Tensor | None = None, ) -> torch.Tensor: - """Fixed-``num_tiles`` W split + nested per-slab rotation (Dynamo-safe). - Both streams share the same W extent at stage 5, so the same ``num_tiles`` gives - identical slab boundaries and ``w_pos`` -- required for their W phases to be comparable - inside the joint softmax. - """ slabs = torch.chunk(x, num_tiles, dim=3) w_off = 0 parts: list[torch.Tensor] = [] @@ -4242,10 +3436,6 @@ def _qkv_nested_rope( x: torch.Tensor, t_pos: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Q/K/V proj + norm/scale + nested abs-RoPE. - ``t_pos`` overrides the integer first-axis positions; the keyframe stream passes - its (possibly fractional) plane times there. - """ q, k, v = attn.project_qkv(x) q = attn.q_norm(q) * attn.scale k = attn.k_norm(k) @@ -4284,15 +3474,6 @@ def full_with_keyframes( keyframe_times: torch.Tensor, keyframe_valid: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Dual-stream AdaLN residual attention: one joint softmax, shared weights. - Plain tensors rather than a ``KeyframeStream`` here: this is the stage-5 hot path, and - rebuilding a dataclass per block inside a compiled region buys nothing. - Both streams take the *same* ``scale``/``shift``. Upstream computes a separate keyframe - modulation, but the two are identical unless per-frame timestep conditioning supplies a - conditioning mask, which we deliberately do not implement yet. - No ``dims >= kernel_size`` floor, unlike :func:`full`: the joint window is - clamp-and-mask, so undersized axes simply mask their out-of-range offsets. - """ if attn.joint_attention_function is None: raise RuntimeError( "keyframe decode needs joint_attention_function installed; build the decoder " @@ -4334,7 +3515,6 @@ def full( scale: torch.Tensor, shift: torch.Tensor, ) -> torch.Tensor: - """``x + NA(modulate(norm(x)))`` with nested full-volume abs-RoPE.""" y = norm(x) * (1.0 + scale) + shift batch, t, h, w, _ = y.shape kt, kh, kw = attn.kernel_size @@ -4366,7 +3546,6 @@ def residual_mlp( shift: torch.Tensor, tile: SwiGLUTileSpec, ) -> torch.Tensor: - """Combined*: ``x + swiglu_tiled(modulate(norm(x), scale, shift))``.""" y = modulate(norm(x), scale, shift) if y.numel() == 0: return x @@ -4377,8 +3556,6 @@ def residual_mlp( class CombinedDiffusionNABlock(DiffusionNABlock): - """Combined-context diffusion block: ``forward`` / ``forward_combined``.""" - def forward_combined_with_keyframes( self, context_and_x: torch.Tensor, @@ -4387,13 +3564,6 @@ def forward_combined_with_keyframes( keyframe_times: torch.Tensor, keyframe_valid: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Dual-stream block; returns the two updated x halves (not the concat buffers). - Each stream gets its own ``context_proj(context)`` injection from its own - ``[context | x]`` buffer, then both meet in one joint attention softmax, then each - runs the shared MLP. Invalid keyframe planes are re-zeroed on the way out -- unlike - the deterministic ``NABlock``, which leaves that to the decoder's post-upsample - masking. Both asymmetries are upstream's. - """ scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) w_proj, b_proj = vram_ready_linear(self.context_proj) x = inject_context(context_and_x, w_proj, b_proj) @@ -4471,22 +3641,6 @@ def forward( class DiffusionVideoDecoder(nn.Module, Disposable, VideoDecoder): - """Diffusion-based video VAE decoder (Neighborhood-Attention backbone). - Minimal port of the reference ``NADiffusionDecoder``. - Stages 1-4 deterministically upsample the latent into a context volume - (same NA-upsample path as the non-diffusion NA decoder). Stage 5 runs - ``DiffusionNABlock``s that denoise the patchified noised pixels ``x_t``, - guided by that context via AdaLN-Zero scale/shift (ungated residuals; - legacy static gates are folded into Linear weights at load time). - Last-frame NATTEN window-shift is mitigated by temporarily replicating the - last latent frame ``(stage1_K_t // 2) * 2`` times through stages 1-4, then - cropping that appendix from context before stage 5 - but only down to - ``max(original_context_T, stage5_kernel[0])`` so undersized clips (e.g. a - single latent frame) still satisfy NATTEN's kernel floor. Latents / tiles - below ``stage_min_tile_sizes`` are edge-padded first via ``diffusion_tiling``; - leftover pad is cropped from the final pixels. - """ - def __init__( # noqa: PLR0913 self, in_channels: int = 128, @@ -4620,7 +3774,6 @@ def __init__( # noqa: PLR0913 self._keyframe_time_strides: Tuple[int, ...] = remaining_time_strides(self.upsamples) def _run_det_stage(self, x: torch.Tensor, stage_i: int, drop_leading_frame: bool) -> torch.Tensor: - """One deterministic stage: NA blocks + upsample.""" if self.mark_dynamic_shapes: for dim in (1, 2, 3): torch._dynamo.mark_dynamic(x, dim) @@ -4633,11 +3786,6 @@ def forward_stages_1_to_3( z_noisy: torch.Tensor, drop_leading_frame: bool = True, ) -> torch.Tensor: - """Stages 1-3 on a full (or already ghost-padded) latent → stage-4 input feature. - Output is channels-last ``(B, T, H, W, C)`` at stage-4 input resolution. - Callers that want NATTEN trailing ghosting should pad the latent first via - ``pad_trailing_latent_for_natten_border``. - """ z_noisy = self.per_channel_statistics.un_normalize(z_noisy) x = z_noisy.permute(0, 2, 3, 4, 1) x = self.conv_in(x) @@ -4651,14 +3799,6 @@ def _keyframe_stream_from_latents( *, valid: torch.Tensor | None = None, ) -> KeyframeStream: - """Keyframe latents to a stage-1-input stream: un-normalize, tag, ``conv_in``, mask. - Stages 1-3 always run on the whole volume, so times stay in the global stage-1 - origin. Tile-local rebasing happens later, at stage 4. - Unlike upstream we un-normalize first, because our ``conv_in`` consumes - un-normalized latents (``forward_stages_1_to_3`` does the same for video) while - upstream un-normalizes above the decoder. ``type_emb`` still lands in exactly the - same place: on the latents, immediately before the shared ``conv_in``. - """ latents = self.per_channel_statistics.un_normalize(keyframes.latents) x = latents.permute(0, 2, 3, 4, 1) x = x + self.type_emb.to(dtype=x.dtype, device=x.device).view(1, 1, 1, 1, -1) @@ -4683,16 +3823,6 @@ def _run_det_stage_with_keyframes( next_time_origin: float, clip_start_frame: int = 0, ) -> tuple[torch.Tensor, KeyframeStream]: - """One deterministic stage over both streams: joint NA blocks + upsample. - The keyframe upsample is spatial-only and always drops its leading frame; the video - stream's ``drop_leading_frame`` is a tiling property and must not leak into it. - Times are rebuilt from the *next* stage's remaining stride after upsampling. - ``next_time_origin`` is in the **next** stage's temporal units, because that is the - scale the times it rebases are expressed in. Zero everywhere except the stage-4 hop of - a tiled decode, whose next stage is 5 and whose origin is therefore a pixel frame. - ``clip_start_frame`` is the first global pixel of *this* video latent (0 for a full - clip; Dist's tile origin for a slice). - """ if self.mark_dynamic_shapes: for dim in (1, 2, 3): torch._dynamo.mark_dynamic(x, dim) @@ -4720,16 +3850,6 @@ def forward_stages_1_to_3_with_keyframes( *, keyframe_valid: torch.Tensor | None = None, ) -> tuple[torch.Tensor, KeyframeStream]: - """Dual-stream stages 1-3: video latent + keyframe planes to stage-4 inputs. - Keyframe counterpart of :meth:`forward_stages_1_to_3`. ``z_noisy`` and - ``keyframes.latents`` must already carry identical spatial padding -- the pad is - applied symmetrically, so padding only one stream would offset every keyframe plane - from the video by half the pad and read as ghosting rather than a failure. - Times are relative to :attr:`DecodeKeyframes.clip_start_frame`. For a full-clip decode - that is 0, so they match global ``t_s``. For a Dist slice they are ``t_s(index) - - t_s(clip_start)`` from stage 1, because this path's "whole volume" *is* the slice. - Additional in-volume tile origins are applied in :meth:`forward_stage_4_with_keyframes`. - """ keyframes.validate() if z_noisy.shape[-2:] != keyframes.latents.shape[-2:]: raise ValueError( @@ -4763,25 +3883,6 @@ def forward_stage_4_with_keyframes( pixel_time_origin: float = 0.0, clip_start_frame: int = 0, ) -> tuple[torch.Tensor, KeyframeStream]: - """Dual-stream stage 4 to stage-5 context. Keyframe counterpart of - :meth:`forward_stage_4`. - The ghost-pad crop applies to the video stream only: the trailing replicate is a - temporal-border workaround and the keyframe planes have no temporal extent to pad. - On the deferred (chunked) pathway neither stream is upsampled here: each stage-5 - block folds ``upsamples[3]`` into its own context inject. The returned stream's - ``times`` are nonetheless the *stage-5* times, because that is where they are - consumed -- the same asymmetry the video stream already has, whose returned - ``x`` is a pre-upsample feature rather than stage-5 context. - **Two origins, at two scales.** The video stream's RoPE is tile-local 0-based at every - stage, so keyframe times must be rebased to whatever the tile's frame 0 is -- and this - method spans two different temporal resolutions. ``stage4_time_origin`` is the tile's - start in stage-4 input units (for the blocks); ``pixel_time_origin`` is its first - global pixel frame (for stage 5). They are taken separately from the tile rather than - derived from one another: ``drop_leading_frame`` and the causal first frame make - ``pixel_origin == stride_t * stage4_origin`` an off-by-one trap, not an identity. Both - are 0.0 for an untiled full-clip decode. ``clip_start_frame`` is subtracted in stage - units as well, so a Dist slice whose first pixel is 56 still sees ``t_s(48) - t_s(56)``. - """ # Rebuild from global indices rather than trusting the caller's stream: stages 1-3 of a # full-clip decode are global, and Dist has already folded clip_start into clip times. keyframes = dataclasses.replace( @@ -4824,10 +3925,6 @@ def _forward_stage_4_deferred_with_keyframes( pixel_time_origin: float, clip_start_frame: int = 0, ) -> tuple[torch.Tensor, KeyframeStream]: - """Stage-4 blocks only, both streams, for the deferred (chunked) pathway. - Mirrors :meth:`forward_stage_4`'s deferred branch: no ``upsamples[3]`` on either - stream, ghost cropped at pre-upsample temporal resolution (video only). - """ if self.mark_dynamic_shapes: for dim in (1, 2, 3): torch._dynamo.mark_dynamic(x, dim) @@ -4859,12 +3956,6 @@ def forward_stage_4( drop_leading_frame: bool = True, pad_trailing: bool = True, ) -> torch.Tensor: - """Stage 4 on a stage-4-input feature tile → stage-5 context (or pre-upsample feat). - ``x`` is channels-last. When ``pad_trailing``, soft-crop the ghosting - appendix before returning (ghost pad must already be present upstream). - When ``deferred_stage4_upsample`` is set, runs NA blocks only (no - ``upsamples[3]``) and crops ghost at pre-upsample temporal resolution. - """ if self.deferred_stage4_upsample: if self.mark_dynamic_shapes: for dim in (1, 2, 3): @@ -4892,7 +3983,6 @@ def forward_stage_4( return x def _context_and_x_for_diff_step(self, context: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor: - """Build block-ready ``[context | conv_in_x_t(patched x)]`` for ``forward_diff_step``.""" noised_pixels_patched = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1) x = self.conv_in_x_t(noised_pixels_patched.permute(0, 2, 3, 4, 1)) return torch.cat([context, x], dim=-1) @@ -4903,24 +3993,16 @@ def _keyframe_context_and_x_for_diff_step( keyframe_x_t: torch.Tensor, keyframe_valid: torch.Tensor, ) -> torch.Tensor: - """Keyframe ``[context | conv_in_x_t(patched x)]``, mask-zeroed. - ``keyframe_x_t`` is ``(B, C_pix, P, H_pix, W_pix)`` -- the keyframe planes' own noised - pixels, one pixel frame per plane, through the *shared* ``conv_in_x_t``. - """ patched = patchify(keyframe_x_t, patch_size_hw=self.patch_size, patch_size_t=1) x = self.conv_in_x_t(patched.permute(0, 2, 3, 4, 1)) x = x * keyframe_valid[None, :, None, None, None] return torch.cat([keyframe_context, x], dim=-1) def _x_for_diff_step(self, x_t: torch.Tensor) -> torch.Tensor: - """Conv-processed noised pixels only (deferred-context path).""" noised_pixels_patched = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1) return self.conv_in_x_t(noised_pixels_patched.permute(0, 2, 3, 4, 1)) def _keyframe_x_for_diff_step(self, keyframe_x_t: torch.Tensor, keyframe_valid: torch.Tensor) -> torch.Tensor: - """Mask-zeroed keyframe noised pixels only (deferred-context path). - Contiguous by construction: the chunked pathway mutates this buffer in place. - """ patched = patchify(keyframe_x_t, patch_size_hw=self.patch_size, patch_size_t=1) x = self.conv_in_x_t(patched.permute(0, 2, 3, 4, 1)) return (x * keyframe_valid[None, :, None, None, None]).contiguous() @@ -4930,13 +4012,6 @@ def forward_diff_step( context_and_x: torch.Tensor, t: torch.Tensor, ) -> torch.Tensor: - """One stage-5 diffusion step. Returns the model prediction in pixel space. - ``context_and_x`` is ``[latent_context | conv_in_x_t(x)]`` (channels-last), built - at the call site via ``_context_and_x_for_diff_step``. That single buffer is - reused across ``diff_blocks``: each block writes its output x-half back - with ``copy_`` (no per-block ``cat``). One-tensor layout keeps Dynamo - T/H/W symbols identical under ``mark_dynamic``. - """ x_half = context_and_x[..., self.context_channels :] t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x_half.dtype) modulation = self.shared_adaln(t_emb) @@ -4957,17 +4032,6 @@ def forward_diff_step_with_keyframes( keyframe_times: torch.Tensor, keyframe_valid: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """One dual-stream stage-5 step. Returns ``(video_pred, keyframe_pred)`` in pixel space. - The keyframe stream at stage 5 is a genuine second *pixel* diffusion stream -- its own - noised pixels through the shared ``conv_in_x_t``, its own per-block - ``context_proj(keyframe_context)``, the same AdaLN modulation -- not a zero tensor and - not the context. It is evolved through the same Euler loop as video so the hidden - state the joint attention reads sits at the noise level it was trained to see, then - discarded: callers use the video prediction only. - Both buffers follow ``forward_diff_step``'s ``[context | x]`` layout and the same - ``copy_``-into-a-view discipline. Video T/H/W stay dynamic; the keyframe plane axis is - specialized, since ``keyframe_times`` / ``keyframe_valid`` pin it. - """ x_half = context_and_x[..., self.context_channels :] keyframe_half = keyframe_context_and_x[..., self.context_channels :] t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x_half.dtype) @@ -4995,7 +4059,6 @@ def forward_diff_step_with_keyframes( return self._pixels_from_stage5(x_half), self._pixels_from_stage5(keyframe_half) def _pixels_from_stage5(self, x: torch.Tensor) -> torch.Tensor: - """Shared stage-5 tail: ``norm_out`` -> ``conv_out`` -> channels-first -> unpatchify.""" x = self.norm_out(x) x = self.conv_out(x) x = x.permute(0, 4, 1, 2, 3).contiguous() @@ -5009,13 +4072,6 @@ def forward_diff_step_deferred( *, drop_leading_frame: bool = True, ) -> torch.Tensor: - """Stage-5 step with deferred context: only ``x`` + low-res ``stage4_feat``. - Marks T/H/W dynamic on both tensors. CHUNKED blocks upsample then - ``context_proj`` on the host before attn+mlp; BLACKWELL_DSL - (``DSLDiffusionBlockChain``) folds that hop into the fused kernel and never - materialises full-resolution context. ``drop_leading_frame`` must match the - flag used for this tile's stage-4 path (origin tile vs non-origin). - """ t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype) modulation = self.shared_adaln(t_emb) @@ -5050,13 +4106,6 @@ def forward_diff_step_deferred_with_keyframes( *, drop_leading_frame: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - """Dual-stream stage-5 step with deferred context. Keyframe counterpart of - :meth:`forward_diff_step_deferred`. - Each stream carries its own pre-upsample stage-4 feature and injects it per block, - so full-resolution context is never materialised for either. ``drop_leading_frame`` - is the video stream's tiling property; the keyframe inject always collapses its - temporal stride. - """ t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype) modulation = self.shared_adaln(t_emb) @@ -5103,9 +4152,6 @@ def forward_diff_step_deferred_with_keyframes( def _euler_step( self, x_t: torch.Tensor, model_out: torch.Tensor, t_now: torch.Tensor, t_next: torch.Tensor ) -> torch.Tensor: - """One reverse-diffusion Euler update: advance ``x_t`` from ``t_now`` to - ``t_next`` given the model's prediction at ``t_now``. - """ compute_dtype = x_t.dtype dt = (t_now - t_next).view(-1, *([1] * (x_t.ndim - 1))).to(torch.float32) x_t_fp32 = x_t.to(torch.float32) @@ -5121,7 +4167,6 @@ def _decode_one_tile( timestep: torch.Tensor, pad_trailing: bool, ) -> torch.Tensor: - """Run stage 4 + diffusion on one stage-4 feature tile (isolation).""" context_tile = self.forward_stage_4( feat_tile, drop_leading_frame=is_origin, @@ -5160,14 +4205,6 @@ def _stage5_canvas_from_context( *, drop_leading_frame: bool, ) -> tuple[int, int, int]: - """``(T, H_pix, W_pix)`` of the stage-5 pixel canvas a context tile implies. - On the combined pathway the context is already at stage-5 resolution and - ``_context_and_x_for_diff_step`` patchifies ``x_t`` by ``patch_size``, so the canvas - is just ``(T, H * patch_size, W * patch_size)``. On the deferred pathway the tile is - still pre-upsample: each stage-5 block folds ``upsamples[3]`` into its inject, so - apply that stride here -- including the leading-frame drop the fold performs when the - temporal stride is 2. - """ t, h, w = context_tile.shape[1], context_tile.shape[2], context_tile.shape[3] if self.deferred_stage4_upsample: # Context is still pre-upsample, so this is the same geometry as a stage-4 @@ -5200,19 +4237,6 @@ def _decode_one_tile_with_keyframes( # noqa: PLR0913 pixel_time_origin: float = 0.0, clip_start_frame: int = 0, ) -> torch.Tensor: - """Stage 4 + dual-stream diffusion on one stage-4 feature tile. - Both streams are Euler-stepped together; only the video pixels are returned. The - keyframe pixel stream exists so the hidden state the joint attention reads stays at - the noise level it was trained on, and is discarded here (upstream exposes it only - through its explicit per-step entry points, which the trainer uses). - Noise is sized from the stage-5 context rather than from a re-derived tile geometry; - see :meth:`_stage5_canvas_from_context` for the deferred-pathway correction. The - keyframe canvas differs only in its frame count -- one pixel frame per plane, since - keyframe upsampling collapses its temporal stride. - ``x_t_tile_init`` lets a tiled decode share one global noise field across tiles (edge - policy applied by the caller, as the plain path does); ``None`` draws fresh noise. The - keyframe stream always draws its own -- its planes are not part of the video canvas. - """ context_tile, keyframes = self.forward_stage_4_with_keyframes( feat_tile, keyframes, @@ -5289,14 +4313,6 @@ def _decode_temporal_group_isolated_with_keyframes( # noqa: PLR0913 complementary: bool, clip_start_frame: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Decode one temporal group's tiles with keyframes and blend into a group buffer. - Keyframe counterpart of :meth:`_decode_temporal_group_isolated`. Each tile carries - only the planes near it -- those inside its pixel-frame span plus the nearest plane on - each *side* of it (:func:`planes_for_tile`) -- and the two origins that plane set has - to be rebased against come from the tile, never from each other. - Tile ``out_coords`` are local to this latent. Dist slices keep global - ``pixel_frame_indices`` and pass ``clip_start_frame`` so selection lines up. - """ group_temporal_len = curr_temporal_slice.stop - curr_temporal_slice.start group_shape = full_video_shape._replace(frames=group_temporal_len) full_torch_shape = full_video_shape.to_torch_shape() @@ -5418,13 +4434,6 @@ def _decode_groups_with_keyframes( # noqa: PLR0913, PLR0915 as_fhwc: bool, clip_start_frame: int = 0, ) -> Iterator[torch.Tensor]: - """Tiled keyframe decode, streaming one temporal group at a time. - Same shape as :meth:`_decode_pixels`: only the trailing overlap of the previous group - is retained between iterations, and a group's exclusive frames are yielded before the - next group decodes. That keeps residency at roughly two tile extents rather than a - whole video, which matters more here than on the plain path -- a keyframe decode also - carries a second pixel stream through stage 5. - """ full_video_shape = ( VideoLatentShape.from_torch_shape(latent.shape) .upscale(self.video_downscale_factors) @@ -5471,7 +4480,6 @@ def _decode_groups_with_keyframes( # noqa: PLR0913, PLR0915 overlap_stub_weights: torch.Tensor | None = None def _emit(buf: torch.Tensor, wts: torch.Tensor | None, global_start: int) -> torch.Tensor | None: - """Finalize, crop to content, and lay out one emitted run of frames.""" if global_start >= content_pixel.frames or buf.shape[2] < 1: return None frames_keep = min(buf.shape[2], content_pixel.frames - global_start) @@ -5558,11 +4566,6 @@ def _decode_pixels_with_keyframes( *, as_fhwc: bool = False, ) -> Iterator[torch.Tensor]: - """Keyframe-aware decode, yielding one chunk of ``(B, C, F, H, W)`` in ``[-1, 1]``. - Stages 1-3 run once on the whole volume for both streams -- so the keyframe stream - reaches stage 4 with *global* times -- then stages 4-5 run per tile with pixel blend. - With ``tiling_config=None`` that is a single tile and the whole thing is one pass. - """ content_shape = VideoLatentShape.from_torch_shape(latent.shape) content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) keyframes.validate(num_frames=content_pixel.frames) @@ -5642,12 +4645,6 @@ def _decode_video_with_keyframes( tiling_config: TilingConfig | None = None, generator: torch.Generator | None = None, ) -> Iterator[torch.Tensor]: - """Keyframe-aware decode, yielding float chunk(s) ``[f, h, w, c]`` in ``[0, 1]``. - Implementation of :meth:`decode_video` when ``keyframes`` is set. ``keyframes`` carries - single-frame latents plus their global pixel frame indices; every video token then - attends to the nearest planes through a joint neighborhood-attention window. - """ - def to_rgb(frames: torch.Tensor) -> torch.Tensor: return frames.add_(1).mul_(0.5).clamp_(0, 1) @@ -5669,7 +4666,6 @@ def _decode_temporal_group_isolated( *, complementary: bool, ) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Decode every tile of one temporal group in isolation and blend.""" group_temporal_len = curr_temporal_slice.stop - curr_temporal_slice.start group_shape = full_video_shape._replace(frames=group_temporal_len) full_torch_shape = full_video_shape.to_torch_shape() @@ -5751,18 +4747,6 @@ def _decode_pixels( # noqa: PLR0912, PLR0915 *, as_fhwc: bool = False, ) -> Iterator[torch.Tensor]: - """Decode latent to pixels, yielding temporal chunks. - Default yields raw ``(B, C, F, H, W)`` in ``[-1, 1]``. With ``as_fhwc=True`` - (used by :meth:`decode_video`), each chunk is materialized once as - contiguous ``[F, H, W, C]`` still in ``[-1, 1]`` - layout copy only; - range mapping stays in ``to_rgb``. - Stages 1-3 run once on the full volume; stages 4-5 run per tile with - pixel blend (one tile / one group when untiled or no real split). - Across temporal groups only the trailing overlap is retained between - iterations; exclusive frames are yielded before the next group decodes. - Peak residency is ~two tile extents (current buffer + still-live emit / - overlap stub), not a single ``tile + overlap`` slab. - """ content_shape = VideoLatentShape.from_torch_shape(latent.shape) content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) @@ -5824,7 +4808,6 @@ def _finalize(buf: torch.Tensor, wts: torch.Tensor | None) -> torch.Tensor: return (buf / wts).to(latent.dtype) def _narrow_content_cfhw(t: torch.Tensor, frames_keep: int) -> torch.Tensor: - """Spatial/temporal content crop as views (no ``.contiguous()``).""" x = t[:, :, :frames_keep] th, tw = content_pixel.height, content_pixel.width scale_h, scale_w = spatial_scale @@ -5945,7 +4928,6 @@ def forward( sample: torch.Tensor, generator: torch.Generator | None = None, ) -> torch.Tensor: - """Decode via ``_decode_pixels`` with ``tiling_config=None`` (single full tile).""" return next(self._decode_pixels(sample, tiling_config=None, generator=generator)) def decode_video( @@ -5956,14 +4938,6 @@ def decode_video( *, keyframes: DecodeKeyframes | None = None, ) -> Iterator[torch.Tensor]: - """Decode latent video, yielding float chunk(s) ``[f, h, w, c]`` in ``[0, 1]``. - Untiled and tiled both go through ``_decode_pixels``. Tiled decode may yield - multiple times when ``tiling_config.frames`` splits the video. - With ``keyframes`` this is :meth:`_decode_video_with_keyframes`; the argument exists on - every decoder so a caller can pass planes without first asking which VAE it holds. - Layout is packed once to contiguous FHWC on emit; ``to_rgb`` only does - inplace ``[-1, 1]→[0, 1]`` (no second realloc). - """ if keyframes is not None: yield from self._decode_video_with_keyframes(latent, keyframes, tiling_config, generator) return @@ -5975,8 +4949,6 @@ def to_rgb(frames: torch.Tensor) -> torch.Tensor: yield to_rgb(chunk) class LTX25DiffusionVideoDecoder(DiffusionVideoDecoder): - """DiffSynth-facing decoder with one full/tiled/keyframe interface.""" - def forward(self, sample, generator=None, keyframes=None): return self.decode(sample, generator=generator, keyframes=keyframes) From 1bd8ddcecb1423f017d66f779f041330f46f7ce3 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 10:48:36 +0800 Subject: [PATCH 11/31] Refine the LTX-2.5 example matrix and default negative prompt The dev weights are the general-purpose checkpoint, so they now cover the full feature set: OneStage/TwoStage T2AV and I2AV examples are added in both the standard and the low-VRAM directories. The distilled and INT8 checkpoints keep a single T2AV example each, so the distilled I2AV scripts are removed; IC-LoRA stays a distilled exception because upstream runs both of its stages on the distilled weights. Docs and README list the resulting ten-row matrix. LTX-2.5 also ships a longer default negative prompt than LTX-2.3: upstream prefixes the shared tag list with has_subtitles, has_blurbox, transition from black, transition to black and speech_ending_short. Add an "LTX-2.5" entry to pipe.default_negative_prompt and switch every 2.5 example, low-VRAM example and validation script to it (the T2A scripts previously used the placeholder "noise"). --- README.md | 9 ++-- diffsynth/pipelines/ltx2_audio_video.py | 14 +++++ docs/en/Model_Details/LTX-2.5.md | 10 ++-- docs/zh/Model_Details/LTX-2.5.md | 10 ++-- .../model_inference/LTX-2.5-A2V-TwoStage.py | 2 +- .../model_inference/LTX-2.5-I2AV-OneStage.py | 52 ++++++++++++++++++ ...edPipeline.py => LTX-2.5-I2AV-TwoStage.py} | 40 ++++++-------- .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 2 +- examples/ltx2/model_inference/LTX-2.5-T2A.py | 2 +- .../LTX-2.5-T2AV-DistilledPipeline.py | 2 +- .../LTX-2.5-T2AV-INT8-ConvRot.py | 2 +- .../model_inference/LTX-2.5-T2AV-OneStage.py | 44 +++++++++++++++ .../LTX-2.5-T2AV-TwoStage-Retake.py | 2 +- .../model_inference/LTX-2.5-T2AV-TwoStage.py | 47 ++++++++++++++++ .../LTX-2.5-A2V-TwoStage.py | 2 +- .../LTX-2.5-I2AV-OneStage.py | 53 +++++++++++++++++++ ...edPipeline.py => LTX-2.5-I2AV-TwoStage.py} | 40 ++++++-------- .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 2 +- .../model_inference_low_vram/LTX-2.5-T2A.py | 2 +- .../LTX-2.5-T2AV-DistilledPipeline.py | 2 +- .../LTX-2.5-T2AV-INT8-ConvRot.py | 2 +- .../LTX-2.5-T2AV-OneStage.py | 45 ++++++++++++++++ .../LTX-2.5-T2AV-TwoStage-Retake.py | 2 +- .../LTX-2.5-T2AV-TwoStage.py | 48 +++++++++++++++++ .../validate_full/LTX-2.5-T2AV.py | 2 +- .../validate_lora/LTX-2.5-T2AV.py | 2 +- 26 files changed, 371 insertions(+), 69 deletions(-) create mode 100644 examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py rename examples/ltx2/model_inference/{LTX-2.5-I2AV-DistilledPipeline.py => LTX-2.5-I2AV-TwoStage.py} (74%) create mode 100644 examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py create mode 100644 examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py rename examples/ltx2/model_inference_low_vram/{LTX-2.5-I2AV-DistilledPipeline.py => LTX-2.5-I2AV-TwoStage.py} (75%) create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py create mode 100644 examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py diff --git a/README.md b/README.md index 7d5445692..2ea0ccfa6 100644 --- a/README.md +++ b/README.md @@ -638,12 +638,15 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44 | JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — | -| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py) | -| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py) | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py) | - | - | - | - | | LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py) | - | - | - | - | | LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py) | - | - | - | - | -| LTX-2.5 | [Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler) | [code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | - | - | - | - | | LTX-2.5 | [Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2A.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler) | [code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | - | - | - | - | | LTX-2.5 | [Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py) | - | - | - | - | | LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) | diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 704014566..1a858135b 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -112,6 +112,20 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ), + "LTX-2.5": ( + "has_subtitles, has_blurbox, transition from black, transition to black, speech_ending_short, " + "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, " + "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, " + "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, " + "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of " + "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent " + "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny " + "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, " + "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, " + "off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward " + "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " + "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." + ), } @staticmethod diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md index 4c0e2de20..748abb27b 100644 --- a/docs/en/Model_Details/LTX-2.5.md +++ b/docs/en/Model_Details/LTX-2.5.md @@ -65,12 +65,15 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |Model ID|Extra parameters|Inference|Low-VRAM inference|Full training|Validate full|LoRA training|Validate LoRA| |-|-|-|-|-|-|-|-| -|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| -|[Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| -|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| ## Inference @@ -89,6 +92,7 @@ For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#in * `auto_duration_min_seconds` / `auto_duration_max_seconds`: lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. * `generate_video`: whether to generate video. Defaults to `True`. Set it to `False` to generate audio only (T2A); the video VAE and latent upsampler are then not required. * `use_diffusion_vae`: video decoder selection. `None` (default) selects by model version (LTX-2.5 uses the DiffVAE diffusion decoder); `False` uses the ConvVAE convolutional decoder (requires `ltx-2.5-video-vae-conv-bf16.safetensors`). +* Default negative prompt: `pipe.default_negative_prompt["LTX-2.5"]` prefixes the LTX-2/2.3 list with the 2.5-specific tags (`has_subtitles`, `has_blurbox`, `transition from black`, `transition to black`, `speech_ending_short`); all example scripts use this key. * `input_images` / `input_images_indexes`: keyframe images and their frame indexes. A single first frame gives image-to-video; first and last (or more) frames give keyframe interpolation. * `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: audio-to-video (A2V) and audio region retake. diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md index 3d0f67a21..a06d35cbf 100644 --- a/docs/zh/Model_Details/LTX-2.5.md +++ b/docs/zh/Model_Details/LTX-2.5.md @@ -65,12 +65,15 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |模型 ID|额外参数|推理|低显存推理|全量训练|全量训练后验证|LoRA 训练|LoRA 训练后验证| |-|-|-|-|-|-|-|-| -|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| -|[Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| -|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| ## 模型推理 @@ -89,6 +92,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 * `generate_video`: 是否生成视频，默认为 `True`。设置为 `False` 时只生成音频（T2A），此时无需加载视频 VAE 与 latent upsampler。 * `use_diffusion_vae`: 视频解码器选择。`None`（默认）表示按模型版本自动选择（LTX-2.5 使用 DiffVAE 扩散解码器），`False` 表示使用 ConvVAE 卷积解码器（需加载 `ltx-2.5-video-vae-conv-bf16.safetensors`）。 +* 默认负向提示词：`pipe.default_negative_prompt["LTX-2.5"]` 在 LTX-2/2.3 的列表之前增加了 2.5 专有标签（`has_subtitles`、`has_blurbox`、`transition from black`、`transition to black`、`speech_ending_short`），示例脚本均使用该键。 * `input_images` / `input_images_indexes`: 关键帧图像及其帧索引。传入首帧即为图生视频，传入首尾（或多帧）即为关键帧插值。 * `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: 音频驱动视频（A2V）与音频区域重生成。 diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py index 1aebfc997..61cd7275b 100644 --- a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -30,7 +30,7 @@ dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The example audio comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 duration = num_frames / frame_rate audio, audio_sample_rate = read_audio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py new file mode 100644 index 000000000..15b3bcc72 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -0,0 +1,52 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from PIL import Image +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example image comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=3.0, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py similarity index 74% rename from examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py rename to examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index c51a3a74f..fc24c939a 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -1,8 +1,8 @@ import torch from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 -from modelscope import dataset_snapshot_download from PIL import Image +from modelscope import dataset_snapshot_download vram_config = { "offload_dtype": torch.bfloat16, @@ -19,35 +19,31 @@ device="cuda", model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), ) - dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example frames come from the shared sample dataset, so reuse their paired prompt. -prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +# The example images come from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) last_frame = Image.open("data/example_video_dataset/ltx2/last_frame.png").convert("RGB").resize((width, height)) - -# Single-image I2AV +# first frame video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, - frame_rate=24, - cfg_scale=1.0, - num_inference_steps=8, - use_distilled_pipeline=True, - use_two_stage_pipeline=True, tiled=True, + cfg_scale=3.0, + use_two_stage_pipeline=True, input_images=[first_frame], input_images_indexes=[0], input_images_strength=1.0, @@ -55,25 +51,23 @@ write_video_audio_ltx2( video=video, audio=audio, - output_path="ltx2.5_distilled_i2av_first.mp4", + output_path="ltx2.5_twostage_i2av_first.mp4", fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, ) +pipe.clear_lora() -# Multi-keyframe interpolation: any frame indexes within num_frames are supported. +# Keyframe interpolation: any frames can be used by setting input_images and input_images_indexes within the range of num_frames. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, - frame_rate=24, - cfg_scale=1.0, - num_inference_steps=8, - use_distilled_pipeline=True, - use_two_stage_pipeline=True, tiled=True, + cfg_scale=3.0, + use_two_stage_pipeline=True, input_images=[first_frame, last_frame], input_images_indexes=[0, num_frames - 1], input_images_strength=1.0, @@ -81,7 +75,7 @@ write_video_audio_ltx2( video=video, audio=audio, - output_path="ltx2.5_distilled_i2av_keyframes.mp4", + output_path="ltx2.5_twostage_i2av_keyframes.mp4", fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, ) diff --git a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 0e3d49a05..9761ebb56 100644 --- a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -36,7 +36,7 @@ dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The reference video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 reference_video = VideoData("data/example_video_dataset/ltx2/video2.mp4", height=height // 4, width=width // 4).raw_data() video, audio = pipe( diff --git a/examples/ltx2/model_inference/LTX-2.5-T2A.py b/examples/ltx2/model_inference/LTX-2.5-T2A.py index 889962e05..a5762f227 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2A.py @@ -25,7 +25,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = "noise" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] _, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py index c1cf54f4c..2b9780367 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -27,7 +27,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width = 512 * 2, 768 * 2 # Automatic duration: one pipe call predicts the clip length from the prompt and generates it. video, audio = pipe( diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py index f5c98f6df..4d21e9708 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py @@ -25,7 +25,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 video, audio = pipe( prompt=prompt, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py new file mode 100644 index 000000000..1ed7fd2d8 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py @@ -0,0 +1,44 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=3.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py index 58772fc72..f1553a88c 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -31,7 +31,7 @@ dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The example video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 path = "data/example_video_dataset/ltx2/video2.mp4" diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py new file mode 100644 index 000000000..83086e722 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -0,0 +1,47 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=3.0, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py index 911817ad2..e59445b5f 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -31,7 +31,7 @@ dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The example audio comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 duration = num_frames / frame_rate audio, audio_sample_rate = read_audio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py new file mode 100644 index 000000000..ebebb388d --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -0,0 +1,53 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from PIL import Image +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +# The example image comes from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=3.0, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py similarity index 75% rename from examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py rename to examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index 2ca5577db..7844c9c10 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -1,8 +1,8 @@ import torch from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 -from modelscope import dataset_snapshot_download from PIL import Image +from modelscope import dataset_snapshot_download vram_config = { "offload_dtype": torch.float8_e5m2, @@ -19,36 +19,32 @@ device="cuda", model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) - dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example frames come from the shared sample dataset, so reuse their paired prompt. -prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +# The example images come from the shared sample dataset, so reuse its paired prompt. +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) last_frame = Image.open("data/example_video_dataset/ltx2/last_frame.png").convert("RGB").resize((width, height)) - -# Single-image I2AV +# first frame video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, - frame_rate=24, - cfg_scale=1.0, - num_inference_steps=8, - use_distilled_pipeline=True, - use_two_stage_pipeline=True, tiled=True, + cfg_scale=3.0, + use_two_stage_pipeline=True, input_images=[first_frame], input_images_indexes=[0], input_images_strength=1.0, @@ -56,25 +52,23 @@ write_video_audio_ltx2( video=video, audio=audio, - output_path="ltx2.5_distilled_i2av_first.mp4", + output_path="ltx2.5_twostage_i2av_first.mp4", fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, ) +pipe.clear_lora() -# Multi-keyframe interpolation: any frame indexes within num_frames are supported. +# Keyframe interpolation: any frames can be used by setting input_images and input_images_indexes within the range of num_frames. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, - frame_rate=24, - cfg_scale=1.0, - num_inference_steps=8, - use_distilled_pipeline=True, - use_two_stage_pipeline=True, tiled=True, + cfg_scale=3.0, + use_two_stage_pipeline=True, input_images=[first_frame, last_frame], input_images_indexes=[0, num_frames - 1], input_images_strength=1.0, @@ -82,7 +76,7 @@ write_video_audio_ltx2( video=video, audio=audio, - output_path="ltx2.5_distilled_i2av_keyframes.mp4", + output_path="ltx2.5_twostage_i2av_keyframes.mp4", fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, ) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 0158873eb..473710b37 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -37,7 +37,7 @@ dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The reference video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 reference_video = VideoData("data/example_video_dataset/ltx2/video2.mp4", height=height // 4, width=width // 4).raw_data() video, audio = pipe( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py index 73e010f39..c18d2f062 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py @@ -26,7 +26,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = "noise" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] _, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index def7f6f52..f31ed8d18 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -28,7 +28,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width = 512 * 2, 768 * 2 # Automatic duration: one pipe call predicts the clip length from the prompt and generates it. video, audio = pipe( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py index dc75e2c78..d70c6b5fa 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py @@ -26,7 +26,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 video, audio = pipe( prompt=prompt, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py new file mode 100644 index 000000000..a128f1067 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -0,0 +1,45 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=3.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py index c68d8164d..2a9ceb95a 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -32,7 +32,7 @@ dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The example video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 path = "data/example_video_dataset/ltx2/video2.mp4" diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py new file mode 100644 index 000000000..d73ca8d43 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -0,0 +1,48 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.float8_e5m2, + "offload_device": "cpu", + "onload_dtype": torch.float8_e5m2, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + cfg_scale=3.0, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py index c07fdcd83..582b1437f 100644 --- a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py @@ -23,7 +23,7 @@ ], ) prompt = "A beautiful sunset over the ocean." -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512, 768, 121 video, audio = pipe( prompt=prompt, diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py index 1f104b6be..019f731d2 100644 --- a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py @@ -24,7 +24,7 @@ ) pipe.load_lora(pipe.dit, "models/train/LTX2.5-T2AV_lora/epoch-4.safetensors") prompt = "A beautiful sunset over the ocean." -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512, 768, 121 video, audio = pipe( prompt=prompt, From a2d4c338cbef2134dc982d944c66e71e40ca1738 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 10:48:47 +0800 Subject: [PATCH 12/31] Reformat the LTX-2.5 Gemma config literal LTX25_GEMMA_CONFIG was a raw pprint dump with single quotes and hanging indentation. Restyle it to match the other text encoder configs (double quotes, four-space indent, one key per line, inline leaf dicts) and collapse the 48-entry layer_types list into its repeating five-sliding-plus-one-full pattern. Formatting only: the resolved Gemma4UnifiedConfig.to_dict() is byte-identical to the previous one. --- diffsynth/models/ltx25_text_encoder.py | 233 ++++++++++--------------- 1 file changed, 97 insertions(+), 136 deletions(-) diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index 7c464c1f5..a869bb9a1 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -15,142 +15,103 @@ ) -LTX25_GEMMA_CONFIG = {'architectures': ['Gemma4UnifiedForConditionalGeneration'], - 'audio_config': {'_name_or_path': '', - 'architectures': None, - 'audio_embed_dim': 640, - 'chunk_size_feed_forward': 0, - 'dtype': 'bfloat16', - 'id2label': {'0': 'LABEL_0', '1': 'LABEL_1'}, - 'initializer_range': 0.02, - 'is_encoder_decoder': False, - 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, - 'model_type': 'gemma4_unified_audio', - 'output_attentions': False, - 'output_hidden_states': False, - 'problem_type': None, - 'return_dict': True, - 'rms_norm_eps': 1e-06}, - 'audio_token_id': 258881, - 'boa_token_id': 256000, - 'boi_token_id': 255999, - 'dtype': 'bfloat16', - 'eoa_token_index': 258883, - 'eoi_token_id': 258882, - 'eos_token_id': [1, 106], - 'gemma_version': 'gemma4-12b-ltx-v1', - 'image_token_id': 258880, - 'initializer_range': 0.02, - 'model_type': 'gemma4_unified', - 'text_config': {'attention_bias': False, - 'attention_dropout': 0.0, - 'attention_k_eq_v': True, - 'bos_token_id': 2, - 'dtype': 'bfloat16', - 'enable_moe_block': False, - 'eos_token_id': 1, - 'final_logit_softcapping': 30.0, - 'global_head_dim': 512, - 'head_dim': 256, - 'hidden_activation': 'gelu_pytorch_tanh', - 'hidden_size': 3840, - 'hidden_size_per_layer_input': 0, - 'initializer_range': 0.02, - 'intermediate_size': 15360, - 'layer_types': ['sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'sliding_attention', - 'full_attention'], - 'max_position_embeddings': 262144, - 'model_type': 'gemma4_unified_text', - 'moe_intermediate_size': None, - 'num_attention_heads': 16, - 'num_experts': None, - 'num_global_key_value_heads': 1, - 'num_hidden_layers': 48, - 'num_key_value_heads': 8, - 'num_kv_shared_layers': 0, - 'pad_token_id': 0, - 'rms_norm_eps': 1e-06, - 'rope_parameters': {'full_attention': {'partial_rotary_factor': 0.25, - 'rope_theta': 1000000.0, - 'rope_type': 'proportional'}, - 'sliding_attention': {'rope_theta': 10000.0, 'rope_type': 'default'}}, - 'sliding_window': 1024, - 'tie_word_embeddings': True, - 'top_k_experts': None, - 'use_bidirectional_attention': 'vision', - 'use_cache': True, - 'use_double_wide_mlp': False, - 'vocab_size': 262144, - 'vocab_size_per_layer_input': 262144}, - 'tie_word_embeddings': True, - 'transformers_version': '5.10.1', - 'video_token_id': 258884, - 'vision_config': {'_name_or_path': '', - 'architectures': None, - 'chunk_size_feed_forward': 0, - 'dtype': 'bfloat16', - 'id2label': {'0': 'LABEL_0', '1': 'LABEL_1'}, - 'initializer_range': 0.02, - 'is_encoder_decoder': False, - 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, - 'mm_embed_dim': 3840, - 'mm_posemb_size': 1120, - 'model_type': 'gemma4_unified_vision', - 'num_soft_tokens': 280, - 'output_attentions': False, - 'output_hidden_states': False, - 'output_proj_dims': 3840, - 'patch_size': 16, - 'pooling_kernel_size': 3, - 'problem_type': None, - 'return_dict': True, - 'rms_norm_eps': 1e-06}} +LTX25_GEMMA_CONFIG = { + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "audio_config": { + "_name_or_path": "", + "architectures": None, + "audio_embed_dim": 640, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, + "initializer_range": 0.02, + "is_encoder_decoder": False, + "label2id": {"LABEL_0": 0, "LABEL_1": 1}, + "model_type": "gemma4_unified_audio", + "output_attentions": False, + "output_hidden_states": False, + "problem_type": None, + "return_dict": True, + "rms_norm_eps": 1e-06, + }, + "audio_token_id": 258881, + "boa_token_id": 256000, + "boi_token_id": 255999, + "dtype": "bfloat16", + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "eos_token_id": [1, 106], + "gemma_version": "gemma4-12b-ltx-v1", + "image_token_id": 258880, + "initializer_range": 0.02, + "model_type": "gemma4_unified", + "text_config": { + "attention_bias": False, + "attention_dropout": 0.0, + "attention_k_eq_v": True, + "bos_token_id": 2, + "dtype": "bfloat16", + "enable_moe_block": False, + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 3840, + "hidden_size_per_layer_input": 0, + "initializer_range": 0.02, + "intermediate_size": 15360, + "layer_types": (["sliding_attention"] * 5 + ["full_attention"]) * 8, + "max_position_embeddings": 262144, + "model_type": "gemma4_unified_text", + "moe_intermediate_size": None, + "num_attention_heads": 16, + "num_experts": None, + "num_global_key_value_heads": 1, + "num_hidden_layers": 48, + "num_key_value_heads": 8, + "num_kv_shared_layers": 0, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": {"partial_rotary_factor": 0.25, "rope_theta": 1000000.0, "rope_type": "proportional"}, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + }, + "sliding_window": 1024, + "tie_word_embeddings": True, + "top_k_experts": None, + "use_bidirectional_attention": "vision", + "use_cache": True, + "use_double_wide_mlp": False, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144, + }, + "tie_word_embeddings": True, + "transformers_version": "5.10.1", + "video_token_id": 258884, + "vision_config": { + "_name_or_path": "", + "architectures": None, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, + "initializer_range": 0.02, + "is_encoder_decoder": False, + "label2id": {"LABEL_0": 0, "LABEL_1": 1}, + "mm_embed_dim": 3840, + "mm_posemb_size": 1120, + "model_type": "gemma4_unified_vision", + "num_soft_tokens": 280, + "output_attentions": False, + "output_hidden_states": False, + "output_proj_dims": 3840, + "patch_size": 16, + "pooling_kernel_size": 3, + "problem_type": None, + "return_dict": True, + "rms_norm_eps": 1e-06, + }, +} class LTX25TextEncoder(torch.nn.Module): From c5828fb14aa14451bcfab9c999ded6385f58cab8 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 12:15:54 +0800 Subject: [PATCH 13/31] Restore upstream preprocessor init timing and stage-2 LoRA default Moving _init_preprocessors into __init__ made the argument preprocessors capture the plain modules before VRAM management replaces them with wrappers, so under CPU offload the preprocessors called stale modules whose weights never onload. PR #1602 built the preprocessors at the start of every forward, after wrapping; restore that timing and drop the OwnerModuleProxy workaround it had made unnecessary. Also restore the upstream stage2_lora_strength default of 0.8, which the LTX-2/2.3 two-stage examples rely on, and pass 1.0 explicitly in the LTX-2.5 two-stage examples whose distilled stage-2 LoRA needs full strength. --- diffsynth/models/ltx2_dit.py | 44 ++----------------- diffsynth/pipelines/ltx2_audio_video.py | 2 +- .../model_inference/LTX-2.5-A2V-TwoStage.py | 1 + .../model_inference/LTX-2.5-I2AV-TwoStage.py | 1 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 1 + .../model_inference/LTX-2.5-T2AV-TwoStage.py | 1 + .../LTX-2.5-A2V-TwoStage.py | 1 + .../LTX-2.5-I2AV-TwoStage.py | 1 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 1 + .../LTX-2.5-T2AV-TwoStage.py | 1 + 10 files changed, 13 insertions(+), 41 deletions(-) diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py index ec50642f8..8600ce5ee 100644 --- a/diffsynth/models/ltx2_dit.py +++ b/diffsynth/models/ltx2_dit.py @@ -613,25 +613,6 @@ class TransformerArgs: ) -class OwnerModuleProxy: - # Preprocessors are plain objects holding references to modules owned by the model. - # VRAM management replaces the owned modules with wrappers, so resolve them lazily - # through the owner to avoid calling stale unwrapped modules. - def __init__(self, owner: torch.nn.Module, name: str): - self.owner = owner - self.name = name - - @property - def module(self): - return getattr(self.owner, self.name) - - def __call__(self, *args, **kwargs): - return self.module(*args, **kwargs) - - def __getattr__(self, item): - return getattr(self.module, item) - - class TransformerArgsPreprocessor: def __init__( # noqa: PLR0913 self, @@ -1368,7 +1349,6 @@ def __init__( # noqa: PLR0913 self.use_prompt_adaln_single = use_prompt_adaln_single self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding self.use_tokenwise_av_ca_scale_shift = use_tokenwise_av_ca_scale_shift - cross_pe_max_pos = None if model_type.is_video_enabled(): if positional_embedding_max_pos is None: positional_embedding_max_pos = [20, 2048, 2048] @@ -1396,12 +1376,10 @@ def __init__( # noqa: PLR0913 ) if model_type.is_video_enabled() and model_type.is_audio_enabled(): - cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier self.audio_cross_attention_dim = audio_cross_attention_dim self._init_audio_video(num_scale_shift_values=4) - self._init_preprocessors(cross_pe_max_pos) # Initialize transformer blocks self._init_transformer_blocks( num_layers=num_layers, @@ -1600,24 +1578,6 @@ def _init_preprocessors( caption_projection=getattr(self, "audio_caption_projection", None), prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), ) - self._bind_preprocessor_modules() - - def _bind_preprocessor_modules(self) -> None: - owned_names = {id(module): name for name, module in self.named_children()} - pending = [ - getattr(self, "video_args_preprocessor", None), - getattr(self, "audio_args_preprocessor", None), - ] - while pending: - preprocessor = pending.pop() - if preprocessor is None: - continue - for attr_name, value in list(vars(preprocessor).items()): - if isinstance(value, torch.nn.Module): - if id(value) in owned_names: - setattr(preprocessor, attr_name, OwnerModuleProxy(self, owned_names[id(value)])) - elif isinstance(value, (TransformerArgsPreprocessor, MultiModalTransformerArgsPreprocessor)): - pending.append(value) def _init_transformer_blocks( self, @@ -1787,6 +1747,10 @@ def forward( video_keyframes_mask=None, perturbations=None, ): + cross_pe_max_pos = None + if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): + cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) + self._init_preprocessors(cross_pe_max_pos) video = ( Modality( video_latents, diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 1a858135b..d05564f2e 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -135,7 +135,7 @@ def from_pretrained( model_configs: list[ModelConfig] = [], tokenizer_config: ModelConfig = ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config: Optional[ModelConfig] = None, - stage2_lora_strength: float = 1.0, + stage2_lora_strength: float = 0.8, vram_limit: float = None, gemma_path: Union[str, Path, None] = None, load_duration_head: bool = False, diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py index 61cd7275b..33f072d4a 100644 --- a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -25,6 +25,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index fc24c939a..03502e8ad 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -25,6 +25,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") # The example images come from the shared sample dataset, so reuse its paired prompt. diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py index f1553a88c..d8e8053cb 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -26,6 +26,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py index 83086e722..f5c57beb9 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -23,6 +23,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py index e59445b5f..d04e31ab2 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -25,6 +25,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index 7844c9c10..200735d85 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -25,6 +25,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py index 2a9ceb95a..bedc597a8 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -26,6 +26,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py index d73ca8d43..ff64e9227 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -23,6 +23,7 @@ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" From db3c9bcec4ff2b7862b8c3f7e1a7aca1f7d40ba8 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 12:49:36 +0800 Subject: [PATCH 14/31] verified dit clean --- diffsynth/models/ltx2_dit.py | 100 ++++++++++------------------------- 1 file changed, 28 insertions(+), 72 deletions(-) diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py index 8600ce5ee..dbcc919a3 100644 --- a/diffsynth/models/ltx2_dit.py +++ b/diffsynth/models/ltx2_dit.py @@ -833,6 +833,10 @@ def prepare( if cross_modality.sigma.ndim != 1: raise ValueError("Cross modality sigma must be a 1D tensor") + cross_timestep = cross_modality.sigma.view( + modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:]) + ) + cross_pe = self.simple_preprocessor._prepare_positional_embeddings( positions=modality.positions[:, 0:1, :], inner_dim=self.audio_cross_attention_dim, @@ -843,11 +847,11 @@ def prepare( ) cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( - modality_timesteps=modality.timesteps, - cross_modality_sigma=cross_modality.sigma, + timestep=modality.timesteps if self.use_tokenwise_av_ca_scale_shift else cross_timestep, timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, batch_size=transformer_args.x.shape[0], hidden_dtype=modality.latent.dtype, + gate_timestep=cross_timestep if self.use_tokenwise_av_ca_scale_shift else None, ) return replace( @@ -859,24 +863,26 @@ def prepare( def _prepare_cross_attention_timestep( self, - modality_timesteps: torch.Tensor, - cross_modality_sigma: torch.Tensor, + timestep: torch.Tensor | None, timestep_scale_multiplier: int, batch_size: int, hidden_dtype: torch.dtype, + gate_timestep: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare A-V cross-attention AdaLN inputs.""" + """Prepare cross attention timestep embeddings.""" + timestep = timestep * timestep_scale_multiplier + av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier - cross_timestep = cross_modality_sigma.view(batch_size, 1, *[1] * (modality_timesteps.ndim - 2)) - scale_shift_input = modality_timesteps if self.use_tokenwise_av_ca_scale_shift else cross_timestep + # LTX-2.5 drives the scale/shift AdaLN per token while the gate stays scalar. + gate_input = (gate_timestep * timestep_scale_multiplier).flatten() if gate_timestep is not None else timestep.flatten() scale_shift_timestep, _ = self.cross_scale_shift_adaln( - (scale_shift_input * timestep_scale_multiplier).flatten(), + timestep.flatten(), hidden_dtype=hidden_dtype, ) scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) gate_noise_timestep, _ = self.cross_gate_adaln( - (cross_timestep * timestep_scale_multiplier * av_ca_factor).flatten(), + gate_input * av_ca_factor, hidden_dtype=hidden_dtype, ) gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) @@ -1400,21 +1406,6 @@ def _adaln_embedding_coefficient(self) -> int: def _keyframes_embedding(self) -> torch.Tensor | None: return getattr(self, "keyframes_abs_pos_embedding", None) - @property - def supports_keyframes_abs_pos_embedding(self) -> bool: - embedding = self._keyframes_embedding() - return embedding is not None and not embedding.is_meta - - def enable_keyframes_abs_pos_embedding(self) -> None: - if not self.model_type.is_video_enabled(): - raise ValueError("The keyframe absolute-position embedding is a video-stream parameter") - existing = self._keyframes_embedding() - if existing is not None and not existing.is_meta: - return - shape = existing.shape if existing is not None else (1, self.inner_dim) - self.use_keyframes_abs_pos_embedding = True - self.keyframes_abs_pos_embedding = torch.nn.Parameter(torch.zeros(shape, dtype=torch.bfloat16)) - def _init_video( self, in_channels: int, @@ -1426,12 +1417,8 @@ def _init_video( # Video input components self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient) - self.prompt_adaln_single = AdaLayerNormSingle( - self.inner_dim, embedding_coefficient=2 - ) if self.cross_attention_adaln and self.use_prompt_adaln_single else None - self.keyframes_abs_pos_embedding = ( - torch.nn.Parameter(torch.zeros(1, self.inner_dim)) if self.use_keyframes_abs_pos_embedding else None - ) + self.prompt_adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln and self.use_prompt_adaln_single else None + self.keyframes_abs_pos_embedding = torch.nn.Parameter(torch.zeros(1, self.inner_dim)) if self.use_keyframes_abs_pos_embedding else None # Video caption projection if caption_channels is not None: @@ -1458,9 +1445,7 @@ def _init_audio( self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) self.audio_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=self._adaln_embedding_coefficient) - self.audio_prompt_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, embedding_coefficient=2 - ) if self.cross_attention_adaln and self.use_prompt_adaln_single else None + self.audio_prompt_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln and self.use_prompt_adaln_single else None # Audio caption projection if caption_channels is not None: @@ -1632,6 +1617,13 @@ def _init_transformer_blocks( ) def set_gradient_checkpointing(self, enable: bool) -> None: + """Enable or disable gradient checkpointing for transformer blocks. + Gradient checkpointing trades compute for memory by recomputing activations + during the backward pass instead of storing them. This can significantly + reduce memory usage at the cost of ~20-30% slower training. + Args: + enable: Whether to enable gradient checkpointing + """ self._enable_gradient_checkpointing = enable def _process_transformer_blocks( @@ -1740,49 +1732,13 @@ def forward( sigma, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, - video_context_mask=None, - audio_context_mask=None, - video_attention_mask=None, - audio_attention_mask=None, video_keyframes_mask=None, - perturbations=None, ): cross_pe_max_pos = None if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) self._init_preprocessors(cross_pe_max_pos) - video = ( - Modality( - video_latents, - sigma, - video_timesteps, - video_positions, - video_context, - context_mask=video_context_mask, - attention_mask=video_attention_mask, - keyframes_mask=video_keyframes_mask, - ) - if video_latents is not None - else None - ) - audio = ( - Modality( - audio_latents, - sigma, - audio_timesteps, - audio_positions, - audio_context, - context_mask=audio_context_mask, - attention_mask=audio_attention_mask, - ) - if audio_latents is not None - else None - ) - vx, ax = self._forward( - video=video, - audio=audio, - perturbations=perturbations, - use_gradient_checkpointing=use_gradient_checkpointing, - use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, - ) + video = Modality(video_latents, sigma, video_timesteps, video_positions, video_context, keyframes_mask=video_keyframes_mask) if video_latents is not None else None + audio = Modality(audio_latents, sigma, audio_timesteps, audio_positions, audio_context) if audio_latents is not None else None + vx, ax = self._forward(video=video, audio=audio, perturbations=None, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload) return vx, ax From 6769df47bca0f4efa72b55e0614098cc557103c0 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 13:35:56 +0800 Subject: [PATCH 15/31] Restore upstream VRAM maps and VAE code, merge the 2.5 tokenizer into the TE file The VRAM module maps and the audio/video VAE modules carried changes that only existed to support our own earlier fine-grained map experiment: PerChannelStatistics, Snake and sibling vocoder entries, the LTX2AudioEncoder entry, the LTXModel block-level entry and the use-site STFT/mel casts. With the maps back to the upstream layout the vocoder and audio encoder fall back to the whole-model wrap exactly as LTX-2.3 runs today, so those code changes are reverted too, together with the string-to-Enum coercions in the video VAE whose only consumers were string extra_kwargs we added for the 2.5 entries. The only remaining map delta is the LTX-2.5 DiffusionVideoDecoder entry, whose flat class path and leaf-level wrapping the 2.5 decoder requires. LTX25GemmaTokenizer moves verbatim into ltx25_text_encoder.py, mirroring how the LTX-2.3 tokenizer lives inside its TE file, and the 2.5 registry entries are regrouped by model hash with concrete dev-file example comments. --- diffsynth/configs/model_configs.py | 52 +++++++++---------- .../configs/vram_management_module_maps.py | 31 +---------- diffsynth/models/ltx25_text_encoder.py | 50 ++++++++++++++++++ diffsynth/models/ltx25_tokenizer.py | 51 ------------------ diffsynth/models/ltx2_audio_vae.py | 7 +-- diffsynth/models/ltx2_video_vae.py | 10 ---- diffsynth/pipelines/ltx2_audio_video.py | 3 +- 7 files changed, 82 insertions(+), 122 deletions(-) delete mode 100644 diffsynth/models/ltx25_tokenizer.py diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py index 7c2ea4f1b..82745fdd4 100644 --- a/diffsynth/configs/model_configs.py +++ b/diffsynth/configs/model_configs.py @@ -896,7 +896,7 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter", }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-*-transformer-bf16.safetensors") + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors") "model_hash": "7960c5dc4626650824e36f65a8e992e9", "model_name": "ltx25_dit", "model_class": "diffsynth.models.ltx2_dit.LTXModel", @@ -904,7 +904,14 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-*-transformer-comfy-int8-convrot.safetensors") + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors") + "model_hash": "7960c5dc4626650824e36f65a8e992e9", + "model_name": "ltx25_embeddings_connectors", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-comfy-int8-convrot.safetensors") "model_hash": "57343d320cac0bbba58a488b8ebe7187", "model_name": "ltx25_dit", "model_class": "diffsynth.models.ltx2_dit.LTXModel", @@ -913,19 +920,19 @@ "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["timestep_embedder.linear_1", "timestep_embedder.linear_2", "adaln_single.linear", "audio_adaln_single.linear", "prompt_adaln_single.linear", "audio_prompt_adaln_single.linear", "av_ca_a2v_gate_adaln_single.linear", "av_ca_audio_scale_shift_adaln_single.linear", "av_ca_v2a_gate_adaln_single.linear", "av_ca_video_scale_shift_adaln_single.linear", "patchify_proj", "audio_patchify_proj", "proj_out", "audio_proj_out", "to_gate_logits"]}, }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") - "model_hash": "055700dc619165899bebb5162f699cd2", - "model_name": "ltx25_text_encoder", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-comfy-int8-convrot.safetensors") + "model_hash": "57343d320cac0bbba58a488b8ebe7187", + "model_name": "ltx25_embeddings_connectors", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["to_gate_logits"]}, }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") - "model_hash": "4743ded7a5725b6589bccdb62512723b", + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") + "model_hash": "055700dc619165899bebb5162f699cd2", "model_name": "ltx25_text_encoder", "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", - "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["lm_head", "embedding_projection", "patch_dense"]}, }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") @@ -935,11 +942,12 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25FeatureExtractorStateDictConverter", }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors") - "model_hash": "7960c5dc4626650824e36f65a8e992e9", - "model_name": "ltx25_embeddings_connectors", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") + "model_hash": "4743ded7a5725b6589bccdb62512723b", + "model_name": "ltx25_text_encoder", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["lm_head", "embedding_projection", "patch_dense"]}, }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") @@ -949,20 +957,12 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25FeatureExtractorStateDictConverter", "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_aggregate_embed", "audio_aggregate_embed"]}, }, - { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors") - "model_hash": "57343d320cac0bbba58a488b8ebe7187", - "model_name": "ltx25_embeddings_connectors", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", - "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["to_gate_logits"]}, - }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors") "model_hash": "e19205490f01801d0a7b6d3aba61e26e", "model_name": "ltx25_video_vae_encoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", - "extra_kwargs": {"encoder_version": "ltx-2.3", "latent_log_var": "constant"}, + "extra_kwargs": {"encoder_version": "ltx-2.3"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", }, { @@ -978,7 +978,7 @@ "model_hash": "a1d642eecae96baa9c31d4e405564f49", "model_name": "ltx25_conv_video_vae_encoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", - "extra_kwargs": {"encoder_version": "ltx-2.3", "latent_log_var": "uniform"}, + "extra_kwargs": {"encoder_version": "ltx-2.3"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", }, { @@ -986,7 +986,7 @@ "model_hash": "a1d642eecae96baa9c31d4e405564f49", "model_name": "ltx25_conv_video_vae_decoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder", - "extra_kwargs": {"decoder_version": "ltx-2.3", "decoder_spatial_padding_mode": "zeros"}, + "extra_kwargs": {"decoder_version": "ltx-2.3"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter", }, { diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index f31749a3c..dc41d9e69 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -271,7 +271,6 @@ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", }, "diffsynth.models.ltx2_dit.LTXModel": { - "diffsynth.models.ltx2_dit.BasicAVTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, @@ -283,15 +282,10 @@ "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors": { - "diffsynth.models.ltx25_text_encoder.LTX25Embeddings1DConnector": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx25_text_encoder.LTX25BasicTransformerBlock1D": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.ltx25_text_encoder.LTX25FeatureExtractorV2": { - "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", - }, "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", @@ -302,38 +296,15 @@ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": { - "diffsynth.models.ltx2_video_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": { - "diffsynth.models.ltx2_video_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder": { - "diffsynth.models.ltx2_audio_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", - }, "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": { - "diffsynth.models.ltx2_audio_vae.PerChannelStatistics": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_audio_vae.LTX2Vocoder": { - "diffsynth.models.ltx2_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.SnakeBeta": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.Activation1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.DownSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule", - }, - "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE": { - "diffsynth.models.ltx2_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.SnakeBeta": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.Activation1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_audio_vae.DownSample1d": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule", }, diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index a869bb9a1..36a140323 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -1,8 +1,14 @@ import copy +import json import math +from pathlib import Path from typing import NamedTuple +import numpy as np import torch +from safetensors import safe_open +from tokenizers import Tokenizer +from transformers import PreTrainedTokenizerFast from .ltx2_common import rms_norm from .ltx2_dit import ( @@ -165,6 +171,50 @@ def forward(self, *args, **kwargs): return self.model(*args, **kwargs) +class LTX25GemmaTokenizer: + def __init__(self, model_path: str | Path, max_length: int = 1024): + model_path = Path(model_path) + with safe_open(model_path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + if "tokenizer_json" not in handle.keys(): + raise ValueError(f"{model_path} does not contain packed tokenizer_json assets.") + tokenizer_bytes = handle.get_tensor("tokenizer_json").detach().cpu().numpy().astype(np.uint8).tobytes() + raw_config = metadata.get("tokenizer_config.json") + if raw_config is None and "hf_asset__tokenizer_config.json" in handle.keys(): + raw_config = handle.get_tensor("hf_asset__tokenizer_config.json").detach().cpu().numpy().astype(np.uint8).tobytes().decode() + config = json.loads(raw_config) if raw_config else {} + ignored = {"tokenizer_class", "auto_map", "model_max_length", "backend", "is_local", "local_files_only", "processor_class", "added_tokens_decoder"} + config = {key: value for key, value in config.items() if key not in ignored} + self.tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer.from_buffer(tokenizer_bytes), + model_max_length=max_length, + **config, + ) + self.tokenizer.model_max_length = max_length + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + self.max_length = max_length + + def tokenize_with_weights(self, text: str) -> dict[str, list[tuple[int, int]]]: + text = text.strip() + bos_id = self.tokenizer.bos_token_id + if bos_id is None: + raise ValueError("Packed Gemma tokenizer has no BOS token id.") + encoded = self.tokenizer(text, padding=False, truncation=True, max_length=self.max_length, return_tensors="pt") + input_ids = encoded.input_ids[0].tolist() + if not input_ids or input_ids[0] != bos_id: + input_ids = [bos_id, *input_ids][: self.max_length] + padded = self.tokenizer.pad( + {"input_ids": [input_ids]}, + padding="max_length", + max_length=self.max_length, + return_tensors="pt", + return_attention_mask=True, + ) + return {"gemma": list(zip(padded.input_ids[0].tolist(), padded.attention_mask[0].tolist(), strict=True))} + + def norm_and_concat_per_token_rms( encoded_text: torch.Tensor, attention_mask: torch.Tensor, diff --git a/diffsynth/models/ltx25_tokenizer.py b/diffsynth/models/ltx25_tokenizer.py deleted file mode 100644 index d6fbbb95a..000000000 --- a/diffsynth/models/ltx25_tokenizer.py +++ /dev/null @@ -1,51 +0,0 @@ -import json -from pathlib import Path - -import numpy as np -from safetensors import safe_open -from tokenizers import Tokenizer -from transformers import PreTrainedTokenizerFast - - -class LTX25GemmaTokenizer: - def __init__(self, model_path: str | Path, max_length: int = 1024): - model_path = Path(model_path) - with safe_open(model_path, framework="pt", device="cpu") as handle: - metadata = handle.metadata() or {} - if "tokenizer_json" not in handle.keys(): - raise ValueError(f"{model_path} does not contain packed tokenizer_json assets.") - tokenizer_bytes = handle.get_tensor("tokenizer_json").detach().cpu().numpy().astype(np.uint8).tobytes() - raw_config = metadata.get("tokenizer_config.json") - if raw_config is None and "hf_asset__tokenizer_config.json" in handle.keys(): - raw_config = handle.get_tensor("hf_asset__tokenizer_config.json").detach().cpu().numpy().astype(np.uint8).tobytes().decode() - config = json.loads(raw_config) if raw_config else {} - ignored = {"tokenizer_class", "auto_map", "model_max_length", "backend", "is_local", "local_files_only", "processor_class", "added_tokens_decoder"} - config = {key: value for key, value in config.items() if key not in ignored} - self.tokenizer = PreTrainedTokenizerFast( - tokenizer_object=Tokenizer.from_buffer(tokenizer_bytes), - model_max_length=max_length, - **config, - ) - self.tokenizer.model_max_length = max_length - self.tokenizer.padding_side = "left" - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - self.max_length = max_length - - def tokenize_with_weights(self, text: str) -> dict[str, list[tuple[int, int]]]: - text = text.strip() - bos_id = self.tokenizer.bos_token_id - if bos_id is None: - raise ValueError("Packed Gemma tokenizer has no BOS token id.") - encoded = self.tokenizer(text, padding=False, truncation=True, max_length=self.max_length, return_tensors="pt") - input_ids = encoded.input_ids[0].tolist() - if not input_ids or input_ids[0] != bos_id: - input_ids = [bos_id, *input_ids][: self.max_length] - padded = self.tokenizer.pad( - {"input_ids": [input_ids]}, - padding="max_length", - max_length=self.max_length, - return_tensors="pt", - return_attention_mask=True, - ) - return {"gemma": list(zip(padded.input_ids[0].tolist(), padded.attention_mask[0].tolist(), strict=True))} diff --git a/diffsynth/models/ltx2_audio_vae.py b/diffsynth/models/ltx2_audio_vae.py index b23fefae9..8a58f9724 100644 --- a/diffsynth/models/ltx2_audio_vae.py +++ b/diffsynth/models/ltx2_audio_vae.py @@ -1285,8 +1285,10 @@ def get_padding(kernel_size: int, dilation: int = 1) -> int: return int((kernel_size * dilation - dilation) / 2) +# --------------------------------------------------------------------------- # Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2 # Adopted from https://github.com/NVIDIA/BigVGAN +# --------------------------------------------------------------------------- def _sinc(x: torch.Tensor) -> torch.Tensor: @@ -1716,8 +1718,7 @@ def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: y = y.unsqueeze(1) # (B, 1, T) left_pad = max(0, self.win_length - self.hop_length) # causal: left-only y = F.pad(y, (left_pad, 0)) - forward_basis = self.forward_basis.to(device=y.device, dtype=y.dtype) - spec = F.conv1d(y, forward_basis, stride=self.hop_length, padding=0) + spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) n_freqs = spec.shape[1] // 2 real, imag = spec[:, :n_freqs], spec[:, n_freqs:] magnitude = torch.sqrt(real**2 + imag**2) @@ -1760,7 +1761,7 @@ def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, """ magnitude, phase = self.stft_fn(y) energy = torch.norm(magnitude, dim=1) - mel = torch.matmul(self.mel_basis.to(device=magnitude.device, dtype=magnitude.dtype), magnitude) + mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) log_mel = torch.log(torch.clamp(mel, min=1e-5)) return log_mel, magnitude, phase, energy diff --git a/diffsynth/models/ltx2_video_vae.py b/diffsynth/models/ltx2_video_vae.py index 80aaa60f9..a70dc80e8 100644 --- a/diffsynth/models/ltx2_video_vae.py +++ b/diffsynth/models/ltx2_video_vae.py @@ -1335,12 +1335,6 @@ def __init__( encoder_version: str = "ltx-2", ): super().__init__() - if isinstance(norm_layer, str): - norm_layer = NormLayerType(norm_layer) - if isinstance(latent_log_var, str): - latent_log_var = LogVarianceType(latent_log_var) - if isinstance(encoder_spatial_padding_mode, str): - encoder_spatial_padding_mode = PaddingModeType(encoder_spatial_padding_mode) if encoder_version == "ltx-2": encoder_blocks = [ ['res_x', {'num_layers': 4}], @@ -1800,10 +1794,6 @@ def __init__( base_channels: int = 128, ): super().__init__() - if isinstance(norm_layer, str): - norm_layer = NormLayerType(norm_layer) - if isinstance(decoder_spatial_padding_mode, str): - decoder_spatial_padding_mode = PaddingModeType(decoder_spatial_padding_mode) # Spatiotemporal downscaling between decoded video space and VAE latents. # According to the LTXV paper, the standard configuration downsamples diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index d05564f2e..818798960 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -19,8 +19,7 @@ from ..models.ltx2_text_encoder import LTX2TextEncoder, LTX2TextEncoderPostModules, LTXVGemmaTokenizer from ..models.ltx2_upsampler import LTX2LatentUpsampler from ..models.ltx2_video_vae import LTX2VideoDecoder, LTX2VideoEncoder, VideoLatentPatchifier -from ..models.ltx25_text_encoder import LTX25TextEncoderPostModules -from ..models.ltx25_tokenizer import LTX25GemmaTokenizer +from ..models.ltx25_text_encoder import LTX25GemmaTokenizer, LTX25TextEncoderPostModules from ..utils.data.audio import convert_to_stereo, resample_waveform from ..utils.data.media_io_ltx2 import ltx2_preprocess From cd4ec93a52801c262f26ff919c7b5e633adb73c4 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 14:56:12 +0800 Subject: [PATCH 16/31] Unify LTX-2.5 model loading with the shared registry and a Repackage checkpoint Register the LTX-2.5 components under the same model names as LTX-2/2.3 so from_pretrained fetches every version through one shared path; only the tokenizer construction stays version-specific. The text encoder post modules are the single component whose weights are scattered across two upstream files (feature extractor in the TE checkpoint, connectors in the transformer checkpoint), so they are packed into DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors in target key layout, needing no state dict converter, and the packed Gemma4 tokenizer assets are unpacked into an HF-style tokenizer directory so tokenizer_config works like LTX-2.3. Pipeline cleanups along the way: the video decoder is selected by availability (conv decoder first, diffusion decoder as fallback), gemma_path and load_duration_head are gone, and the upstream section comments, import layout and upsampler fetch position are restored. --- diffsynth/configs/model_configs.py | 52 +++------ diffsynth/models/ltx25_text_encoder.py | 22 +--- diffsynth/pipelines/ltx2_audio_video.py | 107 +++++++----------- .../ltx25_text_encoder.py | 21 ---- docs/en/Model_Details/LTX-2.5.md | 9 +- docs/zh/Model_Details/LTX-2.5.md | 9 +- .../model_inference/LTX-2.5-A2V-TwoStage.py | 2 + .../model_inference/LTX-2.5-I2AV-OneStage.py | 2 + .../model_inference/LTX-2.5-I2AV-TwoStage.py | 2 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 2 + examples/ltx2/model_inference/LTX-2.5-T2A.py | 3 +- .../LTX-2.5-T2AV-DistilledPipeline.py | 4 +- .../LTX-2.5-T2AV-INT8-ConvRot.py | 2 + .../model_inference/LTX-2.5-T2AV-OneStage.py | 2 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 2 + .../model_inference/LTX-2.5-T2AV-TwoStage.py | 2 + .../LTX-2.5-A2V-TwoStage.py | 2 + .../LTX-2.5-I2AV-OneStage.py | 2 + .../LTX-2.5-I2AV-TwoStage.py | 2 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 2 + .../model_inference_low_vram/LTX-2.5-T2A.py | 3 +- .../LTX-2.5-T2AV-DistilledPipeline.py | 3 +- .../LTX-2.5-T2AV-INT8-ConvRot.py | 2 + .../LTX-2.5-T2AV-OneStage.py | 2 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 2 + .../LTX-2.5-T2AV-TwoStage.py | 2 + .../full/LTX-2.5-T2AV-splited.sh | 8 +- .../lora/LTX-2.5-T2AV-splited-test.sh | 8 +- .../lora/LTX-2.5-T2AV-splited.sh | 8 +- .../scripts/split_model_statedicts_ltx2.5.py | 51 +++++++++ .../validate_full/LTX-2.5-T2AV.py | 2 + .../validate_lora/LTX-2.5-T2AV.py | 2 + 32 files changed, 176 insertions(+), 168 deletions(-) create mode 100644 examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py index 82745fdd4..643268bc3 100644 --- a/diffsynth/configs/model_configs.py +++ b/diffsynth/configs/model_configs.py @@ -898,69 +898,45 @@ { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors") "model_hash": "7960c5dc4626650824e36f65a8e992e9", - "model_name": "ltx25_dit", + "model_name": "ltx2_dit", "model_class": "diffsynth.models.ltx2_dit.LTXModel", "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True, "use_tokenwise_av_ca_scale_shift": True}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", }, - { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors") - "model_hash": "7960c5dc4626650824e36f65a8e992e9", - "model_name": "ltx25_embeddings_connectors", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", - }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-comfy-int8-convrot.safetensors") "model_hash": "57343d320cac0bbba58a488b8ebe7187", - "model_name": "ltx25_dit", + "model_name": "ltx2_dit", "model_class": "diffsynth.models.ltx2_dit.LTXModel", "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True, "use_tokenwise_av_ca_scale_shift": True}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["timestep_embedder.linear_1", "timestep_embedder.linear_2", "adaln_single.linear", "audio_adaln_single.linear", "prompt_adaln_single.linear", "audio_prompt_adaln_single.linear", "av_ca_a2v_gate_adaln_single.linear", "av_ca_audio_scale_shift_adaln_single.linear", "av_ca_v2a_gate_adaln_single.linear", "av_ca_video_scale_shift_adaln_single.linear", "patchify_proj", "audio_patchify_proj", "proj_out", "audio_proj_out", "to_gate_logits"]}, }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-comfy-int8-convrot.safetensors") - "model_hash": "57343d320cac0bbba58a488b8ebe7187", - "model_name": "ltx25_embeddings_connectors", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25EmbeddingsConnectors", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25EmbeddingsConnectorsStateDictConverter", - "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["to_gate_logits"]}, + # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors") + "model_hash": "8f3c146ff3d584392236c5b29d26146c", + "model_name": "ltx2_text_encoder_post_modules", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules", }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") "model_hash": "055700dc619165899bebb5162f699cd2", - "model_name": "ltx25_text_encoder", + "model_name": "ltx2_text_encoder", "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", }, - { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") - "model_hash": "055700dc619165899bebb5162f699cd2", - "model_name": "ltx25_feature_extractor", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25FeatureExtractorV2", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25FeatureExtractorStateDictConverter", - }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") "model_hash": "4743ded7a5725b6589bccdb62512723b", - "model_name": "ltx25_text_encoder", + "model_name": "ltx2_text_encoder", "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["lm_head", "embedding_projection", "patch_dense"]}, }, - { - # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") - "model_hash": "4743ded7a5725b6589bccdb62512723b", - "model_name": "ltx25_feature_extractor", - "model_class": "diffsynth.models.ltx25_text_encoder.LTX25FeatureExtractorV2", - "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25FeatureExtractorStateDictConverter", - "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_aggregate_embed", "audio_aggregate_embed"]}, - }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors") "model_hash": "e19205490f01801d0a7b6d3aba61e26e", - "model_name": "ltx25_video_vae_encoder", + "model_name": "ltx2_video_vae_encoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", "extra_kwargs": {"encoder_version": "ltx-2.3"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", @@ -976,7 +952,7 @@ { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors") "model_hash": "a1d642eecae96baa9c31d4e405564f49", - "model_name": "ltx25_conv_video_vae_encoder", + "model_name": "ltx2_video_vae_encoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", "extra_kwargs": {"encoder_version": "ltx-2.3"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", @@ -984,7 +960,7 @@ { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors") "model_hash": "a1d642eecae96baa9c31d4e405564f49", - "model_name": "ltx25_conv_video_vae_decoder", + "model_name": "ltx2_video_vae_decoder", "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder", "extra_kwargs": {"decoder_version": "ltx-2.3"}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter", @@ -992,21 +968,21 @@ { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") "model_hash": "c2488315f13356abb806f9f217f1e803", - "model_name": "ltx25_audio_vae_decoder", + "model_name": "ltx2_audio_vae_decoder", "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter", }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") "model_hash": "c2488315f13356abb806f9f217f1e803", - "model_name": "ltx25_audio_vocoder", + "model_name": "ltx2_audio_vocoder", "model_class": "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter", }, { # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") "model_hash": "c2488315f13356abb806f9f217f1e803", - "model_name": "ltx25_audio_vae_encoder", + "model_name": "ltx2_audio_vae_encoder", "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder", "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter", }, diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index 36a140323..adb1c3f39 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -1,13 +1,9 @@ import copy -import json import math from pathlib import Path from typing import NamedTuple -import numpy as np import torch -from safetensors import safe_open -from tokenizers import Tokenizer from transformers import PreTrainedTokenizerFast from .ltx2_common import rms_norm @@ -173,23 +169,7 @@ def forward(self, *args, **kwargs): class LTX25GemmaTokenizer: def __init__(self, model_path: str | Path, max_length: int = 1024): - model_path = Path(model_path) - with safe_open(model_path, framework="pt", device="cpu") as handle: - metadata = handle.metadata() or {} - if "tokenizer_json" not in handle.keys(): - raise ValueError(f"{model_path} does not contain packed tokenizer_json assets.") - tokenizer_bytes = handle.get_tensor("tokenizer_json").detach().cpu().numpy().astype(np.uint8).tobytes() - raw_config = metadata.get("tokenizer_config.json") - if raw_config is None and "hf_asset__tokenizer_config.json" in handle.keys(): - raw_config = handle.get_tensor("hf_asset__tokenizer_config.json").detach().cpu().numpy().astype(np.uint8).tobytes().decode() - config = json.loads(raw_config) if raw_config else {} - ignored = {"tokenizer_class", "auto_map", "model_max_length", "backend", "is_local", "local_files_only", "processor_class", "added_tokens_decoder"} - config = {key: value for key, value in config.items() if key not in ignored} - self.tokenizer = PreTrainedTokenizerFast( - tokenizer_object=Tokenizer.from_buffer(tokenizer_bytes), - model_max_length=max_length, - **config, - ) + self.tokenizer = PreTrainedTokenizerFast.from_pretrained(str(model_path), local_files_only=True, model_max_length=max_length) self.tokenizer.model_max_length = max_length self.tokenizer.padding_side = "left" if self.tokenizer.pad_token is None: diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 818798960..c6d347e98 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -1,27 +1,30 @@ -from functools import partial -from pathlib import Path -from typing import Optional, Union - +import torch, types import numpy as np -import torch +from PIL import Image from einops import repeat +from typing import Optional, Union +from einops import rearrange +import numpy as np from PIL import Image from tqdm import tqdm +from typing import Optional from transformers import AutoImageProcessor, Gemma3Processor +from functools import partial -from ..core import ModelConfig from ..core.device.npu_compatible_device import get_device_type from ..diffusion import FlowMatchScheduler +from ..core import ModelConfig from ..diffusion.base_pipeline import BasePipeline, PipelineUnit -from ..models.ltx2_audio_vae import LTX2AudioDecoder, LTX2AudioEncoder, LTX2Vocoder, AudioPatchifier, AudioProcessor -from ..models.ltx2_common import AudioLatentShape, VIDEO_SCALE_FACTORS, VideoLatentShape, VideoPixelShape, get_pixel_coords -from ..models.ltx2_dit import LTXModel + from ..models.ltx2_text_encoder import LTX2TextEncoder, LTX2TextEncoderPostModules, LTXVGemmaTokenizer +from ..models.ltx2_dit import LTXModel +from ..models.ltx2_video_vae import LTX2VideoEncoder, LTX2VideoDecoder, VideoLatentPatchifier +from ..models.ltx2_audio_vae import LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder, AudioPatchifier, AudioProcessor from ..models.ltx2_upsampler import LTX2LatentUpsampler -from ..models.ltx2_video_vae import LTX2VideoDecoder, LTX2VideoEncoder, VideoLatentPatchifier -from ..models.ltx25_text_encoder import LTX25GemmaTokenizer, LTX25TextEncoderPostModules -from ..utils.data.audio import convert_to_stereo, resample_waveform +from ..models.ltx2_common import VideoLatentShape, AudioLatentShape, VideoPixelShape, get_pixel_coords, VIDEO_SCALE_FACTORS +from ..models.ltx25_text_encoder import LTX25GemmaTokenizer from ..utils.data.media_io_ltx2 import ltx2_preprocess +from ..utils.data.audio import convert_to_stereo, resample_waveform class LTX2AudioVideoPipeline(BasePipeline): @@ -44,7 +47,6 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.video_vae_encoder: LTX2VideoEncoder = None self.video_vae_decoder: LTX2VideoDecoder = None self.diffusion_video_vae_decoder = None - self.conv_video_vae_decoder: LTX2VideoDecoder = None self.audio_vae_encoder: LTX2AudioEncoder = None self.audio_vae_decoder: LTX2AudioDecoder = None self.audio_vocoder: LTX2Vocoder = None @@ -136,70 +138,38 @@ def from_pretrained( stage2_lora_config: Optional[ModelConfig] = None, stage2_lora_strength: float = 0.8, vram_limit: float = None, - gemma_path: Union[str, Path, None] = None, - load_duration_head: bool = False, ): + # Initialize pipeline pipe = LTX2AudioVideoPipeline(device=device, torch_dtype=torch_dtype) model_pool = pipe.download_and_load_models(model_configs, vram_limit) - ltx25_text_encoder = model_pool.fetch_model("ltx25_text_encoder") - ltx25_dit = model_pool.fetch_model("ltx25_dit") - pipe.is_ltx25 = ltx25_text_encoder is not None or ltx25_dit is not None + # Fetch models + pipe.text_encoder = model_pool.fetch_model("ltx2_text_encoder") + pipe.dit = model_pool.fetch_model("ltx2_dit") + pipe.is_ltx25 = getattr(pipe.dit, "use_tokenwise_av_ca_scale_shift", False) + tokenizer_config.download_if_necessary() if pipe.is_ltx25: - if ltx25_text_encoder is None or ltx25_dit is None: - raise ValueError("LTX-2.5 requires both ltx25_text_encoder and ltx25_dit components.") - if gemma_path is None: - for model_config in model_configs: - if isinstance(model_config.path, str) and "text_encoders" in model_config.path: - gemma_path = model_config.path - break - if gemma_path is None: - raise ValueError("gemma_path is required for the packed LTX-2.5 Gemma4 tokenizer assets.") - pipe.text_encoder = ltx25_text_encoder - pipe.text_encoder.reset_non_persistent_buffers() - pipe.tokenizer = LTX25GemmaTokenizer(gemma_path) - feature_extractor = model_pool.fetch_model("ltx25_feature_extractor") - connectors = model_pool.fetch_model("ltx25_embeddings_connectors") - if feature_extractor is None or connectors is None: - raise ValueError("LTX-2.5 requires ltx25_feature_extractor and ltx25_embeddings_connectors components.") - pipe.text_encoder_post_modules = LTX25TextEncoderPostModules( - feature_extractor=feature_extractor, - connectors=connectors, - ) - # The container holds VRAM-wrapped modules but is not itself wrapped, so mark it - # for load_models_to_device to offload/onload its wrapped children. - pipe.text_encoder_post_modules.vram_management_enabled = True - pipe.dit = ltx25_dit - pipe.video_vae_encoder = model_pool.fetch_model("ltx25_video_vae_encoder") - if pipe.video_vae_encoder is None: - pipe.video_vae_encoder = model_pool.fetch_model("ltx25_conv_video_vae_encoder") - pipe.diffusion_video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") - pipe.conv_video_vae_decoder = model_pool.fetch_model("ltx25_conv_video_vae_decoder") - pipe.audio_vae_decoder = model_pool.fetch_model("ltx25_audio_vae_decoder") - pipe.audio_vocoder = model_pool.fetch_model("ltx25_audio_vocoder") - pipe.audio_vae_encoder = model_pool.fetch_model("ltx25_audio_vae_encoder") - pipe.duration_head = model_pool.fetch_model("ltx25_duration_head") - if load_duration_head and pipe.duration_head is None: - raise ValueError("load_duration_head=True requires an ltx25_duration_head ModelConfig.") + pipe.tokenizer = LTX25GemmaTokenizer(tokenizer_config.path) else: - pipe.text_encoder = model_pool.fetch_model("ltx2_text_encoder") - tokenizer_config.download_if_necessary() pipe.tokenizer = LTXVGemmaTokenizer(tokenizer_path=tokenizer_config.path) image_processor = AutoImageProcessor.from_pretrained(tokenizer_config.path, local_files_only=True) pipe.processor = Gemma3Processor(image_processor=image_processor, tokenizer=pipe.tokenizer.tokenizer) - pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") - pipe.dit = model_pool.fetch_model("ltx2_dit") - pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") - pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") - pipe.audio_vae_decoder = model_pool.fetch_model("ltx2_audio_vae_decoder") - pipe.audio_vocoder = model_pool.fetch_model("ltx2_audio_vocoder") - pipe.audio_vae_encoder = model_pool.fetch_model("ltx2_audio_vae_encoder") - + pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") + pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") + pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") + pipe.diffusion_video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") + pipe.audio_vae_decoder = model_pool.fetch_model("ltx2_audio_vae_decoder") + pipe.audio_vocoder = model_pool.fetch_model("ltx2_audio_vocoder") pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") + pipe.audio_vae_encoder = model_pool.fetch_model("ltx2_audio_vae_encoder") + pipe.duration_head = model_pool.fetch_model("ltx25_duration_head") + + # Stage 2 if stage2_lora_config is not None: pipe.stage2_lora_config = stage2_lora_config pipe.stage2_lora_strength = stage2_lora_strength + # VRAM Management pipe.vram_management_enabled = pipe.check_vram_management_state() return pipe @@ -295,11 +265,13 @@ def __call__( # progress_bar progress_bar_cmd=tqdm, ): + # Scheduler self.scheduler.set_timesteps( num_inference_steps, denoising_strength=denoising_strength, special_case="distilled_stage1" if use_distilled_pipeline else None, ) + # Inputs inputs_posi = {"prompt": prompt} inputs_nega = {"negative_prompt": negative_prompt} inputs_shared = { @@ -321,9 +293,11 @@ def __call__( "video_patchifier": self.video_patchifier, "audio_patchifier": self.audio_patchifier, "timestep_scale": 1.0 if self.is_ltx25 else 1000.0, } + # Stage 1 inputs_shared, inputs_posi, inputs_nega = self.denoise_stage( inputs_shared, inputs_posi, inputs_nega, self.units, cfg_scale, progress_bar_cmd ) + # Stage 2 inputs_shared, inputs_posi, inputs_nega = self.denoise_stage( inputs_shared, inputs_posi, @@ -333,6 +307,7 @@ def __call__( progress_bar_cmd, not inputs_shared["use_two_stage_pipeline"], ) + # Decode video = None if inputs_shared.get("generate_video", True): video_decoder_name = inputs_shared["video_decoder_name"] @@ -440,10 +415,12 @@ def process( if use_diffusion_vae: raise ValueError("Diffusion VAE decoding is only supported by LTX-2.5 checkpoints.") decoder_name = "video_vae_decoder" - elif use_diffusion_vae is not False: + elif use_diffusion_vae is True: decoder_name = "diffusion_video_vae_decoder" + elif use_diffusion_vae is False: + decoder_name = "video_vae_decoder" else: - decoder_name = "conv_video_vae_decoder" + decoder_name = "video_vae_decoder" if pipe.video_vae_decoder is not None else "diffusion_video_vae_decoder" if getattr(pipe, decoder_name) is None: requested = "DiffusionVAE" if decoder_name == "diffusion_video_vae_decoder" else "ConvVAE" diff --git a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py index fada522c9..b22ca72e9 100644 --- a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py +++ b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py @@ -14,24 +14,3 @@ def LTX25TextEncoderStateDictConverter(state_dict): state_dict_[new_name] = state_dict[name] state_dict_["model.lm_head.weight"] = state_dict_["model.model.language_model.embed_tokens.weight"] return state_dict_ - - -def LTX25FeatureExtractorStateDictConverter(state_dict): - state_dict_ = {} - for name in state_dict: - if name.startswith("text_embedding_projection."): - state_dict_[name.removeprefix("text_embedding_projection.")] = state_dict[name] - return state_dict_ - - -def LTX25EmbeddingsConnectorsStateDictConverter(state_dict): - state_dict_ = {} - for name in state_dict: - if name.startswith("model.diffusion_model.video_embeddings_connector."): - new_name = "video_connector." + name.removeprefix("model.diffusion_model.video_embeddings_connector.") - elif name.startswith("model.diffusion_model.audio_embeddings_connector."): - new_name = "audio_connector." + name.removeprefix("model.diffusion_model.audio_embeddings_connector.") - else: - continue - state_dict_[new_name] = state_dict[name] - return state_dict_ diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md index 748abb27b..97c709b83 100644 --- a/docs/en/Model_Details/LTX-2.5.md +++ b/docs/en/Model_Details/LTX-2.5.md @@ -44,7 +44,6 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - load_duration_head=True, ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.3"] @@ -72,7 +71,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| |[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| -|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| |[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| @@ -82,8 +81,8 @@ Models are loaded with `LTX2AudioVideoPipeline.from_pretrained`; see [Load Model LTX-2.5 related `from_pretrained` arguments: -* `load_duration_head`: require the Duration Head (needed for automatic duration prediction). Defaults to `False`. -* `gemma_path`: path to the Gemma4 checkpoint, used to load the packed tokenizer assets. When omitted, it is derived from the text encoder entry in `model_configs`. +* `tokenizer_config`: source of the tokenizer assets. The LTX-2.5 tokenizer is unpacked into the HF-style directory `DiffSynth-Studio/LTX-2.5-Repackage/tokenizer` (produced by `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py`), matching how LTX-2.3 references its tokenizer. +* `text_encoder_post_modules`: the LTX-2.5 feature extractor and embeddings connectors weights are packed into `DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors` (produced by `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py` from the TE and transformer checkpoints) and must be included in `model_configs`. * `stage2_lora_config`: the stage-2 distilled LoRA used for two-stage inference with the Dev weights. For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#inference). The new or LTX-2.5 specific `LTX2AudioVideoPipeline` arguments are: @@ -91,7 +90,7 @@ For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#in * `auto_duration`: predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. * `auto_duration_min_seconds` / `auto_duration_max_seconds`: lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. * `generate_video`: whether to generate video. Defaults to `True`. Set it to `False` to generate audio only (T2A); the video VAE and latent upsampler are then not required. -* `use_diffusion_vae`: video decoder selection. `None` (default) selects by model version (LTX-2.5 uses the DiffVAE diffusion decoder); `False` uses the ConvVAE convolutional decoder (requires `ltx-2.5-video-vae-conv-bf16.safetensors`). +* `use_diffusion_vae`: video decoder selection. `None` (default) prefers the loaded ConvVAE convolutional decoder and falls back to the DiffVAE diffusion decoder when it is absent; `True`/`False` force DiffVAE/ConvVAE respectively (ConvVAE requires `ltx-2.5-video-vae-conv-bf16.safetensors`). * Default negative prompt: `pipe.default_negative_prompt["LTX-2.5"]` prefixes the LTX-2/2.3 list with the 2.5-specific tags (`has_subtitles`, `has_blurbox`, `transition from black`, `transition to black`, `speech_ending_short`); all example scripts use this key. * `input_images` / `input_images_indexes`: keyframe images and their frame indexes. A single first frame gives image-to-video; first and last (or more) frames give keyframe interpolation. * `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: audio-to-video (A2V) and audio region retake. diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md index a06d35cbf..2f5f204e2 100644 --- a/docs/zh/Model_Details/LTX-2.5.md +++ b/docs/zh/Model_Details/LTX-2.5.md @@ -44,7 +44,6 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - load_duration_head=True, ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.3"] @@ -72,7 +71,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| |[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| -|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`,`load_duration_head=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| |[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| @@ -82,8 +81,8 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 `from_pretrained` 的 LTX-2.5 相关参数： -* `load_duration_head`: 是否要求加载 Duration Head（自动时长预测所需），默认为 `False`。 -* `gemma_path`: Gemma4 权重路径，用于加载内嵌的 tokenizer 资产。留空时自动从 `model_configs` 中的 text encoder 路径推导。 +* `tokenizer_config`：tokenizer 资产来源。LTX-2.5 的 tokenizer 已解包为 HF 目录形式 `DiffSynth-Studio/LTX-2.5-Repackage/tokenizer`（由 `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py` 生成），与 LTX-2.3 的用法一致。 +* `text_encoder_post_modules`：LTX-2.5 的 feature extractor 与 embeddings connectors 权重打包在 `DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors`（由 `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py` 从 TE 与 transformer 权重中提取），需要在 `model_configs` 中一并加载。 * `stage2_lora_config`: Dev 权重两阶段推理时使用的第二阶段 distilled-LoRA。 `LTX2AudioVideoPipeline` 的通用推理参数见 [LTX-2 文档](LTX-2.md#模型推理)，LTX-2.5 新增或特有的参数为： @@ -91,7 +90,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 * `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 * `generate_video`: 是否生成视频，默认为 `True`。设置为 `False` 时只生成音频（T2A），此时无需加载视频 VAE 与 latent upsampler。 -* `use_diffusion_vae`: 视频解码器选择。`None`（默认）表示按模型版本自动选择（LTX-2.5 使用 DiffVAE 扩散解码器），`False` 表示使用 ConvVAE 卷积解码器（需加载 `ltx-2.5-video-vae-conv-bf16.safetensors`）。 +* `use_diffusion_vae`：视频解码器选择。`None`（默认）优先使用已加载的 ConvVAE 卷积解码器，未加载时回退到 DiffVAE 扩散解码器；`True`/`False` 分别强制指定 DiffVAE/ConvVAE（ConvVAE 需加载 `ltx-2.5-video-vae-conv-bf16.safetensors`）。 * 默认负向提示词：`pipe.default_negative_prompt["LTX-2.5"]` 在 LTX-2/2.3 的列表之前增加了 2.5 专有标签（`has_subtitles`、`has_blurbox`、`transition from black`、`transition to black`、`speech_ending_short`），示例脚本均使用该键。 * `input_images` / `input_images_indexes`: 关键帧图像及其帧索引。传入首帧即为图生视频，传入首尾（或多帧）即为关键帧插值。 * `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: 音频驱动视频（A2V）与音频区域重生成。 diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py index 33f072d4a..b33877148 100644 --- a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py index 15b3bcc72..7beddfeab 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index 03502e8ad..791b2bbd6 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 9761ebb56..ef848cc73 100644 --- a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-T2A.py b/examples/ltx2/model_inference/LTX-2.5-T2A.py index a5762f227..39496d6a9 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2A.py @@ -15,13 +15,14 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - load_duration_head=True, ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py index 2b9780367..4b7b947a0 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -15,15 +15,16 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - load_duration_head=True, ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" @@ -53,3 +54,4 @@ fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, ) +print(f"saved to ltx2.5_distilled_t2av.mp4") diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py index 4d21e9708..767ff8078 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py index 1ed7fd2d8..bef059579 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py index d8e8053cb..b71f227ee 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -18,8 +18,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py index f5c57beb9..68d94f734 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py index d04e31ab2..1113fc8d9 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index ebebb388d..5ad07c872 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index 200735d85..ae6e9c75d 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 473710b37..657692b58 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -17,8 +17,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py index c18d2f062..03b78b25a 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py @@ -15,13 +15,14 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - load_duration_head=True, vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index f31ed8d18..c19c37cb9 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -15,15 +15,16 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), ], - load_duration_head=True, vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py index d70c6b5fa..5d75a4072 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py index a128f1067..2d811ffa0 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py index bedc597a8..3ff7c6d9f 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -18,8 +18,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py index ff64e9227..025a1189b 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh index 45ecf8a39..019e91e81 100644 --- a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh +++ b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh @@ -10,7 +10,8 @@ accelerate launch examples/ltx2/model_training/train.py \ --width 768 \ --num_frames 121 \ --dataset_repeat 1 \ - --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-5 \ --num_epochs 2 \ --remove_prefix_in_ckpt "pipe.dit." \ @@ -27,8 +28,9 @@ accelerate launch --config_file examples/ltx2/model_training/full/accelerate_con --width 768 \ --num_frames 121 \ --dataset_repeat 100 \ - --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ - --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-5 \ --num_epochs 2 \ --remove_prefix_in_ckpt "pipe.dit." \ diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh index 54e3a69d2..5fc25d78e 100644 --- a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh +++ b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh @@ -10,7 +10,8 @@ accelerate launch examples/ltx2/model_training/train.py \ --width 768 \ --num_frames 121 \ --dataset_repeat 1 \ - --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-4 \ --num_epochs 1 \ --remove_prefix_in_ckpt "pipe.dit." \ @@ -29,8 +30,9 @@ accelerate launch examples/ltx2/model_training/train.py \ --width 768 \ --num_frames 121 \ --dataset_repeat 1 \ - --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ - --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-4 \ --num_epochs 1 \ --remove_prefix_in_ckpt "pipe.dit." \ diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh index 9888d5867..637e1f18f 100644 --- a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh +++ b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh @@ -10,7 +10,8 @@ accelerate launch examples/ltx2/model_training/train.py \ --width 768 \ --num_frames 121 \ --dataset_repeat 1 \ - --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-4 \ --num_epochs 5 \ --remove_prefix_in_ckpt "pipe.dit." \ @@ -29,8 +30,9 @@ accelerate launch examples/ltx2/model_training/train.py \ --width 768 \ --num_frames 121 \ --dataset_repeat 100 \ - --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ - --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors" \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-4 \ --num_epochs 5 \ --remove_prefix_in_ckpt "pipe.dit." \ diff --git a/examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py b/examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py new file mode 100644 index 000000000..c0d0bf22d --- /dev/null +++ b/examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py @@ -0,0 +1,51 @@ +import os + +import numpy as np + +from safetensors import safe_open +from safetensors.torch import save_file + +from diffsynth import hash_state_dict_keys +from diffsynth.core import load_state_dict +from diffsynth.models.model_loader import ModelPool + +model_pool = ModelPool() +os.makedirs("models/DiffSynth-Studio/LTX-2.5-Repackage", exist_ok=True) + +SOURCES = ( + "models/Lightricks/LTX-2.5/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", + "models/Lightricks/LTX-2.5/diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", +) + +def target_name(name): + if name.startswith("text_embedding_projection."): + return "feature_extractor." + name.removeprefix("text_embedding_projection.") + if name.startswith("model.diffusion_model.video_embeddings_connector."): + return "connectors.video_connector." + name.removeprefix("model.diffusion_model.video_embeddings_connector.") + if name.startswith("model.diffusion_model.audio_embeddings_connector."): + return "connectors.audio_connector." + name.removeprefix("model.diffusion_model.audio_embeddings_connector.") + return None + + +text_encoder_post_modules_state_dict = {} +for path in SOURCES: + with safe_open(path, framework="pt", device="cpu") as handle: + for name in handle.keys(): + new_name = target_name(name) + if new_name is not None: + text_encoder_post_modules_state_dict[new_name] = handle.get_tensor(name) + +save_file(text_encoder_post_modules_state_dict, "models/DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors") +print(f"text_encoder_post_modules keys hash: {hash_state_dict_keys(text_encoder_post_modules_state_dict)}") +model_pool.auto_load_model("models/DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors") + +tokenizer_dir = "models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" +os.makedirs(tokenizer_dir, exist_ok=True) +with safe_open(SOURCES[0], framework="pt", device="cpu") as handle: + tokenizer_bytes = handle.get_tensor("tokenizer_json").detach().cpu().numpy().astype(np.uint8).tobytes() + config_bytes = handle.get_tensor("hf_asset__tokenizer_config.json").detach().cpu().numpy().astype(np.uint8).tobytes() +with open(os.path.join(tokenizer_dir, "tokenizer.json"), "wb") as f: + f.write(tokenizer_bytes) +with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "wb") as f: + f.write(config_bytes) +print(f"tokenizer assets written to {tokenizer_dir}") diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py index 582b1437f..abeac27ac 100644 --- a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(path="./models/train/LTX2.5-T2AV-full/epoch-1.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py index 019f731d2..3fe59258b 100644 --- a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py @@ -15,8 +15,10 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), model_configs=[ ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), From 92b78c1efba7bb703992ffdaea71bf7d1ee645ce Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 15:05:49 +0800 Subject: [PATCH 17/31] Drop the unused Gemma3 processor and unify tokenizer construction The Gemma3Processor was assigned but never read anywhere in the repository, so remove it together with its attribute and imports. Tokenizer construction collapses to picking the version-specific class and instantiating it with tokenizer_path, which is now the parameter name for LTX25GemmaTokenizer as well. --- diffsynth/models/ltx25_text_encoder.py | 4 ++-- diffsynth/pipelines/ltx2_audio_video.py | 10 ++-------- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index adb1c3f39..0e438c911 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -168,8 +168,8 @@ def forward(self, *args, **kwargs): class LTX25GemmaTokenizer: - def __init__(self, model_path: str | Path, max_length: int = 1024): - self.tokenizer = PreTrainedTokenizerFast.from_pretrained(str(model_path), local_files_only=True, model_max_length=max_length) + def __init__(self, tokenizer_path: str | Path, max_length: int = 1024): + self.tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path, local_files_only=True, model_max_length=max_length) self.tokenizer.model_max_length = max_length self.tokenizer.padding_side = "left" if self.tokenizer.pad_token is None: diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index c6d347e98..868cd482a 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -8,7 +8,6 @@ from PIL import Image from tqdm import tqdm from typing import Optional -from transformers import AutoImageProcessor, Gemma3Processor from functools import partial from ..core.device.npu_compatible_device import get_device_type @@ -41,7 +40,6 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.scheduler = FlowMatchScheduler("LTX-2") self.text_encoder: LTX2TextEncoder = None self.tokenizer: LTXVGemmaTokenizer = None - self.processor: Gemma3Processor = None self.text_encoder_post_modules: LTX2TextEncoderPostModules = None self.dit: LTXModel = None self.video_vae_encoder: LTX2VideoEncoder = None @@ -148,12 +146,8 @@ def from_pretrained( pipe.dit = model_pool.fetch_model("ltx2_dit") pipe.is_ltx25 = getattr(pipe.dit, "use_tokenwise_av_ca_scale_shift", False) tokenizer_config.download_if_necessary() - if pipe.is_ltx25: - pipe.tokenizer = LTX25GemmaTokenizer(tokenizer_config.path) - else: - pipe.tokenizer = LTXVGemmaTokenizer(tokenizer_path=tokenizer_config.path) - image_processor = AutoImageProcessor.from_pretrained(tokenizer_config.path, local_files_only=True) - pipe.processor = Gemma3Processor(image_processor=image_processor, tokenizer=pipe.tokenizer.tokenizer) + tokenizer_class = LTX25GemmaTokenizer if pipe.is_ltx25 else LTXVGemmaTokenizer + pipe.tokenizer = tokenizer_class(tokenizer_path=tokenizer_config.path) pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") From 08d2c85b95377f2e421e9444633f51cbb9bde1b5 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 16:03:13 +0800 Subject: [PATCH 18/31] Simplify video decoding and auto duration handling The decoder selector unit is gone: the pipeline keeps a single video_vae_decoder attribute, filled with the ConvVAE decoder when loaded and the DiffVAE decoder otherwise, and __call__ passes the tiling and seed arguments straight to decode(). Both decoders accept **kwargs so the uniform argument set needs no branching; the DiffVAE decoder takes seed/rand_device and builds its sampling generator internally, offset by 42 from the denoise seed, and tiles automatically when tiled is set, matching the official pipeline default. Auto duration validation moves into the AutoDuration unit, which now declares onload_model_names and clamps the requested bounds into (0, 20] instead of raising. --- diffsynth/models/ltx25_diffusion_video_vae.py | 43 +----- diffsynth/models/ltx2_video_vae.py | 1 + diffsynth/pipelines/ltx2_audio_video.py | 144 ++---------------- docs/en/Model_Details/LTX-2.5.md | 2 +- docs/zh/Model_Details/LTX-2.5.md | 2 +- 5 files changed, 25 insertions(+), 167 deletions(-) diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index b26614cef..0cc33f503 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -940,17 +940,6 @@ def enable_size_axis( TilingConfig = TileSizeConfig | TileCountConfig -class AutoTiling: - __slots__ = () - - def __repr__(self) -> str: - return "AUTO_TILING" - - -AUTO_TILING = AutoTiling() - - -PipelineTiling = TilingConfig | AutoTiling | None def _assert_video_on_vae_grid( @@ -4997,33 +4986,17 @@ def decode( self, latent, tiled=False, - tile_size_in_pixels=None, - tile_overlap_in_pixels=None, - tile_size_in_frames=None, - tile_overlap_in_frames=None, - generator=None, + seed=None, + rand_device="cpu", keyframes=None, + **kwargs, ): + generator = torch.Generator(device=rand_device).manual_seed(seed) if seed is not None else None tiling_config = None - if tiled is True or tiled is AUTO_TILING: - if tiled is AUTO_TILING or tile_size_in_pixels is None or tile_overlap_in_pixels is None or tile_size_in_frames is None or tile_overlap_in_frames is None: - tiling_config = self.auto_tiling_config(latent, keyframes=keyframes) - if tiling_config is None: - raise ValueError("Automatic DiffVAE tiling requires a CUDA device with queryable free memory.") - else: - if isinstance(tile_size_in_pixels, Sequence) and not isinstance(tile_size_in_pixels, (str, bytes)): - tile_height, tile_width = tile_size_in_pixels - else: - tile_height = tile_width = int(tile_size_in_pixels) - if isinstance(tile_overlap_in_pixels, Sequence) and not isinstance(tile_overlap_in_pixels, (str, bytes)): - overlap_height, overlap_width = tile_overlap_in_pixels - else: - overlap_height = overlap_width = int(tile_overlap_in_pixels) - tiling_config = TileSizeConfig( - frames=DimensionSizeConfig(int(tile_size_in_frames), int(tile_overlap_in_frames)), - height=DimensionSizeConfig(int(tile_height), int(overlap_height)), - width=DimensionSizeConfig(int(tile_width), int(overlap_width)), - ) + if tiled: + tiling_config = self.auto_tiling_config(latent, keyframes=keyframes) + if tiling_config is None: + raise ValueError("Automatic DiffVAE tiling requires a CUDA device with queryable free memory.") iterator = ( self._decode_pixels_with_keyframes(latent, keyframes, tiling_config, generator=generator) if keyframes is not None diff --git a/diffsynth/models/ltx2_video_vae.py b/diffsynth/models/ltx2_video_vae.py index a70dc80e8..c98dad160 100644 --- a/diffsynth/models/ltx2_video_vae.py +++ b/diffsynth/models/ltx2_video_vae.py @@ -2187,6 +2187,7 @@ def decode( tile_overlap_in_pixels: Optional[int] = 128, tile_size_in_frames: Optional[int] = 128, tile_overlap_in_frames: Optional[int] = 24, + **kwargs, ) -> torch.Tensor: if tiled: tiling_config = TilingConfig( diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 868cd482a..af47fd452 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -17,7 +17,7 @@ from ..models.ltx2_text_encoder import LTX2TextEncoder, LTX2TextEncoderPostModules, LTXVGemmaTokenizer from ..models.ltx2_dit import LTXModel -from ..models.ltx2_video_vae import LTX2VideoEncoder, LTX2VideoDecoder, VideoLatentPatchifier +from ..models.ltx2_video_vae import LTX2VideoEncoder, VideoLatentPatchifier from ..models.ltx2_audio_vae import LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder, AudioPatchifier, AudioProcessor from ..models.ltx2_upsampler import LTX2LatentUpsampler from ..models.ltx2_common import VideoLatentShape, AudioLatentShape, VideoPixelShape, get_pixel_coords, VIDEO_SCALE_FACTORS @@ -43,8 +43,7 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.text_encoder_post_modules: LTX2TextEncoderPostModules = None self.dit: LTXModel = None self.video_vae_encoder: LTX2VideoEncoder = None - self.video_vae_decoder: LTX2VideoDecoder = None - self.diffusion_video_vae_decoder = None + self.video_vae_decoder = None self.audio_vae_encoder: LTX2AudioEncoder = None self.audio_vae_decoder: LTX2AudioDecoder = None self.audio_vocoder: LTX2Vocoder = None @@ -59,7 +58,6 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.in_iteration_models = ("dit",) self.units = [ LTX2AudioVideoUnit_PipelineChecker(), - LTX2AudioVideoUnit_VideoDecoderSelector(), LTX2AudioVideoUnit_PromptEmbedder(), LTX2AudioVideoUnit_AutoDuration(), LTX2AudioVideoUnit_ShapeChecker(), @@ -151,7 +149,8 @@ def from_pretrained( pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") - pipe.diffusion_video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") + if pipe.video_vae_decoder is None: + pipe.video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") pipe.audio_vae_decoder = model_pool.fetch_model("ltx2_audio_vae_decoder") pipe.audio_vocoder = model_pool.fetch_model("ltx2_audio_vocoder") pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") @@ -250,7 +249,6 @@ def __call__( tile_overlap_in_pixels: int = 128, tile_size_in_frames: int = 128, tile_overlap_in_frames: int = 24, - use_diffusion_vae: Optional[bool] = None, # Special Pipelines use_two_stage_pipeline: bool = False, stage2_spatial_upsample_factor: int = 2, @@ -282,7 +280,6 @@ def __call__( "cfg_scale": cfg_scale, "tiled": tiled, "tile_size_in_pixels": tile_size_in_pixels, "tile_overlap_in_pixels": tile_overlap_in_pixels, "tile_size_in_frames": tile_size_in_frames, "tile_overlap_in_frames": tile_overlap_in_frames, - "use_diffusion_vae": use_diffusion_vae, "use_two_stage_pipeline": use_two_stage_pipeline, "use_distilled_pipeline": use_distilled_pipeline, "clear_lora_before_state_two": clear_lora_before_state_two, "stage2_spatial_upsample_factor": stage2_spatial_upsample_factor, "video_patchifier": self.video_patchifier, "audio_patchifier": self.audio_patchifier, "timestep_scale": 1.0 if self.is_ltx25 else 1000.0, @@ -304,10 +301,10 @@ def __call__( # Decode video = None if inputs_shared.get("generate_video", True): - video_decoder_name = inputs_shared["video_decoder_name"] - self.load_models_to_device([video_decoder_name]) - video_decoder = getattr(self, video_decoder_name) - video = video_decoder.decode(inputs_shared["video_latents"], **inputs_shared["video_decode_kwargs"]) + if self.video_vae_decoder is None: + raise ValueError("No video decoder component is loaded.") + self.load_models_to_device(["video_vae_decoder"]) + video = self.video_vae_decoder.decode(inputs_shared["video_latents"], tiled=tiled, tile_size_in_pixels=tile_size_in_pixels, tile_overlap_in_pixels=tile_overlap_in_pixels, tile_size_in_frames=tile_size_in_frames, tile_overlap_in_frames=tile_overlap_in_frames, seed=None if seed is None else seed + 42, rand_device=rand_device) video = self.vae_output_to_video(video) retake_audio = inputs_shared.get("retake_audio") denoise_mask_audio = inputs_shared.get("denoise_mask_audio") @@ -350,113 +347,14 @@ def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inpu raise ValueError("Two-stage pipeline requested, but stage2_lora_config is not set in the pipeline.") if pipe.upsampler is None: raise ValueError("Two-stage pipeline requested, but upsampler model is not loaded in the pipeline.") - if inputs_shared.get("auto_duration", False): - if pipe.duration_head is None: - raise ValueError( - "Automatic duration requires an ltx25_duration_head ModelConfig in from_pretrained()." - ) - min_seconds = inputs_shared["auto_duration_min_seconds"] - max_seconds = inputs_shared["auto_duration_max_seconds"] - if min_seconds <= 0 or max_seconds < min_seconds: - raise ValueError("Automatic duration requires 0 < min_seconds <= max_seconds.") return inputs_shared, inputs_posi, inputs_nega -class LTX2AudioVideoUnit_VideoDecoderSelector(PipelineUnit): - def __init__(self): - super().__init__( - input_params=( - "use_diffusion_vae", - "seed", - "rand_device", - "tiled", - "tile_size_in_pixels", - "tile_overlap_in_pixels", - "tile_size_in_frames", - "tile_overlap_in_frames", - "generate_video", - ), - output_params=("video_decoder_name", "video_decode_kwargs", "noise_generator"), - ) - - def process( - self, - pipe: LTX2AudioVideoPipeline, - use_diffusion_vae, - seed, - rand_device, - tiled, - tile_size_in_pixels, - tile_overlap_in_pixels, - tile_size_in_frames, - tile_overlap_in_frames, - generate_video=True, - ): - if generate_video is False: - return { - "video_decoder_name": None, - "video_decode_kwargs": {}, - "noise_generator": None, - } - if pipe.scheduler.training: - # Caching stages never decode, so the decoder component is not required here. - return { - "video_decoder_name": None, - "video_decode_kwargs": {}, - "noise_generator": None, - } - if not pipe.is_ltx25: - if use_diffusion_vae: - raise ValueError("Diffusion VAE decoding is only supported by LTX-2.5 checkpoints.") - decoder_name = "video_vae_decoder" - elif use_diffusion_vae is True: - decoder_name = "diffusion_video_vae_decoder" - elif use_diffusion_vae is False: - decoder_name = "video_vae_decoder" - else: - decoder_name = "video_vae_decoder" if pipe.video_vae_decoder is not None else "diffusion_video_vae_decoder" - - if getattr(pipe, decoder_name) is None: - requested = "DiffusionVAE" if decoder_name == "diffusion_video_vae_decoder" else "ConvVAE" - raise ValueError(f"{requested} decoder was requested but its model component is not loaded.") - - decode_kwargs = { - "tiled": tiled, - "tile_size_in_pixels": tile_size_in_pixels, - "tile_overlap_in_pixels": tile_overlap_in_pixels, - "tile_size_in_frames": tile_size_in_frames, - "tile_overlap_in_frames": tile_overlap_in_frames, - } - if decoder_name == "diffusion_video_vae_decoder": - conv_defaults = (512, 128, 128, 24) - current_values = ( - tile_size_in_pixels, - tile_overlap_in_pixels, - tile_size_in_frames, - tile_overlap_in_frames, - ) - if any(value is None for value in current_values) or current_values == conv_defaults: - decode_kwargs.update(dict.fromkeys(( - "tile_size_in_pixels", - "tile_overlap_in_pixels", - "tile_size_in_frames", - "tile_overlap_in_frames", - ))) - noise_generator = None - if pipe.is_ltx25 and seed is not None: - noise_generator = torch.Generator(device=rand_device).manual_seed(seed) - if decoder_name == "diffusion_video_vae_decoder": - decode_kwargs["generator"] = noise_generator - return { - "video_decoder_name": decoder_name, - "video_decode_kwargs": decode_kwargs, - "noise_generator": noise_generator, - } class LTX2AudioVideoUnit_AutoDuration(PipelineUnit): def __init__(self): - super().__init__(take_over=True) + super().__init__(take_over=True, onload_model_names=("duration_head",)) @staticmethod def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds): @@ -469,18 +367,12 @@ def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds): return frames def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inputs_nega): - if not inputs_shared.get("auto_duration", False): + if not inputs_shared.get("auto_duration", False) or pipe.duration_head is None: return inputs_shared, inputs_posi, inputs_nega - pipe.load_models_to_device(("duration_head",)) - seconds = float( - pipe.duration_head(inputs_posi["video_context"], inputs_posi["audio_context"]).item() - ) - inputs_shared["num_frames"] = self.seconds_to_num_frames( - seconds, - inputs_shared["frame_rate"], - inputs_shared["auto_duration_min_seconds"], - inputs_shared["auto_duration_max_seconds"], - ) + min_seconds, max_seconds = sorted(min(max(seconds, 1e-6), 20.0) for seconds in (inputs_shared["auto_duration_min_seconds"], inputs_shared["auto_duration_max_seconds"])) + pipe.load_models_to_device(self.onload_model_names) + seconds = float(pipe.duration_head(inputs_posi["video_context"], inputs_posi["audio_context"]).item()) + inputs_shared["num_frames"] = self.seconds_to_num_frames(seconds, inputs_shared["frame_rate"], min_seconds, max_seconds) return inputs_shared, inputs_posi, inputs_nega @@ -581,7 +473,6 @@ def __init__(self): "seed", "rand_device", "frame_rate", - "noise_generator", "generate_video", ), output_params=( @@ -596,7 +487,6 @@ def __init__(self): "video_ancestral_noise_transform", "audio_ancestral_noise_shape", "audio_ancestral_noise_transform", - "noise_generator", ), ) @@ -622,7 +512,6 @@ def process_stage( seed, rand_device, frame_rate=24.0, - noise_generator=None, generate_video=True, ): # The unit runner passes None for params missing from inputs_shared (e.g. in training). @@ -648,7 +537,6 @@ def process_stage( seed=seed, rand_device=rand_device, rand_torch_dtype=noise_dtype, - generator=noise_generator, ) if pipe.is_ltx25: video_noise = pipe.video_patchifier.unpatchify_video( @@ -677,7 +565,6 @@ def process_stage( seed=seed, rand_device=rand_device, rand_torch_dtype=noise_dtype, - generator=noise_generator, ) if pipe.is_ltx25: audio_noise = pipe.audio_patchifier.unpatchify_audio( @@ -723,7 +610,6 @@ def process_stage( "video_ancestral_noise_transform": video_ancestral_noise_transform, "audio_ancestral_noise_shape": audio_ancestral_noise_shape, "audio_ancestral_noise_transform": audio_ancestral_noise_transform, - "noise_generator": noise_generator, } def process( @@ -735,7 +621,6 @@ def process( seed, rand_device, frame_rate=24.0, - noise_generator=None, generate_video=True, ): return self.process_stage( @@ -746,7 +631,6 @@ def process( seed, rand_device, frame_rate, - noise_generator, generate_video, ) diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md index 97c709b83..cc12e3d22 100644 --- a/docs/en/Model_Details/LTX-2.5.md +++ b/docs/en/Model_Details/LTX-2.5.md @@ -90,7 +90,7 @@ For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#in * `auto_duration`: predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. * `auto_duration_min_seconds` / `auto_duration_max_seconds`: lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. * `generate_video`: whether to generate video. Defaults to `True`. Set it to `False` to generate audio only (T2A); the video VAE and latent upsampler are then not required. -* `use_diffusion_vae`: video decoder selection. `None` (default) prefers the loaded ConvVAE convolutional decoder and falls back to the DiffVAE diffusion decoder when it is absent; `True`/`False` force DiffVAE/ConvVAE respectively (ConvVAE requires `ltx-2.5-video-vae-conv-bf16.safetensors`). +* The video decoder is selected by the loaded components: the ConvVAE convolutional decoder is used when `ltx-2.5-video-vae-conv-bf16.safetensors` is loaded, otherwise the DiffVAE diffusion decoder. * Default negative prompt: `pipe.default_negative_prompt["LTX-2.5"]` prefixes the LTX-2/2.3 list with the 2.5-specific tags (`has_subtitles`, `has_blurbox`, `transition from black`, `transition to black`, `speech_ending_short`); all example scripts use this key. * `input_images` / `input_images_indexes`: keyframe images and their frame indexes. A single first frame gives image-to-video; first and last (or more) frames give keyframe interpolation. * `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: audio-to-video (A2V) and audio region retake. diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md index 2f5f204e2..e374af765 100644 --- a/docs/zh/Model_Details/LTX-2.5.md +++ b/docs/zh/Model_Details/LTX-2.5.md @@ -90,7 +90,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 * `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 * `generate_video`: 是否生成视频，默认为 `True`。设置为 `False` 时只生成音频（T2A），此时无需加载视频 VAE 与 latent upsampler。 -* `use_diffusion_vae`：视频解码器选择。`None`（默认）优先使用已加载的 ConvVAE 卷积解码器，未加载时回退到 DiffVAE 扩散解码器；`True`/`False` 分别强制指定 DiffVAE/ConvVAE（ConvVAE 需加载 `ltx-2.5-video-vae-conv-bf16.safetensors`）。 +* 视频解码器按已加载组件自动选择：加载了 `ltx-2.5-video-vae-conv-bf16.safetensors` 就用 ConvVAE 卷积解码器，否则使用 DiffVAE 扩散解码器。 * 默认负向提示词：`pipe.default_negative_prompt["LTX-2.5"]` 在 LTX-2/2.3 的列表之前增加了 2.5 专有标签（`has_subtitles`、`has_blurbox`、`transition from black`、`transition to black`、`speech_ending_short`），示例脚本均使用该键。 * `input_images` / `input_images_indexes`: 关键帧图像及其帧索引。传入首帧即为图生视频，传入首尾（或多帧）即为关键帧插值。 * `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: 音频驱动视频（A2V）与音频区域重生成。 From 8ad3d22ed2190e5f69561bf2042c25a482ce54b7 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Tue, 8 Sep 2026 16:57:08 +0800 Subject: [PATCH 19/31] Keep NoiseInitializer in upstream shape and rename generate_video to audio_only NoiseInitializer goes back to the upstream process_stage/process structure; the only additions are the generate-video guard and a build_video_keyframes_mask helper that returns the LTX-2.5 first-frame marker (None for other versions), so the mask is rebuilt at every stage resolution automatically. The token-layout noise, unpatchify bridges and ancestral transforms are gone: noise is drawn in the 5D latent layout like LTX-2.3. The generate_video flag is renamed to audio_only with inverted polarity, which removes the None-normalization trick: the unit runner passes None for missing keys, and a falsy value now correctly means "generate video" during training caching stages. --- diffsynth/pipelines/ltx2_audio_video.py | 175 +++--------------- docs/en/Model_Details/LTX-2.5.md | 4 +- docs/zh/Model_Details/LTX-2.5.md | 4 +- examples/ltx2/model_inference/LTX-2.5-T2A.py | 2 +- .../model_inference_low_vram/LTX-2.5-T2A.py | 2 +- 5 files changed, 27 insertions(+), 160 deletions(-) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index af47fd452..b390de81e 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -8,7 +8,6 @@ from PIL import Image from tqdm import tqdm from typing import Optional -from functools import partial from ..core.device.npu_compatible_device import get_device_type from ..diffusion import FlowMatchScheduler @@ -191,8 +190,6 @@ def denoise_stage(self, inputs_shared, inputs_posi, inputs_nega, units, cfg_scal noise_pred=noise_pred_video, inpaint_mask=inputs_shared.get("denoise_mask_video", None), input_latents=inputs_shared.get("input_latents_video", None), - ancestral_noise_shape=inputs_shared.get("video_ancestral_noise_shape"), - ancestral_noise_transform=inputs_shared.get("video_ancestral_noise_transform"), ) inputs_shared["audio_latents"] = self.step( self.scheduler, @@ -201,8 +198,6 @@ def denoise_stage(self, inputs_shared, inputs_posi, inputs_nega, units, cfg_scal noise_pred=noise_pred_audio, inpaint_mask=inputs_shared.get("denoise_mask_audio", None), input_latents=inputs_shared.get("input_latents_audio", None), - ancestral_noise_shape=inputs_shared.get("audio_ancestral_noise_shape"), - ancestral_noise_transform=inputs_shared.get("audio_ancestral_noise_transform"), ) return inputs_shared, inputs_posi, inputs_nega @@ -238,7 +233,7 @@ def __call__( auto_duration: bool = False, auto_duration_min_seconds: float = 1.0, auto_duration_max_seconds: float = 20.0, - generate_video: bool = True, + audio_only: bool = False, # Classifier-free guidance cfg_scale: float = 3.0, # Scheduler @@ -276,7 +271,7 @@ def __call__( "auto_duration": auto_duration, "auto_duration_min_seconds": auto_duration_min_seconds, "auto_duration_max_seconds": auto_duration_max_seconds, - "generate_video": generate_video, + "audio_only": audio_only, "cfg_scale": cfg_scale, "tiled": tiled, "tile_size_in_pixels": tile_size_in_pixels, "tile_overlap_in_pixels": tile_overlap_in_pixels, "tile_size_in_frames": tile_size_in_frames, "tile_overlap_in_frames": tile_overlap_in_frames, @@ -300,7 +295,7 @@ def __call__( ) # Decode video = None - if inputs_shared.get("generate_video", True): + if not inputs_shared.get("audio_only", False): if self.video_vae_decoder is None: raise ValueError("No video decoder component is loaded.") self.load_models_to_device(["video_vae_decoder"]) @@ -466,138 +461,33 @@ def process(self, pipe: LTX2AudioVideoPipeline, prompt: str): class LTX2AudioVideoUnit_NoiseInitializer(PipelineUnit): def __init__(self): super().__init__( - input_params=( - "height", - "width", - "num_frames", - "seed", - "rand_device", - "frame_rate", - "generate_video", - ), - output_params=( - "video_noise", - "audio_noise", - "video_positions", - "audio_positions", - "video_latent_shape", - "audio_latent_shape", - "video_keyframes_mask", - "video_ancestral_noise_shape", - "video_ancestral_noise_transform", - "audio_ancestral_noise_shape", - "audio_ancestral_noise_transform", - ), - ) - - @staticmethod - def unpatchify_video_noise(noise, patchifier, latent_shape): - return patchifier.unpatchify_video( - noise, - latent_shape.frames, - latent_shape.height, - latent_shape.width, + input_params=("height", "width", "num_frames", "seed", "rand_device", "frame_rate", "audio_only"), + output_params=("video_noise", "audio_noise", "video_positions", "audio_positions", "video_latent_shape", "audio_latent_shape", "video_keyframes_mask") ) - @staticmethod - def unpatchify_audio_noise(noise, patchifier, latent_shape): - return patchifier.unpatchify_audio(noise, latent_shape.channels, latent_shape.mel_bins) + def build_video_keyframes_mask(self, pipe, video_latent_shape): + if not pipe.is_ltx25: + return None + video_keyframes_mask = torch.zeros(video_latent_shape.batch, 1, video_latent_shape.frames, video_latent_shape.height, video_latent_shape.width, dtype=torch.float32, device=pipe.device) + video_keyframes_mask[:, :, 0] = 1.0 + return video_keyframes_mask - def process_stage( - self, - pipe: LTX2AudioVideoPipeline, - height, - width, - num_frames, - seed, - rand_device, - frame_rate=24.0, - generate_video=True, - ): - # The unit runner passes None for params missing from inputs_shared (e.g. in training). - generate_video = generate_video is not False + def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0, audio_only=False): video_pixel_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) video_latent_shape = VideoLatentShape.from_pixel_shape(shape=video_pixel_shape, latent_channels=128) - noise_dtype = pipe.torch_dtype if pipe.is_ltx25 else torch.float32 - video_noise = None - video_positions = None - video_keyframes_mask = None - video_ancestral_noise_shape = None - video_ancestral_noise_transform = None - if generate_video: - video_noise_shape = video_latent_shape.to_torch_shape() - if pipe.is_ltx25: - video_noise_shape = ( - video_latent_shape.batch, - video_latent_shape.frames * video_latent_shape.height * video_latent_shape.width, - video_latent_shape.channels, - ) - video_noise = pipe.generate_noise( - video_noise_shape, - seed=seed, - rand_device=rand_device, - rand_torch_dtype=noise_dtype, - ) - if pipe.is_ltx25: - video_noise = pipe.video_patchifier.unpatchify_video( - video_noise, - video_latent_shape.frames, - video_latent_shape.height, - video_latent_shape.width, - ) + video_noise = video_positions = video_keyframes_mask = None + if not audio_only: + video_noise = pipe.generate_noise(video_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=video_latent_shape, device=pipe.device) video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, True).float() video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate - if not pipe.is_ltx25: - video_positions = video_positions.to(pipe.torch_dtype) + video_positions = video_positions.to(pipe.torch_dtype) + video_keyframes_mask = self.build_video_keyframes_mask(pipe, video_latent_shape) audio_latent_shape = AudioLatentShape.from_video_pixel_shape(video_pixel_shape) - audio_noise_shape = audio_latent_shape.to_torch_shape() - if pipe.is_ltx25: - audio_noise_shape = ( - audio_latent_shape.batch, - audio_latent_shape.frames, - audio_latent_shape.channels * audio_latent_shape.mel_bins, - ) - audio_noise = pipe.generate_noise( - audio_noise_shape, - seed=seed, - rand_device=rand_device, - rand_torch_dtype=noise_dtype, - ) - if pipe.is_ltx25: - audio_noise = pipe.audio_patchifier.unpatchify_audio( - audio_noise, - audio_latent_shape.channels, - audio_latent_shape.mel_bins, - ) + audio_noise = pipe.generate_noise(audio_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) audio_positions = pipe.audio_patchifier.get_patch_grid_bounds(audio_latent_shape, device=pipe.device) - audio_ancestral_noise_shape = None - audio_ancestral_noise_transform = None - if pipe.is_ltx25 and generate_video: - video_keyframes_mask = torch.zeros( - video_latent_shape.batch, - 1, - video_latent_shape.frames, - video_latent_shape.height, - video_latent_shape.width, - dtype=torch.float32, - device=pipe.device, - ) - video_keyframes_mask[:, :, 0] = 1.0 - video_ancestral_noise_shape = video_noise_shape - video_ancestral_noise_transform = partial( - self.unpatchify_video_noise, - patchifier=pipe.video_patchifier, - latent_shape=video_latent_shape, - ) - audio_ancestral_noise_shape = audio_noise_shape - audio_ancestral_noise_transform = partial( - self.unpatchify_audio_noise, - patchifier=pipe.audio_patchifier, - latent_shape=audio_latent_shape, - ) return { "video_noise": video_noise, "audio_noise": audio_noise, @@ -605,34 +495,11 @@ def process_stage( "audio_positions": audio_positions, "video_latent_shape": video_latent_shape, "audio_latent_shape": audio_latent_shape, - "video_keyframes_mask": video_keyframes_mask, - "video_ancestral_noise_shape": video_ancestral_noise_shape, - "video_ancestral_noise_transform": video_ancestral_noise_transform, - "audio_ancestral_noise_shape": audio_ancestral_noise_shape, - "audio_ancestral_noise_transform": audio_ancestral_noise_transform, + "video_keyframes_mask": video_keyframes_mask } - def process( - self, - pipe: LTX2AudioVideoPipeline, - height, - width, - num_frames, - seed, - rand_device, - frame_rate=24.0, - generate_video=True, - ): - return self.process_stage( - pipe, - height, - width, - num_frames, - seed, - rand_device, - frame_rate, - generate_video, - ) + def process(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0, audio_only=False): + return self.process_stage(pipe, height, width, num_frames, seed, rand_device, frame_rate, audio_only) class LTX2AudioVideoUnit_InputVideoEmbedder(PipelineUnit): diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md index cc12e3d22..f21a1fcf1 100644 --- a/docs/en/Model_Details/LTX-2.5.md +++ b/docs/en/Model_Details/LTX-2.5.md @@ -70,7 +70,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| -|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| |[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| |[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| @@ -89,7 +89,7 @@ For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#in * `auto_duration`: predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. * `auto_duration_min_seconds` / `auto_duration_max_seconds`: lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. -* `generate_video`: whether to generate video. Defaults to `True`. Set it to `False` to generate audio only (T2A); the video VAE and latent upsampler are then not required. +* `audio_only`: whether to generate audio only (T2A). Defaults to `False`; when `True`, the video VAE and latent upsampler are not required. * The video decoder is selected by the loaded components: the ConvVAE convolutional decoder is used when `ltx-2.5-video-vae-conv-bf16.safetensors` is loaded, otherwise the DiffVAE diffusion decoder. * Default negative prompt: `pipe.default_negative_prompt["LTX-2.5"]` prefixes the LTX-2/2.3 list with the 2.5-specific tags (`has_subtitles`, `has_blurbox`, `transition from black`, `transition to black`, `speech_ending_short`); all example scripts use this key. * `input_images` / `input_images_indexes`: keyframe images and their frame indexes. A single first frame gives image-to-video; first and last (or more) frames give keyframe interpolation. diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md index e374af765..39ef46148 100644 --- a/docs/zh/Model_Details/LTX-2.5.md +++ b/docs/zh/Model_Details/LTX-2.5.md @@ -70,7 +70,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| |[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| -|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`generate_video=False`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| |[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| |[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| |[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| @@ -89,7 +89,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 * `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 -* `generate_video`: 是否生成视频，默认为 `True`。设置为 `False` 时只生成音频（T2A），此时无需加载视频 VAE 与 latent upsampler。 +* `audio_only`: 是否只生成音频（T2A），默认为 `False`。设置为 `True` 时无需加载视频 VAE 与 latent upsampler。 * 视频解码器按已加载组件自动选择：加载了 `ltx-2.5-video-vae-conv-bf16.safetensors` 就用 ConvVAE 卷积解码器，否则使用 DiffVAE 扩散解码器。 * 默认负向提示词：`pipe.default_negative_prompt["LTX-2.5"]` 在 LTX-2/2.3 的列表之前增加了 2.5 专有标签（`has_subtitles`、`has_blurbox`、`transition from black`、`transition to black`、`speech_ending_short`），示例脚本均使用该键。 * `input_images` / `input_images_indexes`: 关键帧图像及其帧索引。传入首帧即为图生视频，传入首尾（或多帧）即为关键帧插值。 diff --git a/examples/ltx2/model_inference/LTX-2.5-T2A.py b/examples/ltx2/model_inference/LTX-2.5-T2A.py index 39496d6a9..5a418512f 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2A.py @@ -34,6 +34,6 @@ num_frames=121, frame_rate=24, num_inference_steps=30, - generate_video=False, + audio_only=True, ) save_audio(audio, pipe.audio_vocoder.output_sampling_rate, "ltx2.5_t2a.wav") diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py index 03b78b25a..d8fd3ce04 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py @@ -35,6 +35,6 @@ num_frames=121, frame_rate=24, num_inference_steps=30, - generate_video=False, + audio_only=True, ) save_audio(audio, pipe.audio_vocoder.output_sampling_rate, "ltx2.5_t2a.wav") From c5d46e421da4ae45c6ba8a5b6008288d1cbbee27 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 10:50:50 +0800 Subject: [PATCH 20/31] Give DiffVAE decode explicit tile sizes with a min-tile fallback The decode path accepts tile_size_in_pixels / tile_size_in_frames (overlaps stay halo-derived, since DiffVAE ramps need complementary masks) and validates them; any invalid configuration falls back to the automatic tiling instead of raising. The automatic path now returns the minimum legal tile (spatial floor 512 px, temporal 80 frames), which keeps the decode peak around 7 GB at 1024x1536x121 instead of the recommender's spare-VRAM-hungry choice (~68 GB measured). Inference examples pass tile_size_in_frames=80 and carry a comment pointing at the conv vae decoder as the low-VRAM alternative. --- diffsynth/models/ltx25_diffusion_video_vae.py | 85 ++++++++++--------- .../model_inference/LTX-2.5-A2V-TwoStage.py | 3 + .../model_inference/LTX-2.5-I2AV-OneStage.py | 3 + .../model_inference/LTX-2.5-I2AV-TwoStage.py | 4 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 3 + .../LTX-2.5-T2AV-DistilledPipeline.py | 3 + .../LTX-2.5-T2AV-INT8-ConvRot.py | 3 + .../model_inference/LTX-2.5-T2AV-OneStage.py | 3 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 3 + .../model_inference/LTX-2.5-T2AV-TwoStage.py | 3 + .../LTX-2.5-A2V-TwoStage.py | 3 + .../LTX-2.5-I2AV-OneStage.py | 3 + .../LTX-2.5-I2AV-TwoStage.py | 4 + .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 3 + .../LTX-2.5-T2AV-DistilledPipeline.py | 3 + .../LTX-2.5-T2AV-INT8-ConvRot.py | 3 + .../LTX-2.5-T2AV-OneStage.py | 3 + .../LTX-2.5-T2AV-TwoStage-Retake.py | 3 + .../LTX-2.5-T2AV-TwoStage.py | 3 + 19 files changed, 100 insertions(+), 41 deletions(-) diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index 0cc33f503..44d5fade2 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -4,6 +4,8 @@ import itertools import logging import math + +from tqdm import tqdm from collections.abc import Iterator, Sequence from dataclasses import dataclass, replace from enum import Enum @@ -4313,7 +4315,7 @@ def _decode_temporal_group_isolated_with_keyframes( # noqa: PLR0913 compute_dtype = feat_s4.dtype up3_stride = tuple(self.upsamples[3].stride) - for tile_index, tile in enumerate(tiles): + for tile_index, tile in tqdm(enumerate(tiles), total=len(tiles), desc="DiffVAE decode", leave=False): feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( feat_s4, tile, content_frames=content_s4_frames ) @@ -4667,7 +4669,7 @@ def _decode_temporal_group_isolated( randn_device = generator.device if generator is not None else feat_s4.device up3_stride = tuple(self.upsamples[3].stride) - for tile in tiles: + for tile in tqdm(tiles, total=len(tiles), desc="DiffVAE decode", leave=False): feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( feat_s4, tile, content_frames=content_s4_frames ) @@ -4947,41 +4949,45 @@ def auto_tiling_config(self, latent, keyframes=None): .upscale(self.video_downscale_factors) ._replace(channels=self.out_channels) ) - device = latent.device - if device.type == "cuda": - # Cached allocator blocks from a previous decode would otherwise make the - # free-memory query report a budget of zero for back-to-back decodes. - torch.cuda.empty_cache() - free_bytes = torch.cuda.mem_get_info(device.index)[0] - else: - free_bytes = 0 - if free_bytes <= 0: - return None - - # Budget estimate must not read weight dtype/device: parameters may be meta - # or disk-offloaded here, so assume bf16 storage for the footprint estimate. - model_bytes = sum(parameter.numel() for parameter in self.parameters()) * 2 upsample_strides = [tuple(upsample.stride) for upsample in self.upsamples] - element_size = accumulator_element_size(latent.dtype) - return recommended_decode_tiling_config( - tile_halos=self.tile_halos, - pixel_scale=stage4_to_pixel_scale_factors(upsample_strides[3], self.patch_size), - min_tile_size_s4=self.tile_min_sizes, - patch_size=self.patch_size, - height=pixel_shape.height, - width=pixel_shape.width, - num_frames=pixel_shape.frames, - mode=DiffVAEMode.CHUNKED_EAGER, - free_bytes=free_bytes, - stage5_channels=self.stage_channels[-1], - stage4_channels=self.stage_channels[3], - upsample_strides=upsample_strides, - model_bytes=model_bytes, - element_size=element_size, - natten_trailing_pad_latent_frames=self._natten_trailing_pad_latent_frames, - keyframes=keyframes is not None, + pixel_scale = stage4_to_pixel_scale_factors(upsample_strides[3], self.patch_size) + overlap_t, overlap_hw = recommended_pixel_overlaps(self.tile_halos, pixel_scale) + ft, fh, fw = pixel_scale.time, pixel_scale.height, pixel_scale.width + step_t = math.lcm(ft, VIDEO_SCALE_FACTORS.time) + step_h = math.lcm(fh, VIDEO_SCALE_FACTORS.height) + step_w = math.lcm(fw, VIDEO_SCALE_FACTORS.width) + min_t = _round_up(max(2 * ft, 2 * overlap_t, _round_up(self.tile_min_sizes[0] * ft, ft), 16), step_t) + min_h = _round_up(max(2 * fh, 2 * overlap_hw, _round_up(self.tile_min_sizes[1] * fh, fh), 512), step_h) + min_w = _round_up(max(2 * fw, 2 * overlap_hw, _round_up(self.tile_min_sizes[2] * fw, fw), 512), step_w) + return TileSizeConfig( + frames=DimensionSizeConfig(min_t, overlap_t), + height=DimensionSizeConfig(min_h, overlap_hw), + width=DimensionSizeConfig(min_w, overlap_hw), ) + def _resolve_tiling_config(self, latent, tiled, keyframes=None, tile_size_in_pixels=None, tile_size_in_frames=None): + if not tiled: + return None + if tile_size_in_pixels is not None and tile_size_in_frames is not None: + try: + pixel_scale = stage4_to_pixel_scale_factors([tuple(upsample.stride) for upsample in self.upsamples][3], self.patch_size) + overlap_t, overlap_hw = recommended_pixel_overlaps(self.tile_halos, pixel_scale) + tiling_config = TileSizeConfig( + frames=DimensionSizeConfig(tile_size_in_frames, overlap_t), + height=DimensionSizeConfig(tile_size_in_pixels, overlap_hw), + width=DimensionSizeConfig(tile_size_in_pixels, overlap_hw), + ) + pixel_shape = ( + VideoLatentShape.from_torch_shape(latent.shape) + .upscale(self.video_downscale_factors) + ._replace(channels=self.out_channels) + ) + tiling_config.validate(pixel_scale, pixel_shape) + return tiling_config + except ValueError: + pass + return self.auto_tiling_config(latent, keyframes=keyframes) + def decode( self, latent, @@ -4989,14 +4995,12 @@ def decode( seed=None, rand_device="cpu", keyframes=None, + tile_size_in_pixels=None, + tile_size_in_frames=None, **kwargs, ): - generator = torch.Generator(device=rand_device).manual_seed(seed) if seed is not None else None - tiling_config = None - if tiled: - tiling_config = self.auto_tiling_config(latent, keyframes=keyframes) - if tiling_config is None: - raise ValueError("Automatic DiffVAE tiling requires a CUDA device with queryable free memory.") + generator = torch.Generator(device=rand_device).manual_seed(42 if seed is None else seed + 42) + tiling_config = self._resolve_tiling_config(latent, tiled, keyframes, tile_size_in_pixels, tile_size_in_frames) iterator = ( self._decode_pixels_with_keyframes(latent, keyframes, tiling_config, generator=generator) if keyframes is not None @@ -5006,4 +5010,3 @@ def decode( if not chunks: raise RuntimeError("Diffusion decoder produced no output chunks") return torch.cat(chunks, dim=2) - diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py index b33877148..00b9df838 100644 --- a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -48,6 +50,7 @@ num_frames=num_frames, frame_rate=frame_rate, tiled=True, + tile_size_in_frames=80, use_two_stage_pipeline=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py index 7beddfeab..3a3df1af9 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ], ) @@ -40,6 +42,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, input_images=[first_frame], input_images_indexes=[0], diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index 791b2bbd6..2e3d380f6 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -45,6 +47,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame], @@ -69,6 +72,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame, last_frame], diff --git a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index ef848cc73..f137ed890 100644 --- a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -52,6 +54,7 @@ in_context_videos=[reference_video], in_context_downsample_factor=2, tiled=True, + tile_size_in_frames=80, cfg_scale=1.0, num_inference_steps=8, use_distilled_pipeline=True, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py index 4b7b947a0..7ae96606c 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), @@ -46,6 +48,7 @@ use_distilled_pipeline=True, use_two_stage_pipeline=True, tiled=True, + tile_size_in_frames=80, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py index 767ff8078..cc3f68b09 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -42,6 +44,7 @@ use_distilled_pipeline=True, use_two_stage_pipeline=True, tiled=True, + tile_size_in_frames=80, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py index bef059579..a4a54c224 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ], ) @@ -35,6 +37,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py index b71f227ee..d93ad0b47 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -24,6 +24,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -57,6 +59,7 @@ num_frames=num_frames, frame_rate=frame_rate, tiled=True, + tile_size_in_frames=80, use_two_stage_pipeline=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py index 68d94f734..85a72cdba 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -38,6 +40,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, use_two_stage_pipeline=True, ) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py index 1113fc8d9..b91bd3393 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -49,6 +51,7 @@ num_frames=num_frames, frame_rate=frame_rate, tiled=True, + tile_size_in_frames=80, use_two_stage_pipeline=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index 5ad07c872..40675940f 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ], vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, @@ -41,6 +43,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, input_images=[first_frame], input_images_indexes=[0], diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index ae6e9c75d..8bce88328 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -46,6 +48,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame], @@ -70,6 +73,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame, last_frame], diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 657692b58..f15c09a3f 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -23,6 +23,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -53,6 +55,7 @@ in_context_videos=[reference_video], in_context_downsample_factor=2, tiled=True, + tile_size_in_frames=80, cfg_scale=1.0, num_inference_steps=8, use_distilled_pipeline=True, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index c19c37cb9..b2d692b1c 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), @@ -47,6 +49,7 @@ use_distilled_pipeline=True, use_two_stage_pipeline=True, tiled=True, + tile_size_in_frames=80, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py index 5d75a4072..dec854bcc 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -43,6 +45,7 @@ use_distilled_pipeline=True, use_two_stage_pipeline=True, tiled=True, + tile_size_in_frames=80, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py index 2d811ffa0..6c279699a 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ], vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, @@ -36,6 +38,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py index 3ff7c6d9f..15472cfb8 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -24,6 +24,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -58,6 +60,7 @@ num_frames=num_frames, frame_rate=frame_rate, tiled=True, + tile_size_in_frames=80, use_two_stage_pipeline=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py index 025a1189b..b884791d7 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -21,6 +21,8 @@ ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), ], @@ -39,6 +41,7 @@ width=width, num_frames=num_frames, tiled=True, + tile_size_in_frames=80, cfg_scale=3.0, use_two_stage_pipeline=True, ) From 074758c09dd829e17f51fea0d6d1781b5bf9ea35 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 16:30:21 +0800 Subject: [PATCH 21/31] Add AncestralFlowMatchScheduler and revert non-essential pipeline deviations - Move the rectified-flow ancestral Euler step into AncestralFlowMatchScheduler (flow_match.py base classes stay byte-identical to upstream); installed in __call__ for LTX-2.5 distilled runs, stage 2 restores the base scheduler. - Revert base_pipeline.py to upstream: drop the dead generate_noise generator parameter and handle disabled-modality predictions with a 0 placeholder at the model_fn_ltx2 return site instead of a CFG None guard. - Remove LTX-2.5 special-casing from denoise_stage (timestep dtype/sigmas switch, cfg_scale read, timestep_scale), the PipelineChecker raises, and cosmetic reformats; denoise_stage and PipelineChecker now match upstream exactly. - Point LTX25TextEncoder.forward at the inner hidden-states pass so both text encoder generations share one call site in the prompt embedder unit. - Run LTX-2.5 one-stage examples at half resolution, matching LTX-2/2.3. - Ignore *.mp4 and *.wav test outputs. --- .gitignore | 2 + diffsynth/diffusion/__init__.py | 2 +- diffsynth/diffusion/base_pipeline.py | 15 +- diffsynth/diffusion/flow_match.py | 89 ++++------- diffsynth/models/ltx25_diffusion_video_vae.py | 4 +- diffsynth/models/ltx25_text_encoder.py | 4 +- diffsynth/pipelines/ltx2_audio_video.py | 143 +++++------------- .../model_inference/LTX-2.5-I2AV-OneStage.py | 2 +- .../LTX-2.5-T2AV-DistilledPipeline.py | 2 +- .../model_inference/LTX-2.5-T2AV-OneStage.py | 2 +- .../model_inference/LTX-2.5-T2AV-TwoStage.py | 2 +- .../LTX-2.5-I2AV-OneStage.py | 2 +- .../LTX-2.5-T2AV-OneStage.py | 2 +- 13 files changed, 81 insertions(+), 190 deletions(-) diff --git a/.gitignore b/.gitignore index dfe725a05..80c33d33b 100644 --- a/.gitignore +++ b/.gitignore @@ -19,6 +19,8 @@ *.mv log*.txt AGENTS.md +*.mp4 +*.wav # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/diffsynth/diffusion/__init__.py b/diffsynth/diffusion/__init__.py index d285482de..f5d9ea194 100644 --- a/diffsynth/diffusion/__init__.py +++ b/diffsynth/diffusion/__init__.py @@ -1,4 +1,4 @@ -from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler +from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler, AncestralFlowMatchScheduler from .training_module import DiffusionTrainingModule from .logger import ModelLogger from .runner import launch_training_task, launch_data_process_task diff --git a/diffsynth/diffusion/base_pipeline.py b/diffsynth/diffusion/base_pipeline.py index 1a2ee13e5..5af48a110 100644 --- a/diffsynth/diffusion/base_pipeline.py +++ b/diffsynth/diffusion/base_pipeline.py @@ -180,10 +180,9 @@ def load_models_to_device(self, model_names): module.onload() - def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None, generator=None): + def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None): # Initialize Gaussian noise - if generator is None: - generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) + generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype) noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) return noise @@ -224,12 +223,7 @@ def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, if inpaint_mask is not None: noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents) noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask) - scheduler_kwargs = {} - if "ancestral_noise_shape" in kwargs: - scheduler_kwargs["ancestral_noise_shape"] = kwargs["ancestral_noise_shape"] - if "ancestral_noise_transform" in kwargs: - scheduler_kwargs["ancestral_noise_transform"] = kwargs["ancestral_noise_transform"] - latents_next = scheduler.step(noise_pred, timestep, latents, **scheduler_kwargs) + latents_next = scheduler.step(noise_pred, timestep, latents) return latents_next @@ -352,9 +346,8 @@ def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, i if isinstance(noise_pred_posi, tuple): # Separately handling different output types of latents, eg. video and audio latents. - # Disabled modalities return None and stay None under CFG. noise_pred = tuple( - None if n_posi is None or n_nega is None else n_nega + cfg_scale * (n_posi - n_nega) + n_nega + cfg_scale * (n_posi - n_nega) for n_posi, n_nega in zip(noise_pred_posi, noise_pred_nega) ) else: diff --git a/diffsynth/diffusion/flow_match.py b/diffsynth/diffusion/flow_match.py index ea0db6311..8e0f3664e 100644 --- a/diffsynth/diffusion/flow_match.py +++ b/diffsynth/diffusion/flow_match.py @@ -26,32 +26,6 @@ def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "Z "SenseNova-U1": FlowMatchScheduler.set_timesteps_sensenova_u1, }.get(template, FlowMatchScheduler.set_timesteps_flux) self.num_train_timesteps = 1000 - self.step_mode = "euler" - self.ancestral_eta = 1.0 - self.ancestral_s_noise = 1.0 - self.ancestral_generator = None - self.roundtrip_denoised = False - - def set_step_mode( - self, - mode="euler", - eta=1.0, - s_noise=1.0, - noise_seed=None, - device="cpu", - roundtrip_denoised=False, - ): - if mode not in ("euler", "euler_ancestral"): - raise ValueError(f"Unsupported flow-matching step mode: {mode}") - self.step_mode = mode - self.ancestral_eta = eta - self.ancestral_s_noise = s_noise - self.ancestral_generator = None - self.roundtrip_denoised = roundtrip_denoised - if mode == "euler_ancestral": - if noise_seed is None: - raise ValueError("noise_seed is required for ancestral Euler sampling.") - self.ancestral_generator = torch.Generator(device=device).manual_seed(noise_seed) @staticmethod def set_timesteps_flux(num_inference_steps=100, denoising_strength=1.0, shift=None): @@ -394,7 +368,6 @@ def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, trainin denoising_strength=denoising_strength, **kwargs, ) - self.set_step_mode("euler") if training: self.set_training_weight() self.training = True @@ -407,42 +380,11 @@ def step(self, model_output, timestep, sample, to_final=False, **kwargs): timestep_id = torch.argmin((self.timesteps - timestep).abs()) sigma = self.sigmas[timestep_id] if to_final or timestep_id + 1 >= len(self.timesteps): - sigma_ = torch.zeros_like(sigma) + sigma_ = 0 else: sigma_ = self.sigmas[timestep_id + 1] - denoised = sample.float() - model_output.float() * sigma.float() - if self.roundtrip_denoised: - denoised = denoised.to(sample.dtype).float() - if self.step_mode == "euler": - if not self.roundtrip_denoised: - return sample + model_output * (sigma_ - sigma) - velocity = ((sample.float() - denoised) / sigma.float()).to(sample.dtype) - return (sample.float() + velocity.float() * (sigma_ - sigma).float()).to(sample.dtype) - - if sigma_ == 0: - return denoised.to(sample.dtype) - downstep_ratio = 1.0 + (sigma_ / sigma - 1.0) * self.ancestral_eta - sigma_down = sigma_ * downstep_ratio - sigma_down_ratio = sigma_down / sigma - prev_sample = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised - alpha_next = 1.0 - sigma_ - alpha_down = 1.0 - sigma_down - renoise_coeff = ( - sigma_ ** 2 - sigma_down ** 2 * alpha_next ** 2 / alpha_down ** 2 - ).clamp(min=0).sqrt() - noise_shape = kwargs.get("ancestral_noise_shape", sample.shape) - noise = torch.randn( - noise_shape, - generator=self.ancestral_generator, - dtype=sample.dtype, - device=sample.device, - ) - noise_transform = kwargs.get("ancestral_noise_transform") - if noise_transform is not None: - noise = noise_transform(noise) - prev_sample = alpha_next / alpha_down * prev_sample - prev_sample = prev_sample + noise.float() * self.ancestral_s_noise * renoise_coeff - return prev_sample.to(sample.dtype) + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample def return_to_timestep(self, timestep, sample, sample_stablized): if isinstance(timestep, torch.Tensor): @@ -508,3 +450,28 @@ def step(self, model_output, timestep, sample): noise = self.clip_noise(torch.randn(denoised.shape, device=denoised.device, dtype=denoised.dtype)) sample = sigma_ * noise * self.noise_scale_schedule[timestep_id] + (1.0 - sigma_) * denoised return sample + + +class AncestralFlowMatchScheduler(FlowMatchScheduler): + + def __init__(self, template="LTX-2", eta=1.0, s_noise=1.0, noise_seed=0, rand_device="cpu"): + super().__init__(template) + self.eta = eta + self.s_noise = s_noise + self.generator = torch.Generator(device=rand_device).manual_seed(noise_seed) + + def step(self, model_output, timestep, sample, **kwargs): + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + sigma_ = self.sigmas[timestep_id + 1] if timestep_id + 1 < len(self.timesteps) else torch.zeros_like(sigma) + denoised = sample.float() - model_output.float() * sigma.float() + if sigma_ == 0: + return denoised.to(sample.dtype) + sigma_down = sigma_ * (1.0 + (sigma_ / sigma - 1.0) * self.eta) + ratio = sigma_down / sigma + prev_sample = ratio * sample.float() + (1.0 - ratio) * denoised + alpha_next, alpha_down = 1.0 - sigma_, 1.0 - sigma_down + renoise_coeff = (sigma_ ** 2 - sigma_down ** 2 * alpha_next ** 2 / alpha_down ** 2).clamp(min=0).sqrt() + noise = torch.randn(sample.shape, generator=self.generator, dtype=sample.dtype, device=self.generator.device).to(sample.device) + prev_sample = alpha_next / alpha_down * prev_sample + noise.float() * self.s_noise * renoise_coeff + return prev_sample.to(sample.dtype) diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index 44d5fade2..9d77da8a7 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -4315,7 +4315,7 @@ def _decode_temporal_group_isolated_with_keyframes( # noqa: PLR0913 compute_dtype = feat_s4.dtype up3_stride = tuple(self.upsamples[3].stride) - for tile_index, tile in tqdm(enumerate(tiles), total=len(tiles), desc="DiffVAE decode", leave=False): + for tile_index, tile in tqdm(enumerate(tiles), total=len(tiles), desc="DiffVAE decode"): feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( feat_s4, tile, content_frames=content_s4_frames ) @@ -4669,7 +4669,7 @@ def _decode_temporal_group_isolated( randn_device = generator.device if generator is not None else feat_s4.device up3_stride = tuple(self.upsamples[3].stride) - for tile in tqdm(tiles, total=len(tiles), desc="DiffVAE decode", leave=False): + for tile in tqdm(tiles, total=len(tiles), desc="DiffVAE decode"): feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( feat_s4, tile, content_frames=content_s4_frames ) diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index 0e438c911..93fcfa9e7 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -163,8 +163,8 @@ def reset_non_persistent_buffers(self): delattr(embed_tokens, "embed_scale") embed_tokens.register_buffer("embed_scale", embed_scale, persistent=False) - def forward(self, *args, **kwargs): - return self.model(*args, **kwargs) + def forward(self, input_ids=None, attention_mask=None, output_hidden_states=False, **kwargs): + return self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=output_hidden_states, **kwargs) class LTX25GemmaTokenizer: diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index b390de81e..2777e1375 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -10,7 +10,7 @@ from typing import Optional from ..core.device.npu_compatible_device import get_device_type -from ..diffusion import FlowMatchScheduler +from ..diffusion import AncestralFlowMatchScheduler, FlowMatchScheduler from ..core import ModelConfig from ..diffusion.base_pipeline import BasePipeline, PipelineUnit @@ -60,7 +60,6 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): LTX2AudioVideoUnit_PromptEmbedder(), LTX2AudioVideoUnit_AutoDuration(), LTX2AudioVideoUnit_ShapeChecker(), - LTX25AudioVideoUnit_SetScheduleStage1Ancestral(), LTX2AudioVideoUnit_NoiseInitializer(), LTX2AudioVideoUnit_VideoRetakeEmbedder(), LTX2AudioVideoUnit_AudioRetakeEmbedder(), @@ -170,35 +169,19 @@ def denoise_stage(self, inputs_shared, inputs_posi, inputs_nega, units, cfg_scal return inputs_shared, inputs_posi, inputs_nega for unit in units: inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) - cfg_scale = inputs_shared.get("cfg_scale", cfg_scale) self.load_models_to_device(self.in_iteration_models) models = {name: getattr(self, name) for name in self.in_iteration_models} - timestep_dtype = torch.float32 if self.is_ltx25 else self.torch_dtype for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): - if self.is_ltx25: - timestep = self.scheduler.sigmas[progress_id] - timestep = timestep.unsqueeze(0).to(dtype=timestep_dtype, device=self.device) + timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device) noise_pred_video, noise_pred_audio = self.cfg_guided_model_fn( self.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) - if inputs_shared.get("video_latents") is not None and noise_pred_video is not None: - inputs_shared["video_latents"] = self.step( - self.scheduler, - inputs_shared["video_latents"], - progress_id=progress_id, - noise_pred=noise_pred_video, - inpaint_mask=inputs_shared.get("denoise_mask_video", None), - input_latents=inputs_shared.get("input_latents_video", None), - ) - inputs_shared["audio_latents"] = self.step( - self.scheduler, - inputs_shared["audio_latents"], - progress_id=progress_id, - noise_pred=noise_pred_audio, - inpaint_mask=inputs_shared.get("denoise_mask_audio", None), - input_latents=inputs_shared.get("input_latents_audio", None), - ) + if inputs_shared.get("video_latents") is not None: + inputs_shared["video_latents"] = self.step(self.scheduler, inputs_shared["video_latents"], progress_id=progress_id, noise_pred=noise_pred_video, + inpaint_mask=inputs_shared.get("denoise_mask_video", None), input_latents=inputs_shared.get("input_latents_video", None), **inputs_shared) + inputs_shared["audio_latents"] = self.step(self.scheduler, inputs_shared["audio_latents"], progress_id=progress_id, noise_pred=noise_pred_audio, + inpaint_mask=inputs_shared.get("denoise_mask_audio", None), input_latents=inputs_shared.get("input_latents_audio", None), **inputs_shared) return inputs_shared, inputs_posi, inputs_nega @torch.no_grad() @@ -253,11 +236,9 @@ def __call__( progress_bar_cmd=tqdm, ): # Scheduler - self.scheduler.set_timesteps( - num_inference_steps, - denoising_strength=denoising_strength, - special_case="distilled_stage1" if use_distilled_pipeline else None, - ) + if self.is_ltx25 and use_distilled_pipeline: + self.scheduler = AncestralFlowMatchScheduler(noise_seed=seed + 10000, rand_device=rand_device) + self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, special_case="distilled_stage1" if use_distilled_pipeline else None) # Inputs inputs_posi = {"prompt": prompt} inputs_nega = {"negative_prompt": negative_prompt} @@ -268,38 +249,25 @@ def __call__( "in_context_videos": in_context_videos, "in_context_downsample_factor": in_context_downsample_factor, "seed": seed, "rand_device": rand_device, "height": height, "width": width, "num_frames": num_frames, "frame_rate": frame_rate, - "auto_duration": auto_duration, - "auto_duration_min_seconds": auto_duration_min_seconds, - "auto_duration_max_seconds": auto_duration_max_seconds, + "auto_duration": auto_duration, "auto_duration_min_seconds": auto_duration_min_seconds, "auto_duration_max_seconds": auto_duration_max_seconds, "audio_only": audio_only, "cfg_scale": cfg_scale, "tiled": tiled, "tile_size_in_pixels": tile_size_in_pixels, "tile_overlap_in_pixels": tile_overlap_in_pixels, "tile_size_in_frames": tile_size_in_frames, "tile_overlap_in_frames": tile_overlap_in_frames, "use_two_stage_pipeline": use_two_stage_pipeline, "use_distilled_pipeline": use_distilled_pipeline, "clear_lora_before_state_two": clear_lora_before_state_two, "stage2_spatial_upsample_factor": stage2_spatial_upsample_factor, "video_patchifier": self.video_patchifier, "audio_patchifier": self.audio_patchifier, - "timestep_scale": 1.0 if self.is_ltx25 else 1000.0, } # Stage 1 - inputs_shared, inputs_posi, inputs_nega = self.denoise_stage( - inputs_shared, inputs_posi, inputs_nega, self.units, cfg_scale, progress_bar_cmd - ) + inputs_shared, inputs_posi, inputs_nega = self.denoise_stage(inputs_shared, inputs_posi, inputs_nega, self.units, cfg_scale, progress_bar_cmd) # Stage 2 - inputs_shared, inputs_posi, inputs_nega = self.denoise_stage( - inputs_shared, - inputs_posi, - inputs_nega, - self.stage2_units, - 1.0, - progress_bar_cmd, - not inputs_shared["use_two_stage_pipeline"], - ) + inputs_shared, inputs_posi, inputs_nega = self.denoise_stage(inputs_shared, inputs_posi, inputs_nega, self.stage2_units, 1.0, progress_bar_cmd, not inputs_shared["use_two_stage_pipeline"]) # Decode video = None if not inputs_shared.get("audio_only", False): - if self.video_vae_decoder is None: - raise ValueError("No video decoder component is loaded.") self.load_models_to_device(["video_vae_decoder"]) - video = self.video_vae_decoder.decode(inputs_shared["video_latents"], tiled=tiled, tile_size_in_pixels=tile_size_in_pixels, tile_overlap_in_pixels=tile_overlap_in_pixels, tile_size_in_frames=tile_size_in_frames, tile_overlap_in_frames=tile_overlap_in_frames, seed=None if seed is None else seed + 42, rand_device=rand_device) + video = self.video_vae_decoder.decode( + inputs_shared["video_latents"], tiled=tiled, tile_size_in_pixels=tile_size_in_pixels, tile_overlap_in_pixels=tile_overlap_in_pixels, + tile_size_in_frames=tile_size_in_frames, tile_overlap_in_frames=tile_overlap_in_frames, seed=seed, rand_device=rand_device) video = self.vae_output_to_video(video) retake_audio = inputs_shared.get("retake_audio") denoise_mask_audio = inputs_shared.get("denoise_mask_audio") @@ -323,24 +291,22 @@ def __call__( class LTX2AudioVideoUnit_PipelineChecker(PipelineUnit): def __init__(self): - super().__init__(take_over=True) + super().__init__( + take_over=True, + input_params=("use_distilled_pipeline", "use_two_stage_pipeline"), + output_params=("use_two_stage_pipeline", "cfg_scale") + ) def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inputs_nega): - use_distilled_pipeline = inputs_shared.get("use_distilled_pipeline", False) - use_two_stage_pipeline = inputs_shared.get("use_two_stage_pipeline", False) - if use_distilled_pipeline: + if inputs_shared.get("use_distilled_pipeline", False): inputs_shared["cfg_scale"] = 1.0 - print("Distilled pipeline requested, disable CFG by setting cfg_scale to 1.0.") - if pipe.is_ltx25 and use_distilled_pipeline: - if not use_two_stage_pipeline: - raise ValueError("LTX-2.5 distilled inference requires use_two_stage_pipeline=True.") - if inputs_shared.get("seed") is None: - raise ValueError("LTX-2.5 distilled ancestral sampling requires an explicit seed.") - if use_two_stage_pipeline: - if not use_distilled_pipeline: + print(f"Distilled pipeline requested, disable CFG by setting cfg_scale to 1.0.") + if inputs_shared.get("use_two_stage_pipeline", False): + # distill pipeline also uses two-stage, but it does not needs lora + if not inputs_shared.get("use_distilled_pipeline", False): if not (hasattr(pipe, "stage2_lora_config") and pipe.stage2_lora_config is not None): raise ValueError("Two-stage pipeline requested, but stage2_lora_config is not set in the pipeline.") - if pipe.upsampler is None: + if not (hasattr(pipe, "upsampler") and pipe.upsampler is not None): raise ValueError("Two-stage pipeline requested, but upsampler model is not loaded in the pipeline.") return inputs_shared, inputs_posi, inputs_nega @@ -371,27 +337,6 @@ def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inpu return inputs_shared, inputs_posi, inputs_nega -class LTX25AudioVideoUnit_SetScheduleStage1Ancestral(PipelineUnit): - def __init__(self): - super().__init__(input_params=("use_distilled_pipeline", "seed")) - - def process(self, pipe: LTX2AudioVideoPipeline, use_distilled_pipeline, seed): - if not pipe.is_ltx25: - return {} - if use_distilled_pipeline: - pipe.scheduler.set_step_mode( - "euler_ancestral", - eta=1.0, - s_noise=1.0, - noise_seed=seed + 10000, - device=pipe.device, - roundtrip_denoised=True, - ) - else: - pipe.scheduler.set_step_mode("euler", roundtrip_denoised=True) - return {} - - class LTX2AudioVideoUnit_ShapeChecker(PipelineUnit): """ For two-stage pipelines, the resolution must be divisible by 64. @@ -431,20 +376,9 @@ def _preprocess_text( text: str, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: token_pairs = pipe.tokenizer.tokenize_with_weights(text)["gemma"] - input_ids = torch.tensor([[token_id for token_id, _ in token_pairs]], device=pipe.device) - attention_mask = torch.tensor([[weight for _, weight in token_pairs]], device=pipe.device) - if pipe.is_ltx25: - outputs = pipe.text_encoder.model.model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - else: - outputs = pipe.text_encoder( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) + input_ids = torch.tensor([[t[0] for t in token_pairs]], device=pipe.device) + attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=pipe.device) + outputs = pipe.text_encoder(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) return outputs.hidden_states, attention_mask def encode_prompt(self, pipe, text, padding_side="left"): hidden_states, attention_mask = self._preprocess_text(pipe, text) @@ -758,17 +692,16 @@ class LTX2AudioVideoUnit_SetScheduleStage2(PipelineUnit): def __init__(self): super().__init__( input_params=("video_latents", "video_noise", "audio_latents", "audio_noise"), - output_params=("video_latents", "audio_latents", "cfg_scale"), + output_params=("video_latents", "audio_latents"), ) def process(self, pipe: LTX2AudioVideoPipeline, video_latents, video_noise, audio_latents, audio_noise): + pipe.scheduler = FlowMatchScheduler("LTX-2") pipe.scheduler.set_timesteps(special_case="stage2") - pipe.scheduler.set_step_mode("euler", roundtrip_denoised=pipe.is_ltx25) if video_latents is not None and video_noise is not None: video_latents = pipe.scheduler.add_noise(video_latents, video_noise, pipe.scheduler.timesteps[0]) audio_latents = pipe.scheduler.add_noise(audio_latents, audio_noise, pipe.scheduler.timesteps[0]) - # The refinement stage runs without classifier-free guidance. - return {"video_latents": video_latents, "audio_latents": audio_latents, "cfg_scale": 1.0} + return {"video_latents": video_latents, "audio_latents": audio_latents} class LTX2AudioVideoUnit_LatentsUpsampler(PipelineUnit): @@ -815,13 +748,12 @@ def model_fn_ltx2( denoise_mask_audio=None, # LTX-2.5 keyframe class embedding video_keyframes_mask=None, - timestep_scale=1000.0, # Gradient Checkpointing use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, **kwargs, ): - timestep = timestep.float() / timestep_scale + timestep = timestep.float() / 1000. video_timesteps = None if video_latents is not None: @@ -890,9 +822,6 @@ def model_fn_ltx2( use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, ) - if vx is not None: - vx = vx[:, :seq_len_video, ...] - # unpatchify - vx = video_patchifier.unpatchify_video(vx, f, h, w) - ax = audio_patchifier.unpatchify_audio(ax, c_a, mel_bins) if ax is not None else None + vx = video_patchifier.unpatchify_video(vx[:, :seq_len_video, ...], f, h, w) if vx is not None else 0 + ax = audio_patchifier.unpatchify_audio(ax, c_a, mel_bins) if ax is not None else 0 return vx, ax diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py index 3a3df1af9..136ac8609 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -32,7 +32,7 @@ # The example image comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] -height, width, num_frames = 512 * 2, 768 * 2, 121 +height, width, num_frames = 512, 768, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) video, audio = pipe( prompt=prompt, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py index 7ae96606c..0078a7472 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -36,7 +36,7 @@ video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, frame_rate=24, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py index a4a54c224..50754ea0b 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py @@ -28,7 +28,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] -height, width, num_frames = 512 * 2, 768 * 2, 121 +height, width, num_frames = 512, 768, 121 video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py index 85a72cdba..2e510be3a 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -35,7 +35,7 @@ video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index 40675940f..2019bac24 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -33,7 +33,7 @@ # The example image comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] -height, width, num_frames = 512 * 2, 768 * 2, 121 +height, width, num_frames = 512, 768, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) video, audio = pipe( prompt=prompt, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py index 6c279699a..4660966b1 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -29,7 +29,7 @@ ) prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] -height, width, num_frames = 512 * 2, 768 * 2, 121 +height, width, num_frames = 512, 768, 121 video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, From bcbe4bc1cc5ea8807d7cca6ef5830ddc89bc8fec Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 17:37:42 +0800 Subject: [PATCH 22/31] Inherit LTX25TextEncoder from Gemma4UnifiedForConditionalGeneration - Drop the wrapper layout: the class now inherits the transformers model, the config literal lives in __init__ (no deepcopy, no module-level constant), and the forward override is gone since the inherited forward returns hidden states. - Delete reset_non_persistent_buffers: rope values match transformers' own init bit-for-bit and the embed_scale difference cancels in RMSNorm, verified by end-to-end bit-identical encodings with and without it. - Update the state dict converter key prefixes for the inherited layout (model.language_model.*, model.embed_vision.*, model.embed_audio.*, lm_head). --- diffsynth/models/ltx25_text_encoder.py | 249 +++++++----------- .../ltx25_text_encoder.py | 10 +- 2 files changed, 105 insertions(+), 154 deletions(-) diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index 93fcfa9e7..6e6d8e2fa 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -1,10 +1,9 @@ -import copy import math from pathlib import Path from typing import NamedTuple import torch -from transformers import PreTrainedTokenizerFast +from transformers import Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration, PreTrainedTokenizerFast from .ltx2_common import rms_norm from .ltx2_dit import ( @@ -17,154 +16,106 @@ ) -LTX25_GEMMA_CONFIG = { - "architectures": ["Gemma4UnifiedForConditionalGeneration"], - "audio_config": { - "_name_or_path": "", - "architectures": None, - "audio_embed_dim": 640, - "chunk_size_feed_forward": 0, - "dtype": "bfloat16", - "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, - "initializer_range": 0.02, - "is_encoder_decoder": False, - "label2id": {"LABEL_0": 0, "LABEL_1": 1}, - "model_type": "gemma4_unified_audio", - "output_attentions": False, - "output_hidden_states": False, - "problem_type": None, - "return_dict": True, - "rms_norm_eps": 1e-06, - }, - "audio_token_id": 258881, - "boa_token_id": 256000, - "boi_token_id": 255999, - "dtype": "bfloat16", - "eoa_token_index": 258883, - "eoi_token_id": 258882, - "eos_token_id": [1, 106], - "gemma_version": "gemma4-12b-ltx-v1", - "image_token_id": 258880, - "initializer_range": 0.02, - "model_type": "gemma4_unified", - "text_config": { - "attention_bias": False, - "attention_dropout": 0.0, - "attention_k_eq_v": True, - "bos_token_id": 2, - "dtype": "bfloat16", - "enable_moe_block": False, - "eos_token_id": 1, - "final_logit_softcapping": 30.0, - "global_head_dim": 512, - "head_dim": 256, - "hidden_activation": "gelu_pytorch_tanh", - "hidden_size": 3840, - "hidden_size_per_layer_input": 0, - "initializer_range": 0.02, - "intermediate_size": 15360, - "layer_types": (["sliding_attention"] * 5 + ["full_attention"]) * 8, - "max_position_embeddings": 262144, - "model_type": "gemma4_unified_text", - "moe_intermediate_size": None, - "num_attention_heads": 16, - "num_experts": None, - "num_global_key_value_heads": 1, - "num_hidden_layers": 48, - "num_key_value_heads": 8, - "num_kv_shared_layers": 0, - "pad_token_id": 0, - "rms_norm_eps": 1e-06, - "rope_parameters": { - "full_attention": {"partial_rotary_factor": 0.25, "rope_theta": 1000000.0, "rope_type": "proportional"}, - "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, - }, - "sliding_window": 1024, - "tie_word_embeddings": True, - "top_k_experts": None, - "use_bidirectional_attention": "vision", - "use_cache": True, - "use_double_wide_mlp": False, - "vocab_size": 262144, - "vocab_size_per_layer_input": 262144, - }, - "tie_word_embeddings": True, - "transformers_version": "5.10.1", - "video_token_id": 258884, - "vision_config": { - "_name_or_path": "", - "architectures": None, - "chunk_size_feed_forward": 0, - "dtype": "bfloat16", - "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, - "initializer_range": 0.02, - "is_encoder_decoder": False, - "label2id": {"LABEL_0": 0, "LABEL_1": 1}, - "mm_embed_dim": 3840, - "mm_posemb_size": 1120, - "model_type": "gemma4_unified_vision", - "num_soft_tokens": 280, - "output_attentions": False, - "output_hidden_states": False, - "output_proj_dims": 3840, - "patch_size": 16, - "pooling_kernel_size": 3, - "problem_type": None, - "return_dict": True, - "rms_norm_eps": 1e-06, - }, -} - - -class LTX25TextEncoder(torch.nn.Module): +class LTX25TextEncoder(Gemma4UnifiedForConditionalGeneration): def __init__(self): - super().__init__() - from transformers import Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration - - self.config = Gemma4UnifiedConfig(**copy.deepcopy(LTX25_GEMMA_CONFIG)) - self.model = Gemma4UnifiedForConditionalGeneration(self.config) - self.reset_non_persistent_buffers() - - def reset_non_persistent_buffers(self): - from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS - - language_model = self.model.model.language_model - text_config = self.config.text_config - rotary_embedding = language_model.rotary_emb - # VRAM management replaces submodules with wrappers; buffers live on the inner module. - rotary_embedding = getattr(rotary_embedding, "module", rotary_embedding) - for layer_type in dict.fromkeys(text_config.layer_types): - rope_parameters = text_config.rope_parameters[layer_type] - if rope_parameters is None: - continue - rope_type = rope_parameters["rope_type"] - if rope_type == "default": - inv_freq, attention_scaling = rotary_embedding.compute_default_rope_parameters( - text_config, layer_type=layer_type - ) - else: - init_kwargs = {"layer_type": layer_type} - if layer_type == "full_attention" and rope_type == "proportional": - init_kwargs["head_dim_key"] = "global_head_dim" - inv_freq, attention_scaling = ROPE_INIT_FUNCTIONS[rope_type](text_config, **init_kwargs) - for buffer_name, buffer_value in ( - (f"{layer_type}_inv_freq", inv_freq), - (f"{layer_type}_original_inv_freq", inv_freq.clone()), - ): - if hasattr(rotary_embedding, buffer_name): - delattr(rotary_embedding, buffer_name) - rotary_embedding.register_buffer(buffer_name, buffer_value, persistent=False) - setattr(rotary_embedding, f"{layer_type}_attention_scaling", attention_scaling) - - embed_scale = torch.tensor(text_config.hidden_size**0.5, device="cpu") - embed_tokens = language_model.embed_tokens - embed_tokens = getattr(embed_tokens, "module", embed_tokens) - if hasattr(embed_tokens, "embed_scale"): - delattr(embed_tokens, "embed_scale") - embed_tokens.register_buffer("embed_scale", embed_scale, persistent=False) - - def forward(self, input_ids=None, attention_mask=None, output_hidden_states=False, **kwargs): - return self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=output_hidden_states, **kwargs) + config = { + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "audio_config": { + "_name_or_path": "", + "architectures": None, + "audio_embed_dim": 640, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, + "initializer_range": 0.02, + "is_encoder_decoder": False, + "label2id": {"LABEL_0": 0, "LABEL_1": 1}, + "model_type": "gemma4_unified_audio", + "output_attentions": False, + "output_hidden_states": False, + "problem_type": None, + "return_dict": True, + "rms_norm_eps": 1e-06, + }, + "audio_token_id": 258881, + "boa_token_id": 256000, + "boi_token_id": 255999, + "dtype": "bfloat16", + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "eos_token_id": [1, 106], + "gemma_version": "gemma4-12b-ltx-v1", + "image_token_id": 258880, + "initializer_range": 0.02, + "model_type": "gemma4_unified", + "text_config": { + "attention_bias": False, + "attention_dropout": 0.0, + "attention_k_eq_v": True, + "bos_token_id": 2, + "dtype": "bfloat16", + "enable_moe_block": False, + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 3840, + "hidden_size_per_layer_input": 0, + "initializer_range": 0.02, + "intermediate_size": 15360, + "layer_types": (["sliding_attention"] * 5 + ["full_attention"]) * 8, + "max_position_embeddings": 262144, + "model_type": "gemma4_unified_text", + "moe_intermediate_size": None, + "num_attention_heads": 16, + "num_experts": None, + "num_global_key_value_heads": 1, + "num_hidden_layers": 48, + "num_key_value_heads": 8, + "num_kv_shared_layers": 0, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": {"partial_rotary_factor": 0.25, "rope_theta": 1000000.0, "rope_type": "proportional"}, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + }, + "sliding_window": 1024, + "tie_word_embeddings": True, + "top_k_experts": None, + "use_bidirectional_attention": "vision", + "use_cache": True, + "use_double_wide_mlp": False, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144, + }, + "tie_word_embeddings": True, + "transformers_version": "5.10.1", + "video_token_id": 258884, + "vision_config": { + "_name_or_path": "", + "architectures": None, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, + "initializer_range": 0.02, + "is_encoder_decoder": False, + "label2id": {"LABEL_0": 0, "LABEL_1": 1}, + "mm_embed_dim": 3840, + "mm_posemb_size": 1120, + "model_type": "gemma4_unified_vision", + "num_soft_tokens": 280, + "output_attentions": False, + "output_hidden_states": False, + "output_proj_dims": 3840, + "patch_size": 16, + "pooling_kernel_size": 3, + "problem_type": None, + "return_dict": True, + "rms_norm_eps": 1e-06, + }, + } + super().__init__(Gemma4UnifiedConfig(**config)) class LTX25GemmaTokenizer: diff --git a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py index b22ca72e9..e27f0f29d 100644 --- a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py +++ b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py @@ -2,15 +2,15 @@ def LTX25TextEncoderStateDictConverter(state_dict): state_dict_ = {} for name in state_dict: if name.startswith("model."): - new_name = "model.model.language_model." + name.removeprefix("model.") + new_name = "model.language_model." + name.removeprefix("model.") elif name.startswith("vision_model."): - new_name = "model.model.embed_vision." + name.removeprefix("vision_model.") + new_name = "model.embed_vision." + name.removeprefix("vision_model.") elif name.startswith("multi_modal_projector."): - new_name = "model.model.embed_vision.multimodal_embedder." + name.removeprefix("multi_modal_projector.") + new_name = "model.embed_vision.multimodal_embedder." + name.removeprefix("multi_modal_projector.") elif name.startswith("audio_projector."): - new_name = "model.model.embed_audio." + name.removeprefix("audio_projector.") + new_name = "model.embed_audio." + name.removeprefix("audio_projector.") else: continue state_dict_[new_name] = state_dict[name] - state_dict_["model.lm_head.weight"] = state_dict_["model.model.language_model.embed_tokens.weight"] + state_dict_["lm_head.weight"] = state_dict_["model.language_model.embed_tokens.weight"] return state_dict_ From ae9f33f9152e21526b335a5e1ebfa5a7111d3207 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 20:02:28 +0800 Subject: [PATCH 23/31] Mark appended reference tokens as non-keyframes in the keyframes mask Upstream extends the keyframes mask with zeros for given-content conditioning (keyframe_cond.py:86, reference_video_cond.py:102-104: "Reference tokens are never keyframes"); we concatenated ones, which added the learned keyframe embedding to tokens that must not receive it. Verified against the target library with per-channel probes (appended tokens stay unmarked, first latent frame stays marked unconditionally). --- diffsynth/pipelines/ltx2_audio_video.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 2777e1375..083f44cee 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -429,7 +429,7 @@ def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, "audio_positions": audio_positions, "video_latent_shape": video_latent_shape, "audio_latent_shape": audio_latent_shape, - "video_keyframes_mask": video_keyframes_mask + "video_keyframes_mask": video_keyframes_mask, } def process(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0, audio_only=False): @@ -786,14 +786,7 @@ def model_fn_ltx2( video_positions = torch.cat([video_positions, ref_frames_position], dim=2) video_timesteps = torch.cat([video_timesteps, ref_frames_timestep], dim=1) if video_keyframes_mask is not None: - # Target marks appended single-frame guiding latents as keyframe tokens too. - ref_keyframes_mask = torch.ones( - ref_frames_latent.shape[0], - ref_frames_latent.shape[1], - 1, - dtype=video_keyframes_mask.dtype, - device=video_keyframes_mask.device, - ) + ref_keyframes_mask = torch.zeros(ref_frames_latent.shape[0], ref_frames_latent.shape[1], 1, dtype=video_keyframes_mask.dtype, device=video_keyframes_mask.device) video_keyframes_mask = torch.cat([video_keyframes_mask, ref_keyframes_mask], dim=1) if audio_latents is not None: From b2aa1fc378e1f82060636726f84ddc4cfd4195c5 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 20:17:01 +0800 Subject: [PATCH 24/31] Decode frozen retake audio through the audio VAE like upstream The audio_fully_frozen shortcut returned the resampled input waveform and skipped the audio decoder when the retake mask was all-zero; upstream always decodes the (frozen) audio latent (retake.py:326). Restore the unconditional decode path and drop the now-unused resample_waveform import. --- diffsynth/pipelines/ltx2_audio_video.py | 25 +++++-------------------- 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 083f44cee..391bf96f9 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -22,7 +22,7 @@ from ..models.ltx2_common import VideoLatentShape, AudioLatentShape, VideoPixelShape, get_pixel_coords, VIDEO_SCALE_FACTORS from ..models.ltx25_text_encoder import LTX25GemmaTokenizer from ..utils.data.media_io_ltx2 import ltx2_preprocess -from ..utils.data.audio import convert_to_stereo, resample_waveform +from ..utils.data.audio import convert_to_stereo class LTX2AudioVideoPipeline(BasePipeline): @@ -269,23 +269,10 @@ def __call__( inputs_shared["video_latents"], tiled=tiled, tile_size_in_pixels=tile_size_in_pixels, tile_overlap_in_pixels=tile_overlap_in_pixels, tile_size_in_frames=tile_size_in_frames, tile_overlap_in_frames=tile_overlap_in_frames, seed=seed, rand_device=rand_device) video = self.vae_output_to_video(video) - retake_audio = inputs_shared.get("retake_audio") - denoise_mask_audio = inputs_shared.get("denoise_mask_audio") - audio_fully_frozen = ( - retake_audio is not None - and denoise_mask_audio is not None - and float(denoise_mask_audio.abs().max()) == 0.0 - ) - if audio_fully_frozen: - waveform, waveform_sample_rate = retake_audio - decoded_audio = resample_waveform(waveform, waveform_sample_rate, self.audio_vocoder.output_sampling_rate) - num_samples = int(inputs_shared["num_frames"] / inputs_shared["frame_rate"] * self.audio_vocoder.output_sampling_rate) - decoded_audio = self.output_audio_format_check(decoded_audio[..., :num_samples]) - else: - self.load_models_to_device(["audio_vae_decoder", "audio_vocoder"]) - decoded_audio = self.audio_vae_decoder(inputs_shared["audio_latents"]) - decoded_audio = self.audio_vocoder(decoded_audio) - decoded_audio = self.output_audio_format_check(decoded_audio) + self.load_models_to_device(["audio_vae_decoder", "audio_vocoder"]) + decoded_audio = self.audio_vae_decoder(inputs_shared["audio_latents"]) + decoded_audio = self.audio_vocoder(decoded_audio) + decoded_audio = self.output_audio_format_check(decoded_audio) return video, decoded_audio @@ -761,8 +748,6 @@ def model_fn_ltx2( b, c_v, f, h, w = video_latents.shape video_latents = video_patchifier.patchify(video_latents) if video_keyframes_mask is not None: - # Target LTX-2.5 keeps patchified video tokens as a channel-first view. - # Preserve that layout because BF16 GEMM reduction order depends on strides. video_latents = video_latents.transpose(1, 2).contiguous().transpose(1, 2) video_keyframes_mask = video_patchifier.patchify(video_keyframes_mask) seq_len_video = video_latents.shape[1] From 3438abec486f34fcb8fc59263917afd3c6c2d097 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 20:19:58 +0800 Subject: [PATCH 25/31] Revert .gitignore to the upstream version The *.mp4 / *.wav ignore rules were local-only convenience for test outputs and must not ship in the integration branch. --- .gitignore | 2 -- 1 file changed, 2 deletions(-) diff --git a/.gitignore b/.gitignore index 80c33d33b..dfe725a05 100644 --- a/.gitignore +++ b/.gitignore @@ -19,8 +19,6 @@ *.mv log*.txt AGENTS.md -*.mp4 -*.wav # Byte-compiled / optimized / DLL files __pycache__/ From d3932aaefc1aa7291c32ee14ec04653abf374ef6 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Wed, 9 Sep 2026 22:05:38 +0800 Subject: [PATCH 26/31] Merge LTX-2.5 docs into the LTX-2 page and keep position precision per generation - Treat LTX-2.5 as a version update: fold its content into docs/{zh,en}/Model_Details/LTX-2.md, delete the standalone LTX-2.5 pages, point README news/links at the merged page and drop the duplicated doc link entry. - Keep video positions in fp32 for LTX-2.5 and cast them to the model dtype only for LTX-2/2.3; bf16 time coordinates made the 2.5 outputs flicker while fp32 breaks the older generations. - Drop the dead ff_bias plumbing from LTX2TextEncoder and the redundant prose comments in the LTX-2.5 example scripts. --- README.md | 3 +- README_zh.md | 2 +- diffsynth/models/ltx2_text_encoder.py | 4 - diffsynth/pipelines/ltx2_audio_video.py | 9 +- docs/en/Model_Details/LTX-2.5.md | 128 ------------------ docs/en/Model_Details/LTX-2.md | 18 ++- docs/zh/Model_Details/LTX-2.5.md | 128 ------------------ docs/zh/Model_Details/LTX-2.md | 19 ++- .../model_inference/LTX-2.5-A2V-TwoStage.py | 3 +- .../model_inference/LTX-2.5-I2AV-OneStage.py | 1 - .../model_inference/LTX-2.5-I2AV-TwoStage.py | 2 - .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 1 - .../LTX-2.5-T2AV-DistilledPipeline.py | 2 - .../LTX-2.5-T2AV-TwoStage-Retake.py | 2 - .../LTX-2.5-A2V-TwoStage.py | 1 - .../LTX-2.5-I2AV-OneStage.py | 1 - .../LTX-2.5-I2AV-TwoStage.py | 2 - .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 1 - .../LTX-2.5-T2AV-DistilledPipeline.py | 1 - .../LTX-2.5-T2AV-TwoStage-Retake.py | 2 - 20 files changed, 41 insertions(+), 289 deletions(-) delete mode 100644 docs/en/Model_Details/LTX-2.5.md delete mode 100644 docs/zh/Model_Details/LTX-2.5.md diff --git a/README.md b/README.md index 2ea0ccfa6..b73142c38 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ See also: > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand. -- **September 7, 2026** We have integrated [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5), the latest audio-video generation model from Lightricks. The features include text-to-audio/video with automatic duration prediction, image-to-audio/video with keyframe interpolation, audio-to-video, audio-video retake, IC-LoRA pixel spatial upscaling, text-to-audio, INT8 quantized inference, low VRAM inference, and training. For details, please refer to the [documentation](/docs/en/Model_Details/LTX-2.5.md) and [code](/examples/ltx2/). +- **September 10, 2026** We have integrated [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5), the latest audio-video generation model from Lightricks. The features include text-to-audio/video with automatic duration prediction, image-to-audio/video with keyframe interpolation, audio-to-video, audio-video retake, IC-LoRA pixel spatial upscaling, text-to-audio, INT8 quantized inference, low VRAM inference, and training. For details, please refer to the [documentation](/docs/en/Model_Details/LTX-2.md) and [code](/examples/ltx2/). - **September 1, 2026** We have integrated [SenseNova-U1.5](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT), SenseTime's unified multimodal model, for which we provide text-to-image generation, image editing, low VRAM inference, and training support. For details, please refer to the [documentation](/docs/en/Model_Details/SenseNova-U1.md) and [example code](/examples/sensenova_u1/). @@ -329,7 +329,6 @@ Model overview: - Video generation - MiniMax-H3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-H3.html), [Example code](/examples/minimax_h3/) - LingBot-Video: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LingBot-Video.html), [Example code](/examples/lingbot_video/) - - LTX-2.5: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.5.html), [Example code](/examples/ltx2/) - LTX-2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.html), [Example code](/examples/ltx2/) - Wan: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html), [Example code](/examples/wanvideo/) - Audio generation diff --git a/README_zh.md b/README_zh.md index fb4ea949f..c1742164c 100644 --- a/README_zh.md +++ b/README_zh.md @@ -40,7 +40,7 @@ > 目前本项目的开发人员有限，大部分工作由 [Artiprocher](https://github.com/Artiprocher) 和 [mi804](https://github.com/mi804) 负责，因此新功能的开发进展会比较缓慢，issue 的回复和解决速度有限，我们对此感到非常抱歉，请各位开发者理解。 -- **2026年9月7日** 我们接入了 [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5)，这是 Lightricks 最新的音视频联合生成模型。支持的功能包括自动时长预测的文生音视频、关键帧插值的图生音视频、音频驱动视频、音视频区域重生成、IC-LoRA 像素空间上采样、文生音频、INT8 量化推理、低显存推理以及模型训练。详情请参考[文档](/docs/zh/Model_Details/LTX-2.5.md)和[示例代码](/examples/ltx2/)。 +- **2026年9月10日** 我们接入了 [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5)，这是 Lightricks 最新的音视频联合生成模型。支持的功能包括自动时长预测的文生音视频、关键帧插值的图生音视频、音频驱动视频、音视频区域重生成、IC-LoRA 像素空间上采样、文生音频、INT8 量化推理、低显存推理以及模型训练。详情请参考[文档](/docs/zh/Model_Details/LTX-2.md)和[示例代码](/examples/ltx2/)。 - **2026年9月1日** 我们接入了 [SenseNova-U1.5](https://www.modelscope.cn/models/SenseNova/SenseNova-U1.5-8B-MoT)，这是商汤科技开源的统一多模态模型，我们为其提供了文生图、图像编辑、低显存推理和训练支持。详情请参考[文档](/docs/zh/Model_Details/SenseNova-U1.md)和[示例代码](/examples/sensenova_u1/)。 diff --git a/diffsynth/models/ltx2_text_encoder.py b/diffsynth/models/ltx2_text_encoder.py index 3570df568..e4f3b1a3e 100644 --- a/diffsynth/models/ltx2_text_encoder.py +++ b/diffsynth/models/ltx2_text_encoder.py @@ -225,7 +225,6 @@ def __init__( dim_head: int, rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, apply_gated_attention: bool = False, - ff_bias: bool = True, ): super().__init__() @@ -240,7 +239,6 @@ def __init__( self.ff = FeedForward( dim, dim_out=dim, - bias=ff_bias, ) def forward( @@ -309,7 +307,6 @@ def __init__( rope_type: LTXRopeType = LTXRopeType.SPLIT, double_precision_rope: bool = True, apply_gated_attention: bool = False, - ff_bias: bool = True, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -329,7 +326,6 @@ def __init__( dim_head=attention_head_dim, rope_type=rope_type, apply_gated_attention=apply_gated_attention, - ff_bias=ff_bias, ) for _ in range(num_layers) ] diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 391bf96f9..472f971a6 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -403,7 +403,8 @@ def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=video_latent_shape, device=pipe.device) video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, True).float() video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate - video_positions = video_positions.to(pipe.torch_dtype) + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) video_keyframes_mask = self.build_video_keyframes_mask(pipe, video_latent_shape) audio_latent_shape = AudioLatentShape.from_video_pixel_shape(video_pixel_shape) @@ -595,7 +596,8 @@ def process( latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=VideoLatentShape.from_torch_shape(latents.shape), device=pipe.device) video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, False).float() video_positions[:, 0, ...] = (video_positions[:, 0, ...] + index) / frame_rate - video_positions = video_positions.to(pipe.torch_dtype) + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) frame_conditions["ref_frames_latents"].append(latents) frame_conditions["ref_frames_positions"].append(video_positions) if len(frame_conditions["ref_frames_latents"]) == 0: @@ -642,7 +644,8 @@ def process(self, pipe: LTX2AudioVideoPipeline, in_context_videos, height, width video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate video_positions[:, 1, ...] *= in_context_downsample_factor # height axis video_positions[:, 2, ...] *= in_context_downsample_factor # width axis - video_positions = video_positions.to(pipe.torch_dtype) + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) latents.append(in_context_latents) positions.append(video_positions) diff --git a/docs/en/Model_Details/LTX-2.5.md b/docs/en/Model_Details/LTX-2.5.md deleted file mode 100644 index f21a1fcf1..000000000 --- a/docs/en/Model_Details/LTX-2.5.md +++ /dev/null @@ -1,128 +0,0 @@ -# LTX-2.5 - -LTX-2.5 is the joint audio-video generation model released by Lightricks. DiffSynth-Studio supports its inference and training through `LTX2AudioVideoPipeline` (the same pipeline class used by LTX-2 / LTX-2.3). Compared with LTX-2.3, LTX-2.5 introduces a fine-tuned Gemma4 12B text encoder (with packed tokenizer assets and dual audio/video connectors), the DiffVAE diffusion video decoder, automatic duration prediction via a Duration Head, and INT8 quantized checkpoints. - -## Installation - -Before using this project for inference or training, please install DiffSynth-Studio. - -```shell -git clone https://github.com/modelscope/DiffSynth-Studio.git -cd DiffSynth-Studio -pip install -e . -``` - -For more information, see [Setup](../Pipeline_Usage/Setup.md). The LTX-2.5 Gemma4 text encoder requires `transformers>=5.8,<5.15`. - -## Quickstart - -The code below loads [Lightricks/LTX-2.5](https://www.modelscope.cn/models/Lightricks/LTX-2.5) and runs inference. With `auto_duration` enabled, the pipeline first predicts the clip length from the prompt with the Duration Head and then generates audio and video, all in a single `pipe(...)` call. - -```python -import torch -from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig -from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 - -vram_config = { - "offload_dtype": torch.bfloat16, - "offload_device": "cpu", - "onload_dtype": torch.bfloat16, - "onload_device": "cuda", - "preparing_dtype": torch.bfloat16, - "preparing_device": "cuda", - "computation_dtype": torch.bfloat16, - "computation_device": "cuda", -} -pipe = LTX2AudioVideoPipeline.from_pretrained( - torch_dtype=torch.bfloat16, - device="cuda", - model_configs=[ - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), - ], -) -prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] -video, audio = pipe( - prompt=prompt, - negative_prompt=negative_prompt, - seed=43, - height=1024, width=1536, frame_rate=24, - auto_duration=True, - cfg_scale=1.0, num_inference_steps=8, - use_distilled_pipeline=True, use_two_stage_pipeline=True, - tiled=True, -) -write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate) -``` - -## Models - -|Model ID|Extra parameters|Inference|Low-VRAM inference|Full training|Validate full|LoRA training|Validate LoRA| -|-|-|-|-|-|-|-|-| -|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| -|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| -|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| -|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| -|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| -|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| - -## Inference - -Models are loaded with `LTX2AudioVideoPipeline.from_pretrained`; see [Load Models](../Pipeline_Usage/Model_Inference.md#load-models). LTX-2.5 shares `LTX2AudioVideoPipeline` with LTX-2.3, and the framework detects the model version from the loaded weights. - -LTX-2.5 related `from_pretrained` arguments: - -* `tokenizer_config`: source of the tokenizer assets. The LTX-2.5 tokenizer is unpacked into the HF-style directory `DiffSynth-Studio/LTX-2.5-Repackage/tokenizer` (produced by `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py`), matching how LTX-2.3 references its tokenizer. -* `text_encoder_post_modules`: the LTX-2.5 feature extractor and embeddings connectors weights are packed into `DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors` (produced by `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py` from the TE and transformer checkpoints) and must be included in `model_configs`. -* `stage2_lora_config`: the stage-2 distilled LoRA used for two-stage inference with the Dev weights. - -For the arguments shared with LTX-2.3, see the [LTX-2 documentation](LTX-2.md#inference). The new or LTX-2.5 specific `LTX2AudioVideoPipeline` arguments are: - -* `auto_duration`: predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. -* `auto_duration_min_seconds` / `auto_duration_max_seconds`: lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. -* `audio_only`: whether to generate audio only (T2A). Defaults to `False`; when `True`, the video VAE and latent upsampler are not required. -* The video decoder is selected by the loaded components: the ConvVAE convolutional decoder is used when `ltx-2.5-video-vae-conv-bf16.safetensors` is loaded, otherwise the DiffVAE diffusion decoder. -* Default negative prompt: `pipe.default_negative_prompt["LTX-2.5"]` prefixes the LTX-2/2.3 list with the 2.5-specific tags (`has_subtitles`, `has_blurbox`, `transition from black`, `transition to black`, `speech_ending_short`); all example scripts use this key. -* `input_images` / `input_images_indexes`: keyframe images and their frame indexes. A single first frame gives image-to-video; first and last (or more) frames give keyframe interpolation. -* `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: audio-to-video (A2V) and audio region retake. - -Geometry constraints: `num_frames % 8 == 1`; height and width must be multiples of 32 for the one-stage pipeline and multiples of 64 for the two-stage pipeline. - -If you are short of VRAM, enable [VRAM management](../Pipeline_Usage/VRAM_management.md). Each example script ships a recommended low-VRAM configuration (FP8 CPU weight offload plus fine-grained VRAM management); see the table above. - -## Training - -LTX-2.5 shares the training script [`examples/ltx2/model_training/train.py`](/examples/ltx2/model_training/train.py) with LTX-2 / LTX-2.3. The general training arguments are documented in the [LTX-2 documentation](LTX-2.md#training). - -The 22B DiT and the 12B Gemma4 encoder do not fit on a single GPU together, so the LTX-2.5 training scripts use the two-stage (splited) scheme: - -1. `--task "sft:data_process"`: run text encoding and VAE encoding and cache the results to disk. -2. `--task "sft:train"`: read the cached results and train the DiT only. - -Both stages keep the full `--model_id_with_origin_paths` list and use `--fp8_models` to declare the models that are not forwarded in that stage (stage two loads the text encoder and the VAEs in FP8). The training dataset columns are `video,prompt,input_audio,frame_rate`, which map to `--data_file_keys "video,input_audio"` and `--extra_inputs "input_audio"`. - -A sample dataset is available for testing: - -```shell -modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset -``` - -After training, use `examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py` (LoRA) or `examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py` (full) to run inference with the trained checkpoint. For more details on writing training scripts, see [Model Training](../Pipeline_Usage/Model_Training.md). - -## Unsupported features - -The following official LTX-2.5 capabilities are not integrated yet: - -* DFR (Diffusion Frame Rate) -* Native HDR / EXR output -* HDR IC-LoRA -* Dub-It dubbing diff --git a/docs/en/Model_Details/LTX-2.md b/docs/en/Model_Details/LTX-2.md index 29f5654ca..bda9a5b45 100644 --- a/docs/en/Model_Details/LTX-2.md +++ b/docs/en/Model_Details/LTX-2.md @@ -85,6 +85,16 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Up.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Up.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Down.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Down.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Static](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Static.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Static.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| +|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| ## Model Inference @@ -94,16 +104,19 @@ Input parameters for `LTX2AudioVideoPipeline` inference include: * `prompt`: Prompt describing the content appearing in the video. * `negative_prompt`: Negative prompt describing content that should not appear in the video, default value is `""`. -* `cfg_scale`: Classifier-free guidance parameter, default value is 3.0. +* `denoising_strength`: Denoising strength, range is 0～1, default value is 1.0. * `input_images`: List of input images for image-to-video generation. * `input_images_indexes`: Frame index list of input images in the video. * `input_images_strength`: Strength of input images, default value is 1.0. -* `denoising_strength`: Denoising strength, range is 0～1, default value is 1.0. * `seed`: Random seed. Default is `None`, which means completely random. * `rand_device`: Computing device for generating random Gaussian noise matrix, default is `"cpu"`. When set to `cuda`, different results will be generated on different GPUs. * `height`: Video height, must be a multiple of 32 (single-stage) or 64 (two-stage). * `width`: Video width, must be a multiple of 32 (single-stage) or 64 (two-stage). * `num_frames`: Number of video frames, default value is 121, must be a multiple of 8 + 1. +* `auto_duration`: Predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. +* `auto_duration_min_seconds` / `auto_duration_max_seconds`: Lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. +* `audio_only`: Whether to generate audio only (T2A). Defaults to `False`; when `True`, the video VAE and latent upsampler are not required. +* `cfg_scale`: Classifier-free guidance parameter, default value is 3.0. * `num_inference_steps`: Number of inference steps, default value is 40. * `tiled`: Whether to enable VAE tiling inference, default is `True`. When set to `True`, it can significantly reduce VRAM usage during VAE encoding/decoding stages, with slight errors and minor inference time extension. * `tile_size_in_pixels`: Pixel tiling size during VAE encoding/decoding stages, default is 512. @@ -114,6 +127,7 @@ Input parameters for `LTX2AudioVideoPipeline` inference include: * `use_distilled_pipeline`: Whether to use distilled pipeline, default is `False`. * `progress_bar_cmd`: Progress bar, default is `tqdm.tqdm`. Can be set to `lambda x:x` to hide the progress bar. + If VRAM is insufficient, please enable [VRAM Management](../Pipeline_Usage/VRAM_management.md). We provide recommended low VRAM configurations for each model in the example code, see the table in the previous "Supported Inference Scripts" section. ## Model Training diff --git a/docs/zh/Model_Details/LTX-2.5.md b/docs/zh/Model_Details/LTX-2.5.md deleted file mode 100644 index 39ef46148..000000000 --- a/docs/zh/Model_Details/LTX-2.5.md +++ /dev/null @@ -1,128 +0,0 @@ -# LTX-2.5 - -LTX-2.5 是 Lightricks 发布的音视频联合生成模型。DiffSynth-Studio 通过 `LTX2AudioVideoPipeline` 提供其推理与训练支持（与 LTX-2 / LTX-2.3 共用同一个 Pipeline 类）。相比 LTX-2.3，LTX-2.5 引入了微调版 Gemma4 12B 文本编码器（内嵌 tokenizer 资产与音视频双 connector）、DiffVAE 扩散视频解码器、Duration Head 自动时长预测，以及 INT8 量化权重。 - -## 安装 - -在使用本项目进行模型推理和训练前，请先安装 DiffSynth-Studio。 - -```shell -git clone https://github.com/modelscope/DiffSynth-Studio.git -cd DiffSynth-Studio -pip install -e . -``` - -更多关于安装的信息，请参考[安装依赖](../Pipeline_Usage/Setup.md)。LTX-2.5 的 Gemma4 文本编码器需要 `transformers>=5.8,<5.15`。 - -## 快速开始 - -运行以下代码可以快速加载 [Lightricks/LTX-2.5](https://www.modelscope.cn/models/Lightricks/LTX-2.5) 模型并进行推理。开启 `auto_duration` 后，Pipeline 会先用 Duration Head 从提示词预测视频时长，再完成音视频生成，整个过程只需一次 `pipe(...)` 调用。 - -```python -import torch -from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig -from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 - -vram_config = { - "offload_dtype": torch.bfloat16, - "offload_device": "cpu", - "onload_dtype": torch.bfloat16, - "onload_device": "cuda", - "preparing_dtype": torch.bfloat16, - "preparing_device": "cuda", - "computation_dtype": torch.bfloat16, - "computation_device": "cuda", -} -pipe = LTX2AudioVideoPipeline.from_pretrained( - torch_dtype=torch.bfloat16, - device="cuda", - model_configs=[ - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), - ], -) -prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" -negative_prompt = pipe.default_negative_prompt["LTX-2.3"] -video, audio = pipe( - prompt=prompt, - negative_prompt=negative_prompt, - seed=43, - height=1024, width=1536, frame_rate=24, - auto_duration=True, - cfg_scale=1.0, num_inference_steps=8, - use_distilled_pipeline=True, use_two_stage_pipeline=True, - tiled=True, -) -write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate) -``` - -## 模型总览 - -|模型 ID|额外参数|推理|低显存推理|全量训练|全量训练后验证|LoRA 训练|LoRA 训练后验证| -|-|-|-|-|-|-|-|-| -|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| -|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| -|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| -|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| -|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| -|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| -|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| - -## 模型推理 - -模型通过 `LTX2AudioVideoPipeline.from_pretrained` 加载，详见[加载模型](../Pipeline_Usage/Model_Inference.md#加载模型)。LTX-2.5 与 LTX-2.3 共用 `LTX2AudioVideoPipeline`，框架根据加载到的权重自动识别模型版本。 - -`from_pretrained` 的 LTX-2.5 相关参数： - -* `tokenizer_config`：tokenizer 资产来源。LTX-2.5 的 tokenizer 已解包为 HF 目录形式 `DiffSynth-Studio/LTX-2.5-Repackage/tokenizer`（由 `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py` 生成），与 LTX-2.3 的用法一致。 -* `text_encoder_post_modules`：LTX-2.5 的 feature extractor 与 embeddings connectors 权重打包在 `DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors`（由 `examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py` 从 TE 与 transformer 权重中提取），需要在 `model_configs` 中一并加载。 -* `stage2_lora_config`: Dev 权重两阶段推理时使用的第二阶段 distilled-LoRA。 - -`LTX2AudioVideoPipeline` 的通用推理参数见 [LTX-2 文档](LTX-2.md#模型推理)，LTX-2.5 新增或特有的参数为： - -* `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 -* `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 -* `audio_only`: 是否只生成音频（T2A），默认为 `False`。设置为 `True` 时无需加载视频 VAE 与 latent upsampler。 -* 视频解码器按已加载组件自动选择：加载了 `ltx-2.5-video-vae-conv-bf16.safetensors` 就用 ConvVAE 卷积解码器，否则使用 DiffVAE 扩散解码器。 -* 默认负向提示词：`pipe.default_negative_prompt["LTX-2.5"]` 在 LTX-2/2.3 的列表之前增加了 2.5 专有标签（`has_subtitles`、`has_blurbox`、`transition from black`、`transition to black`、`speech_ending_short`），示例脚本均使用该键。 -* `input_images` / `input_images_indexes`: 关键帧图像及其帧索引。传入首帧即为图生视频，传入首尾（或多帧）即为关键帧插值。 -* `retake_audio` / `audio_sample_rate` / `retake_audio_regions`: 音频驱动视频（A2V）与音频区域重生成。 - -几何约束：`num_frames % 8 == 1`；单阶段的宽高为 32 的倍数，两阶段的宽高为 64 的倍数。 - -如果显存不足，请开启[显存管理](../Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置（FP8 CPU 权重卸载 + 细粒度显存管理），详见前文"模型总览"中的表格。 - -## 模型训练 - -LTX-2.5 与 LTX-2 / LTX-2.3 共用训练脚本 [`examples/ltx2/model_training/train.py`](/examples/ltx2/model_training/train.py)，通用训练参数的说明见 [LTX-2 文档](LTX-2.md#模型训练)。 - -由于 22B DiT 与 12B Gemma4 编码器无法同时放入单卡，LTX-2.5 的训练脚本采用双阶段（splited）方案： - -1. `--task "sft:data_process"`：运行文本编码与 VAE 编码，把结果缓存到硬盘。 -2. `--task "sft:train"`：从缓存读取前处理结果，只训练 DiT。 - -两个阶段都保留完整的 `--model_id_with_origin_paths`，并用 `--fp8_models` 声明该阶段不需要前向的模型（阶段二的 TextEncoder 与 VAE 使用 FP8 加载）。训练数据集的字段为 `video,prompt,input_audio,frame_rate`，对应 `--data_file_keys "video,input_audio"` 与 `--extra_inputs "input_audio"`。 - -我们构建了一个样例视频数据集，以方便您进行测试，通过以下命令可以下载这个数据集： - -```shell -modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset -``` - -训练完成后，可以使用 `examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py`（LoRA）或 `examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py`（全量）加载训练产物进行推理验证。关于如何编写模型训练脚本，请参考[模型训练](../Pipeline_Usage/Model_Training.md)。 - -## 暂不支持的功能 - -以下 LTX-2.5 官方能力暂未接入： - -* DFR（Diffusion Frame Rate） -* Native HDR / EXR 输出 -* HDR IC-LoRA -* Dub-It 配音 diff --git a/docs/zh/Model_Details/LTX-2.md b/docs/zh/Model_Details/LTX-2.md index baa674347..0261f103a 100644 --- a/docs/zh/Model_Details/LTX-2.md +++ b/docs/zh/Model_Details/LTX-2.md @@ -85,6 +85,16 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Up.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Up.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Down.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Down.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Static](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Static.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Static.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| +|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| ## 模型推理 @@ -94,16 +104,19 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `prompt`: 提示词，描述视频中出现的内容。 * `negative_prompt`: 负向提示词，描述视频中不应该出现的内容，默认值为 `""`。 -* `cfg_scale`: Classifier-free guidance 的参数，默认值为 3.0。 +* `denoising_strength`: 去噪强度，范围是 0～1，默认值为 1.0。 * `input_images`: 输入图像列表，用于图生视频。 * `input_images_indexes`: 输入图像在视频中的帧索引列表。 * `input_images_strength`: 输入图像的强度，默认值为 1.0。 -* `denoising_strength`: 去噪强度，范围是 0～1，默认值为 1.0。 * `seed`: 随机种子。默认为 `None`，即完全随机。 * `rand_device`: 生成随机高斯噪声矩阵的计算设备，默认为 `"cpu"`。当设置为 `cuda` 时，在不同 GPU 上会导致不同的生成结果。 * `height`: 视频高度，需保证高度为 32 的倍数（单阶段）或 64 的倍数（两阶段）。 * `width`: 视频宽度，需保证宽度为 32 的倍数（单阶段）或 64 的倍数（两阶段）。 * `num_frames`: 视频帧数，默认值为 121，需保证为 8 的倍数 + 1。 +* `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 +* `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 +* `audio_only`: 是否只生成音频（T2A），默认为 `False`。设置为 `True` 时无需加载视频 VAE 与 latent upsampler。 +* `cfg_scale`: Classifier-free guidance 的参数，默认值为 3.0。 * `num_inference_steps`: 推理次数，默认值为 40。 * `tiled`: 是否启用 VAE 分块推理，默认为 `True`。设置为 `True` 时可显著减少 VAE 编解码阶段的显存占用，会产生少许误差，以及少量推理时间延长。 * `tile_size_in_pixels`: VAE 编解码阶段的像素分块大小，默认为 512。 @@ -114,7 +127,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `use_distilled_pipeline`: 是否使用蒸馏管道，默认为 `False`。 * `progress_bar_cmd`: 进度条，默认为 `tqdm.tqdm`。可通过设置为 `lambda x:x` 来屏蔽进度条。 -如果显存不足，请开启[显存管理](../Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置，详见前文"支持的推理脚本"中的表格。 +如果显存不足，请开启[显存管理](../Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置，详见前文"模型总览"中的表格。 ## 模型训练 diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py index 00b9df838..64b5c6e7b 100644 --- a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -33,7 +33,6 @@ ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example audio comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 @@ -44,7 +43,7 @@ negative_prompt=negative_prompt, retake_audio=audio, audio_sample_rate=audio_sample_rate, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py index 136ac8609..e3c925769 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -29,7 +29,6 @@ ], ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example image comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512, 768, 121 diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index 2e3d380f6..174f2ef48 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -32,7 +32,6 @@ stage2_lora_strength=1.0, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example images come from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 @@ -63,7 +62,6 @@ ) pipe.clear_lora() -# Keyframe interpolation: any frames can be used by setting input_images and input_images_indexes within the range of num_frames. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index f137ed890..456ab29e6 100644 --- a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -38,7 +38,6 @@ ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The reference video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py index 0078a7472..786d0ace5 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -32,7 +32,6 @@ prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width = 512 * 2, 768 * 2 -# Automatic duration: one pipe call predicts the clip length from the prompt and generates it. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, @@ -57,4 +56,3 @@ fps=24, audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, ) -print(f"saved to ltx2.5_distilled_t2av.mp4") diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py index d93ad0b47..099a96259 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -34,7 +34,6 @@ ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] @@ -44,7 +43,6 @@ assert len(video) == num_frames, f"Input video has {len(video)} frames, but expected {num_frames} frames based on the specified num_frames argument." audio, audio_sample_rate = read_audio(path) -# Regenerate the video within time regions. Retake regions are in seconds. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py index b91bd3393..f9e6bf6ee 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -34,7 +34,6 @@ ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example audio comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index 2019bac24..28cc46ed3 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -30,7 +30,6 @@ vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example image comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512, 768, 121 diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index 8bce88328..a00b3774b 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -33,7 +33,6 @@ vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example images come from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 @@ -64,7 +63,6 @@ ) pipe.clear_lora() -# Keyframe interpolation: any frames can be used by setting input_images and input_images_indexes within the range of num_frames. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index f15c09a3f..22a7488ca 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -39,7 +39,6 @@ ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The reference video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index b2d692b1c..c31e24708 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -33,7 +33,6 @@ prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width = 512 * 2, 768 * 2 -# Automatic duration: one pipe call predicts the clip length from the prompt and generates it. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py index 15472cfb8..a14f53cb2 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -35,7 +35,6 @@ ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -# The example video comes from the shared sample dataset, so reuse its paired prompt. prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] @@ -45,7 +44,6 @@ assert len(video) == num_frames, f"Input video has {len(video)} frames, but expected {num_frames} frames based on the specified num_frames argument." audio, audio_sample_rate = read_audio(path) -# Regenerate the video within time regions. Retake regions are in seconds. video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, From ae1f2c47d88933992cc2139a6fb53c23f7e5d6fc Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Thu, 10 Sep 2026 10:17:08 +0800 Subject: [PATCH 27/31] revert gitignore --- .gitignore | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.gitignore b/.gitignore index dfe725a05..7cce2c12d 100644 --- a/.gitignore +++ b/.gitignore @@ -2,10 +2,6 @@ /models /scripts /diffusers -# Local inference media and isolated upstream/environment artifacts. -# These can contain machine-specific paths, authenticated inspection logs, or large binaries. -/outputs -/packages /.vscode /.opencode *.pkl From 1f9866da1987c907de5ba17032e070c59fe4f139 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Thu, 10 Sep 2026 15:48:54 +0800 Subject: [PATCH 28/31] refactor ltx vram --- .../configs/vram_management_module_maps.py | 13 +++++---- diffsynth/models/ltx25_diffusion_video_vae.py | 10 ++++--- diffsynth/models/ltx25_text_encoder.py | 6 ++++ diffsynth/models/ltx2_audio_vae.py | 29 +++++++++++++------ diffsynth/models/ltx2_video_vae.py | 10 ++++--- diffsynth/pipelines/ltx2_audio_video.py | 2 +- .../LTX-2.5-A2V-TwoStage.py | 10 +++---- .../LTX-2.5-I2AV-OneStage.py | 8 ++--- .../LTX-2.5-I2AV-TwoStage.py | 8 ++--- .../LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py | 8 ++--- .../model_inference_low_vram/LTX-2.5-T2A.py | 8 ++--- .../LTX-2.5-T2AV-DistilledPipeline.py | 8 ++--- .../LTX-2.5-T2AV-INT8-ConvRot.py | 8 ++--- .../LTX-2.5-T2AV-OneStage.py | 8 ++--- .../LTX-2.5-T2AV-TwoStage-Retake.py | 8 ++--- .../LTX-2.5-T2AV-TwoStage.py | 10 +++---- 16 files changed, 88 insertions(+), 66 deletions(-) diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index dc41d9e69..249c72f8f 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -273,6 +273,7 @@ "diffsynth.models.ltx2_dit.LTXModel": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_dit.BasicAVTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", @@ -285,24 +286,24 @@ "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_text_encoder.LTX25Embeddings1DConnector": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffusion_video_vae.DiffusionVideoDecoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": { - "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": { - "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": { - "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": { - "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_audio_vae.LTX2Vocoder": { "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py index 9d77da8a7..d340ec3a1 100644 --- a/diffsynth/models/ltx25_diffusion_video_vae.py +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -1078,16 +1078,18 @@ def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> to class PerChannelStatistics(nn.Module): def __init__(self, latent_channels: int = 128): super().__init__() - self.register_buffer("std-of-means", torch.ones(latent_channels)) - self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + # Frozen parameters, not buffers: disk-offload reload restores only named_parameters(), + # so persistent buffers would be skipped and strict load_state_dict would report them missing. + self.register_parameter("std-of-means", nn.Parameter(torch.ones(latent_channels), requires_grad=False)) + self.register_parameter("mean-of-means", nn.Parameter(torch.zeros(latent_channels), requires_grad=False)) def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + return (x * getattr(self, "std-of-means").view(1, -1, 1, 1, 1).to(x)) + getattr(self, "mean-of-means").view( 1, -1, 1, 1, 1 ).to(x) def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + return (x - getattr(self, "mean-of-means").view(1, -1, 1, 1, 1).to(x)) / getattr(self, "std-of-means").view( 1, -1, 1, 1, 1 ).to(x) diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py index 6e6d8e2fa..f0d05a0af 100644 --- a/diffsynth/models/ltx25_text_encoder.py +++ b/diffsynth/models/ltx25_text_encoder.py @@ -116,6 +116,12 @@ def __init__(self): }, } super().__init__(Gemma4UnifiedConfig(**config)) + # Gemma4 registers a constant `layer_scalar` as a persistent buffer on each decoder layer. + # Disk-offload reload restores only named_parameters(), so a persistent buffer is never fetched + # from disk and strict load_state_dict fails; register it as a frozen parameter instead. + for module in self.modules(): + if "layer_scalar" in module._buffers: + module.register_parameter("layer_scalar", torch.nn.Parameter(module._buffers.pop("layer_scalar"), requires_grad=False)) class LTX25GemmaTokenizer: diff --git a/diffsynth/models/ltx2_audio_vae.py b/diffsynth/models/ltx2_audio_vae.py index 8a58f9724..0c48583a3 100644 --- a/diffsynth/models/ltx2_audio_vae.py +++ b/diffsynth/models/ltx2_audio_vae.py @@ -821,14 +821,16 @@ class PerChannelStatistics(nn.Module): def __init__(self, latent_channels: int = 128) -> None: super().__init__() - self.register_buffer("std-of-means", torch.empty(latent_channels)) - self.register_buffer("mean-of-means", torch.empty(latent_channels)) + # Frozen parameters, not buffers: disk-offload reload restores only named_parameters(), + # so persistent buffers would be skipped and strict load_state_dict would report them missing. + self.register_parameter("std-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) + self.register_parameter("mean-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x) + return (x * getattr(self, "std-of-means").to(x)) + getattr(self, "mean-of-means").to(x) def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x) + return (x - getattr(self, "mean-of-means").to(x)) / getattr(self, "std-of-means").to(x) LATENT_DOWNSAMPLE_FACTOR = 4 @@ -1342,7 +1344,8 @@ def __init__( self.stride = stride self.padding = padding self.padding_mode = padding_mode - self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + # Parameter, not buffer, so disk-offload reload (named_parameters() only) restores it from the checkpoint. + self.register_parameter("filter", nn.Parameter(kaiser_sinc_filter1d(cutoff, half_width, kernel_size), requires_grad=False)) def forward(self, x: torch.Tensor) -> torch.Tensor: _, n_channels, _ = x.shape @@ -1388,7 +1391,12 @@ def __init__( kernel_size=self.kernel_size, ) - self.register_buffer("filter", sinc_filter, persistent=persistent) + if persistent: + # Checkpoint-stored: parameter so disk-offload reload (named_parameters() only) restores it. + self.register_parameter("filter", nn.Parameter(sinc_filter, requires_grad=False)) + else: + # Not stored in the checkpoint: keep as a non-persistent buffer (recomputed above). + self.register_buffer("filter", sinc_filter, persistent=False) def forward(self, x: torch.Tensor) -> torch.Tensor: _, n_channels, _ = x.shape @@ -1701,8 +1709,10 @@ def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None self.hop_length = hop_length self.win_length = win_length n_freqs = filter_length // 2 + 1 - self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)) - self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length)) + # Parameters, not buffers: disk-offload reload restores only named_parameters(); + # the zeros here are overwritten by load_state_dict from the checkpoint. + self.register_parameter("forward_basis", nn.Parameter(torch.zeros(n_freqs * 2, 1, filter_length), requires_grad=False)) + self.register_parameter("inverse_basis", nn.Parameter(torch.zeros(n_freqs * 2, 1, filter_length), requires_grad=False)) def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Compute magnitude and phase spectrogram from a batch of waveforms. @@ -1746,8 +1756,9 @@ def __init__( # Initialized to zeros; load_state_dict overwrites with the checkpoint's # exact bfloat16 filterbank (vocoder.mel_stft.mel_basis, shape [n_mels, n_freqs]). + # Parameter, not buffer, so disk-offload reload (named_parameters() only) restores it. n_freqs = filter_length // 2 + 1 - self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) + self.register_parameter("mel_basis", nn.Parameter(torch.zeros(n_mel_channels, n_freqs), requires_grad=False)) def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Compute log-mel spectrogram and auxiliary spectral quantities. diff --git a/diffsynth/models/ltx2_video_vae.py b/diffsynth/models/ltx2_video_vae.py index c98dad160..026c63ff1 100644 --- a/diffsynth/models/ltx2_video_vae.py +++ b/diffsynth/models/ltx2_video_vae.py @@ -553,15 +553,17 @@ class PerChannelStatistics(nn.Module): def __init__(self, latent_channels: int = 128): super().__init__() - self.register_buffer("std-of-means", torch.empty(latent_channels)) - self.register_buffer("mean-of-means", torch.empty(latent_channels)) + # Frozen parameters, not buffers: disk-offload reload restores only named_parameters(), + # so persistent buffers would be skipped and strict load_state_dict would report them missing. + self.register_parameter("std-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) + self.register_parameter("mean-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + return (x * getattr(self, "std-of-means").view(1, -1, 1, 1, 1).to(x)) + getattr(self, "mean-of-means").view( 1, -1, 1, 1, 1).to(x) def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + return (x - getattr(self, "mean-of-means").view(1, -1, 1, 1, 1).to(x)) / getattr(self, "std-of-means").view( 1, -1, 1, 1, 1).to(x) diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 472f971a6..46c86d784 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -699,7 +699,7 @@ def __init__(self): super().__init__( input_params=("video_latents",), output_params=("video_latents",), - onload_model_names=("upsampler",), + onload_model_names=("upsampler", "video_vae_encoder"), ) def process(self, pipe: LTX2AudioVideoPipeline, video_latents): diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py index f9e6bf6ee..95d3d24e9 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -5,10 +5,10 @@ from modelscope import dataset_snapshot_download vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, @@ -44,7 +44,7 @@ negative_prompt=negative_prompt, retake_audio=audio, audio_sample_rate=audio_sample_rate, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index 28cc46ed3..e0136d4c1 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -5,10 +5,10 @@ from modelscope import dataset_snapshot_download vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index a00b3774b..8999acd0f 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -5,10 +5,10 @@ from modelscope import dataset_snapshot_download vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py index 22a7488ca..56a7ce198 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -5,10 +5,10 @@ from modelscope import dataset_snapshot_download vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py index d8fd3ce04..c7fa4c517 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py @@ -3,10 +3,10 @@ from diffsynth.utils.data.audio import save_audio vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index c31e24708..9088fa6eb 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -3,9 +3,9 @@ from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.bfloat16, "onload_device": "cpu", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", @@ -36,7 +36,7 @@ video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, frame_rate=24, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py index dec854bcc..d144681a4 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py @@ -3,10 +3,10 @@ from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py index 4660966b1..fa33a3f31 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -3,10 +3,10 @@ from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py index a14f53cb2..fffae6d2b 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -6,10 +6,10 @@ from modelscope import dataset_snapshot_download vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py index b884791d7..903c11247 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -3,10 +3,10 @@ from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 vram_config = { - "offload_dtype": torch.float8_e5m2, - "offload_device": "cpu", - "onload_dtype": torch.float8_e5m2, - "onload_device": "cpu", + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, @@ -36,7 +36,7 @@ video, audio = pipe( prompt=prompt, negative_prompt=negative_prompt, - seed=43, + seed=42, height=height, width=width, num_frames=num_frames, From 7e24c8058fb26a65ae00c47b31be46b20655dde8 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Thu, 10 Sep 2026 18:05:03 +0800 Subject: [PATCH 29/31] final check --- .../configs/vram_management_module_maps.py | 4 +- diffsynth/pipelines/ltx2_audio_video.py | 2 - .../model_inference/LTX-2.5-I2AV-OneStage.py | 2 +- .../model_inference/LTX-2.5-I2AV-TwoStage.py | 2 +- .../LTX-2.5-I2AV-OneStage.py | 2 +- .../LTX-2.5-I2AV-TwoStage.py | 2 +- .../LTX-2.5-T2AV-DistilledPipeline.py | 4 +- .../full/LTX-2.5-I2AV-splited.sh | 41 ++++++++++++++ ...plited-test.sh => LTX-2.5-I2AV-splited.sh} | 24 ++++----- .../validate_full/LTX-2.5-I2AV.py | 52 ++++++++++++++++++ .../validate_lora/LTX-2.5-I2AV.py | 53 +++++++++++++++++++ 11 files changed, 166 insertions(+), 22 deletions(-) create mode 100644 examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh rename examples/ltx2/model_training/lora/{LTX-2.5-T2AV-splited-test.sh => LTX-2.5-I2AV-splited.sh} (81%) create mode 100644 examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py create mode 100644 examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index 249c72f8f..60b2e68fd 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -273,13 +273,13 @@ "diffsynth.models.ltx2_dit.LTXModel": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", - "diffsynth.models.ltx2_dit.BasicAVTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_dit.BasicAVTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule", }, "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", - "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextDecoderLayer": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextDecoderLayer": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule", "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", }, diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 46c86d784..61f24dc22 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -298,8 +298,6 @@ def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inpu return inputs_shared, inputs_posi, inputs_nega - - class LTX2AudioVideoUnit_AutoDuration(PipelineUnit): def __init__(self): super().__init__(take_over=True, onload_model_names=("duration_head",)) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py index e3c925769..7e8a99519 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -29,7 +29,7 @@ ], ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512, 768, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index 174f2ef48..182c0b5ed 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -32,7 +32,7 @@ stage2_lora_strength=1.0, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index e0136d4c1..b55113656 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -30,7 +30,7 @@ vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512, 768, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index 8999acd0f..16785e67c 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -33,7 +33,7 @@ vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") -prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" negative_prompt = pipe.default_negative_prompt["LTX-2.5"] height, width, num_frames = 512 * 2, 768 * 2, 121 first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py index 9088fa6eb..80af71a7f 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -5,8 +5,8 @@ vram_config = { "offload_dtype": "disk", "offload_device": "disk", - "onload_dtype": torch.bfloat16, - "onload_device": "cpu", + "onload_dtype": "disk", + "onload_device": "disk", "preparing_dtype": torch.bfloat16, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, diff --git a/examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh new file mode 100644 index 000000000..00ccf3fc8 --- /dev/null +++ b/examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh @@ -0,0 +1,41 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-I2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio,input_image" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-I2AV-full-splited-cache" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch --config_file examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-I2AV-full-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio,input_image" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-I2AV-full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --task "sft:train" diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh b/examples/ltx2/model_training/lora/LTX-2.5-I2AV-splited.sh similarity index 81% rename from examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh rename to examples/ltx2/model_training/lora/LTX-2.5-I2AV-splited.sh index 5fc25d78e..3f53b3a02 100644 --- a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited-test.sh +++ b/examples/ltx2/model_training/lora/LTX-2.5-I2AV-splited.sh @@ -1,11 +1,11 @@ -modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-I2AV-splited/*" --local_dir ./data/diffsynth_example_dataset -# Splited Training (debug: 1 epoch, dataset_repeat 1) +# Splited Training accelerate launch examples/ltx2/model_training/train.py \ - --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited \ - --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited/metadata.csv \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/metadata.csv \ --data_file_keys "video,input_audio" \ - --extra_inputs "input_audio" \ + --extra_inputs "input_audio,input_image" \ --height 512 \ --width 768 \ --num_frames 121 \ @@ -13,9 +13,9 @@ accelerate launch examples/ltx2/model_training/train.py \ --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-4 \ - --num_epochs 1 \ + --num_epochs 5 \ --remove_prefix_in_ckpt "pipe.dit." \ - --output_path "./models/train/LTX2.5-T2AV_lora-splited-cache-test" \ + --output_path "./models/train/LTX2.5-I2AV_lora-splited-cache" \ --lora_base_model "dit" \ --lora_target_modules "to_k,to_q,to_v,to_out.0" \ --lora_rank 32 \ @@ -23,20 +23,20 @@ accelerate launch examples/ltx2/model_training/train.py \ --task "sft:data_process" accelerate launch examples/ltx2/model_training/train.py \ - --dataset_base_path ./models/train/LTX2.5-T2AV_lora-splited-cache-test \ + --dataset_base_path ./models/train/LTX2.5-I2AV_lora-splited-cache \ --data_file_keys "video,input_audio" \ - --extra_inputs "input_audio" \ + --extra_inputs "input_audio,input_image" \ --height 512 \ --width 768 \ --num_frames 121 \ - --dataset_repeat 1 \ + --dataset_repeat 100 \ --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ --learning_rate 1e-4 \ - --num_epochs 1 \ + --num_epochs 5 \ --remove_prefix_in_ckpt "pipe.dit." \ - --output_path "./models/train/LTX2.5-T2AV_lora-test" \ + --output_path "./models/train/LTX2.5-I2AV_lora" \ --lora_base_model "dit" \ --lora_target_modules "to_k,to_q,to_v,to_out.0" \ --lora_rank 32 \ diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py new file mode 100644 index 000000000..41c3360e4 --- /dev/null +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py @@ -0,0 +1,52 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data import VideoData + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(path="./models/train/LTX2.5-I2AV-full/epoch-1.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +image = VideoData("data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/video.mp4", height=height, width=width)[0] +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + cfg_scale=3.0, + input_images=[image], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_i2av_full.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py new file mode 100644 index 000000000..217bea747 --- /dev/null +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py @@ -0,0 +1,53 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data import VideoData + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +pipe.load_lora(pipe.dit, "models/train/LTX2.5-I2AV_lora/epoch-4.safetensors") +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +image = VideoData("data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/video.mp4", height=height, width=width)[0] +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + cfg_scale=3.0, + input_images=[image], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_i2av_lora.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) From 462b35e588dc24fe822072d0d109c47efcb1ebc8 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Thu, 10 Sep 2026 18:35:51 +0800 Subject: [PATCH 30/31] Drop redundant cfg_scale from LTX-2.5 examples to use the pipeline default The LTX2AudioVideoPipeline defaults cfg_scale to 3.0. Remove the explicit cfg_scale=3.0 from the LTX-2.5 inference and validate scripts (no behavior change) and remove cfg_scale=4.0 from the T2AV validate scripts so they also fall back to 3.0. Keep cfg_scale=1.0 in the distilled / INT8-ConvRot / IC-LoRA-Pixel-Spatial-Upscaler scripts, where CFG is intentionally disabled. --- examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py | 1 - examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py | 2 -- examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py | 1 - examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py | 1 - examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py | 1 - examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py | 2 -- examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py | 1 - examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py | 1 - examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py | 1 - examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py | 1 - examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py | 1 - examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py | 1 - 12 files changed, 14 deletions(-) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py index 7e8a99519..e342d0339 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -42,7 +42,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, input_images=[first_frame], input_images_indexes=[0], input_images_strength=1.0, diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py index 182c0b5ed..5c21db200 100644 --- a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -47,7 +47,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame], input_images_indexes=[0], @@ -71,7 +70,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame, last_frame], input_images_indexes=[0, num_frames - 1], diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py index 50754ea0b..7fba1b295 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py @@ -38,7 +38,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py index 2e510be3a..e2234a9c7 100644 --- a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -41,7 +41,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, use_two_stage_pipeline=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py index b55113656..4c937759c 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -43,7 +43,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, input_images=[first_frame], input_images_indexes=[0], input_images_strength=1.0, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py index 16785e67c..5d7dbe15c 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -48,7 +48,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame], input_images_indexes=[0], @@ -72,7 +71,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, use_two_stage_pipeline=True, input_images=[first_frame, last_frame], input_images_indexes=[0, num_frames - 1], diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py index fa33a3f31..7e5227d02 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -39,7 +39,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py index 903c11247..be12094d4 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -42,7 +42,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, use_two_stage_pipeline=True, ) write_video_audio_ltx2( diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py index 41c3360e4..0f5afe452 100644 --- a/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py @@ -38,7 +38,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, input_images=[image], input_images_indexes=[0], input_images_strength=1.0, diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py index abeac27ac..0d338e9f0 100644 --- a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py @@ -35,7 +35,6 @@ width=width, num_frames=num_frames, tiled=True, - cfg_scale=4.0, ) write_video_audio_ltx2( video=video, diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py index 217bea747..50296faf6 100644 --- a/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py @@ -39,7 +39,6 @@ num_frames=num_frames, tiled=True, tile_size_in_frames=80, - cfg_scale=3.0, input_images=[image], input_images_indexes=[0], input_images_strength=1.0, diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py index 3fe59258b..d2932fcba 100644 --- a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py @@ -36,7 +36,6 @@ width=width, num_frames=num_frames, tiled=True, - cfg_scale=4.0, ) write_video_audio_ltx2( video=video, From bda5443f5debd86cd96ef2b4f9e92086691b0981 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Thu, 10 Sep 2026 19:40:15 +0800 Subject: [PATCH 31/31] Point LTX-2.5 T2AV training at its own dataset instead of LTX-2.3's The LTX-2.5 T2AV split-training scripts downloaded and read from ltx2/LTX-2.3-T2AV-splited. Give LTX-2.5 a dedicated example dataset (ltx2/LTX-2.5-T2AV-splited), matching the I2AV scripts, so each model version references its own dataset. --- examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh | 6 +++--- examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh index 019e91e81..6cb69b7c5 100644 --- a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh +++ b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh @@ -1,9 +1,9 @@ -modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset # Splited Training accelerate launch examples/ltx2/model_training/train.py \ - --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited \ - --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited/metadata.csv \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited/metadata.csv \ --data_file_keys "video,input_audio" \ --extra_inputs "input_audio" \ --height 512 \ diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh index 637e1f18f..af873ff64 100644 --- a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh +++ b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh @@ -1,9 +1,9 @@ -modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.3-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset # Splited Training accelerate launch examples/ltx2/model_training/train.py \ - --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited \ - --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.3-T2AV-splited/metadata.csv \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited/metadata.csv \ --data_file_keys "video,input_audio" \ --extra_inputs "input_audio" \ --height 512 \