From 4d7bf35fa347991b8ee35bd84b48570f9edfac04 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Wed, 7 Jan 2026 03:22:48 -0800 Subject: [PATCH] Revert "dit" This reverts commit a6a9c9ca07fbd58a94771b2e8b65e33f48346256. --- fastvideo/__init__.py | 18 +- fastvideo/configs/models/dits/__init__.py | 3 +- fastvideo/configs/models/dits/ltx2.py | 84 - fastvideo/configs/models/vaes/__init__.py | 2 - fastvideo/configs/models/vaes/ltx2vae.py | 43 - fastvideo/layers/rotary_embedding.py | 1 + fastvideo/models/dits/ltx2.py | 1381 ----------------- fastvideo/models/registry.py | 4 +- fastvideo/models/vaes/ltx2vae.py | 90 -- fastvideo/tests/conftest.py | 3 - fastvideo/tests/transformers/test_ltx2.py | 458 ------ fastvideo/tests/vaes/test_ltx2_vae.py | 164 -- .../convert_ltx2_weights.py | 264 ---- 13 files changed, 7 insertions(+), 2508 deletions(-) delete mode 100644 fastvideo/configs/models/dits/ltx2.py delete mode 100644 fastvideo/configs/models/vaes/ltx2vae.py delete mode 100644 fastvideo/models/dits/ltx2.py delete mode 100644 fastvideo/models/vaes/ltx2vae.py delete mode 100644 fastvideo/tests/transformers/test_ltx2.py delete mode 100644 fastvideo/tests/vaes/test_ltx2_vae.py delete mode 100644 scripts/checkpoint_conversion/convert_ltx2_weights.py diff --git a/fastvideo/__init__.py b/fastvideo/__init__.py index 2910084bc..44c7938e9 100644 --- a/fastvideo/__init__.py +++ b/fastvideo/__init__.py @@ -1,16 +1,6 @@ -import os - +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.configs.sample import SamplingParam +from fastvideo.entrypoints.video_generator import VideoGenerator from fastvideo.version import __version__ -_LIGHT_IMPORT = os.getenv("FASTVIDEO_LIGHT_IMPORT", "0") == "1" - -if not _LIGHT_IMPORT: - from fastvideo.configs.pipelines import PipelineConfig - from fastvideo.configs.sample import SamplingParam - from fastvideo.entrypoints.video_generator import VideoGenerator - - __all__ = [ - "VideoGenerator", "PipelineConfig", "SamplingParam", "__version__" - ] -else: - __all__ = ["__version__"] +__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"] diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index c4b21a714..958e9bc71 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -3,12 +3,11 @@ from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config from fastvideo.configs.models.dits.longcat import LongCatVideoConfig -from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.stepvideo import StepVideoConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig", "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", - "LongCatVideoConfig", "LTX2VideoConfig" + "LongCatVideoConfig" ] diff --git a/fastvideo/configs/models/dits/ltx2.py b/fastvideo/configs/models/dits/ltx2.py deleted file mode 100644 index 2c280a38a..000000000 --- a/fastvideo/configs/models/dits/ltx2.py +++ /dev/null @@ -1,84 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LTX-2 Transformer configuration for native FastVideo integration. -""" - -from dataclasses import dataclass, field - -from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig - - -def is_ltx2_blocks(name: str, _module) -> bool: - """FSDP shard condition for LTX-2 transformer blocks.""" - return "transformer_blocks" in name - - -@dataclass -class LTX2VideoArchConfig(DiTArchConfig): - """Architecture configuration for LTX-2 video transformer.""" - - _fsdp_shard_conditions: list = field( - default_factory=lambda: [is_ltx2_blocks]) - _compile_conditions: list = field(default_factory=lambda: [is_ltx2_blocks]) - - # Parameter name mapping for weight conversion (hf/comfy -> FastVideo) - param_names_mapping: dict = field( - default_factory=lambda: { - r"^model\.diffusion_model\.(.*)$": r"model.\1", - r"^diffusion_model\.(.*)$": r"model.\1", - r"^model\.(.*)$": r"model.\1", - r"^(.*)$": r"model.\1", - }) - - reverse_param_names_mapping: dict = field(default_factory=lambda: {}) - lora_param_names_mapping: dict = field(default_factory=lambda: {}) - - # Core transformer settings (defaults from LTX-2 metadata) - num_attention_heads: int = 32 - attention_head_dim: int = 128 - num_layers: int = 48 - cross_attention_dim: int = 4096 - caption_channels: int = 3840 - norm_eps: float = 1e-6 - attention_type: str = "default" - rope_type: str = "split" - double_precision_rope: bool = True - - positional_embedding_theta: float = 10000.0 - positional_embedding_max_pos: list[int] = field( - default_factory=lambda: [20, 2048, 2048]) - timestep_scale_multiplier: int = 1000 - use_middle_indices_grid: bool = True - - # Patchification (video-only path) - patch_size: tuple[int, int, int] = (1, 1, 1) - num_channels_latents: int = 128 - in_channels: int | None = None - out_channels: int | None = None - - # Audio defaults (reserved for joint AV ports) - audio_num_attention_heads: int = 32 - audio_attention_head_dim: int = 64 - audio_in_channels: int = 128 - audio_out_channels: int = 128 - audio_cross_attention_dim: int = 2048 - audio_positional_embedding_max_pos: list[int] = field( - default_factory=lambda: [20]) - av_ca_timestep_scale_multiplier: int = 1 - - def __post_init__(self): - super().__post_init__() - patch_volume = self.patch_size[0] * self.patch_size[ - 1] * self.patch_size[2] - if self.in_channels is None: - self.in_channels = self.num_channels_latents * patch_volume - if self.out_channels is None: - self.out_channels = self.in_channels - - -@dataclass -class LTX2VideoConfig(DiTConfig): - """Main configuration for LTX-2 transformer.""" - - arch_config: DiTArchConfig = field(default_factory=LTX2VideoArchConfig) - prefix: str = "ltx2" diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index 044cda18b..d39254d21 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -1,7 +1,6 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig -from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig from fastvideo.configs.models.vaes.wanvae import WanVAEConfig @@ -11,5 +10,4 @@ __all__ = [ "StepVideoVAEConfig", "CosmosVAEConfig", "Hunyuan15VAEConfig", - "LTX2VAEConfig", ] diff --git a/fastvideo/configs/models/vaes/ltx2vae.py b/fastvideo/configs/models/vaes/ltx2vae.py deleted file mode 100644 index 2bc911acf..000000000 --- a/fastvideo/configs/models/vaes/ltx2vae.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LTX-2 VAE configuration. -""" - -from dataclasses import dataclass, field - -from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig - - -@dataclass -class LTX2VAEArchConfig(VAEArchConfig): - # Mirrors LTX-2 safetensors metadata config under "vae" - _class_name: str = "CausalVideoAutoencoder" - dims: int = 3 - in_channels: int = 3 - out_channels: int = 3 - latent_channels: int = 128 - encoder_blocks: list = field(default_factory=list) - decoder_blocks: list = field(default_factory=list) - patch_size: int = 4 - norm_layer: str = "pixel_norm" - latent_log_var: str = "uniform" - encoder_spatial_padding_mode: str = "zeros" - decoder_spatial_padding_mode: str = "reflect" - causal_decoder: bool = False - timestep_conditioning: bool = True - use_quant_conv: bool = False - scaling_factor: float = 1.0 - normalize_latent_channels: bool = False - - # Match FastVideo naming for compression ratios (LTX-2 default) - temporal_compression_ratio: int = 8 - spatial_compression_ratio: int = 32 - - def __post_init__(self) -> None: - super().__post_init__() - # Keep parent invariants if needed - - -@dataclass -class LTX2VAEConfig(VAEConfig): - arch_config: VAEArchConfig = field(default_factory=LTX2VAEArchConfig) diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index facae95ee..19b02b71d 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -24,6 +24,7 @@ # limitations under the License. """Rotary Positional Embeddings.""" from typing import Any + import torch from fastvideo.distributed.parallel_state import get_sp_group diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py deleted file mode 100644 index f5d0a4b62..000000000 --- a/fastvideo/models/dits/ltx2.py +++ /dev/null @@ -1,1381 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Native LTX-2 transformer implementation (video-only). -""" - -from dataclasses import dataclass, replace -from enum import Enum -import functools -import math -from typing import Any, Optional, Tuple, Callable - -import numpy as np - -import torch -import torch.nn as nn -from einops import rearrange, repeat - -from fastvideo.attention.layer import LocalAttention -from fastvideo.configs.models.dits import LTX2VideoConfig -from fastvideo.models.dits.base import CachableDiT -from fastvideo.platforms import AttentionBackendEnum - - -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: - """Sinusoidal timestep embedding used by LTX-2 AdaLN.""" - if len(timesteps.shape) != 1: - raise ValueError("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): - """Two-layer MLP to project timestep embeddings.""" - 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) - self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) if cond_proj_dim is not None else 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) - self.post_act = None if post_act_fn is None else 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): - """Sinusoidal timestep embedding wrapper with scaling knobs.""" - 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: - return 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, - ) - - -class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): - """PixArt-Alpha timestep embedding used by LTX-2 AdaLN.""" - 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 - - -class AdaLayerNormSingle(torch.nn.Module): - """AdaLN-single modulation that emits scale/shift/gates for LTX-2.""" - def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): - super().__init__() - self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( - embedding_dim, - size_emb_dim=embedding_dim // 3, - ) - self.silu = torch.nn.SiLU() - self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True) - - def forward( - self, - timestep: torch.Tensor, - hidden_dtype: Optional[torch.dtype] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) - return self.linear(self.silu(embedded_timestep)), embedded_timestep - - -class PixArtAlphaTextProjection(torch.nn.Module): - """Caption projection MLP used by LTX-2.""" - def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): - super().__init__() - if out_features is None: - out_features = hidden_size - self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) - if act_fn == "gelu_tanh": - self.act_1 = torch.nn.GELU(approximate="tanh") - elif act_fn == "silu": - self.act_1 = torch.nn.SiLU() - else: - raise ValueError(f"Unknown activation function: {act_fn}") - self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) - - def forward(self, caption: torch.Tensor) -> torch.Tensor: - hidden_states = self.linear_1(caption) - hidden_states = self.act_1(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -class GELUApprox(nn.Module): - """Linear + tanh-approximate GELU used by LTX-2 FFN.""" - def __init__(self, in_features: int, out_features: int): - super().__init__() - self.proj = nn.Linear(in_features, out_features) - self.act = nn.GELU(approximate="tanh") - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.act(self.proj(x)) - - -class FeedForward(nn.Module): - """LTX-2 FFN: GELUApprox -> Identity -> Linear.""" - def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: - super().__init__() - inner_dim = int(dim * mult) - project_in = GELUApprox(dim, inner_dim) - self.net = nn.Sequential(project_in, nn.Identity(), nn.Linear(inner_dim, dim_out)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class VideoLatentShape(tuple): - """Helper for (B, C, T, H, W) latent shapes.""" - @property - def batch(self) -> int: - return self[0] - - @property - def channels(self) -> int: - return self[1] - - @property - def frames(self) -> int: - return self[2] - - @property - def height(self) -> int: - return self[3] - - @property - def width(self) -> int: - return self[4] - - @staticmethod - def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": - return VideoLatentShape(shape) - - def to_torch_shape(self) -> torch.Size: - return torch.Size(self) - - -class VideoLatentPatchifier: - """Patchify/unpatchify latent tokens for LTX-2.""" - def __init__(self, patch_size: int): - self._patch_size = (1, patch_size, patch_size) - - @property - def patch_size(self) -> Tuple[int, int, int]: - return self._patch_size - - def get_token_count(self, tgt_shape: VideoLatentShape) -> int: - return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size) - - def patchify(self, latents: torch.Tensor) -> torch.Tensor: - return rearrange( - latents, - "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", - p1=self._patch_size[0], - p2=self._patch_size[1], - p3=self._patch_size[2], - ) - - def unpatchify(self, latents: torch.Tensor, output_shape: VideoLatentShape) -> torch.Tensor: - patch_grid_frames = output_shape.frames // self._patch_size[0] - patch_grid_height = output_shape.height // self._patch_size[1] - patch_grid_width = output_shape.width // self._patch_size[2] - return rearrange( - latents, - "b (f h w) (c p q) -> b c f (h p) (w q)", - f=patch_grid_frames, - h=patch_grid_height, - w=patch_grid_width, - p=self._patch_size[1], - q=self._patch_size[2], - ) - - def get_patch_grid_bounds( - self, - output_shape: VideoLatentShape, - device: Optional[torch.device] = None, - ) -> torch.Tensor: - frames = output_shape.frames - height = output_shape.height - width = output_shape.width - batch_size = output_shape.batch - - grid_coords = torch.meshgrid( - torch.arange(start=0, end=frames, step=self._patch_size[0], device=device), - torch.arange(start=0, end=height, step=self._patch_size[1], device=device), - torch.arange(start=0, end=width, step=self._patch_size[2], device=device), - indexing="ij", - ) - - patch_starts = torch.stack(grid_coords, dim=0) - patch_size_delta = torch.tensor( - self._patch_size, - device=patch_starts.device, - dtype=patch_starts.dtype, - ).view(3, 1, 1, 1) - patch_ends = patch_starts + patch_size_delta - latent_coords = torch.stack((patch_starts, patch_ends), dim=-1) - latent_coords = repeat( - latent_coords, - "c f h w bounds -> b c (f h w) bounds", - b=batch_size, - bounds=2, - ) - return latent_coords - - -class LTXRopeType(Enum): - """LTX-2 rotary variants (interleaved vs split).""" - INTERLEAVED = "interleaved" - SPLIT = "split" - - -def apply_ltx_rotary_emb( - input_tensor: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, -) -> torch.Tensor: - """Apply LTX-2 rotary embeddings to a tensor.""" - if rope_type == LTXRopeType.INTERLEAVED: - return _apply_ltx_interleaved_rotary_emb(input_tensor, *freqs_cis) - if rope_type == LTXRopeType.SPLIT: - return _apply_ltx_split_rotary_emb(input_tensor, *freqs_cis) - raise ValueError(f"Invalid rope type: {rope_type}") - - -def _apply_ltx_interleaved_rotary_emb( - input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor -) -> torch.Tensor: - t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) - t1, t2 = t_dup.unbind(dim=-1) - t_dup = torch.stack((-t2, t1), dim=-1) - input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") - return input_tensor * cos_freqs + input_tensor_rot * sin_freqs - - -def _apply_ltx_split_rotary_emb( - input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor -) -> torch.Tensor: - needs_reshape = False - if input_tensor.ndim != 4 and cos_freqs.ndim == 4: - b, h, t, _ = cos_freqs.shape - input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2) - needs_reshape = True - - split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) - first_half_input = split_input[..., :1, :] - second_half_input = split_input[..., 1:, :] - - output = split_input * cos_freqs.unsqueeze(-2) - first_half_output = output[..., :1, :] - second_half_output = output[..., 1:, :] - - first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) - second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) - - output = rearrange(output, "... d r -> ... (d r)") - if needs_reshape: - output = output.swapaxes(1, 2).reshape(b, t, -1) - return output - - -@functools.lru_cache(maxsize=5) -def generate_ltx_freq_grid_np( - positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int -) -> torch.Tensor: - """Generate LTX-2 rotary frequencies with high-precision numpy.""" - theta = positional_embedding_theta - start = 1 - end = theta - n_elem = 2 * positional_embedding_max_pos_count - pow_indices = np.power( - theta, - np.linspace( - np.log(start) / np.log(theta), - np.log(end) / np.log(theta), - inner_dim // n_elem, - dtype=np.float64, - ), - ) - return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32) - - -@functools.lru_cache(maxsize=5) -def generate_ltx_freq_grid_pytorch( - positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int -) -> torch.Tensor: - """Generate LTX-2 rotary frequencies in torch for speed.""" - theta = positional_embedding_theta - start = 1 - end = theta - n_elem = 2 * positional_embedding_max_pos_count - indices = theta ** ( - torch.linspace( - math.log(start, theta), - math.log(end, theta), - inner_dim // n_elem, - dtype=torch.float32, - ) - ) - indices = indices.to(dtype=torch.float32) - return indices * math.pi / 2 - - -def _ltx_get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor: - n_pos_dims = indices_grid.shape[1] - if n_pos_dims != len(max_pos): - raise ValueError( - f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})" - ) - return torch.stack( - [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], - dim=-1, - ) - - -def _ltx_generate_freqs( - indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool -) -> torch.Tensor: - if use_middle_indices_grid: - indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1] - indices_grid = (indices_grid_start + indices_grid_end) / 2.0 - elif len(indices_grid.shape) == 4: - indices_grid = indices_grid[..., 0] - - fractional_positions = _ltx_get_fractional_positions(indices_grid, max_pos) - indices = indices.to(device=fractional_positions.device) - freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2) - return freqs - - -def _ltx_split_freqs_cis( - freqs: torch.Tensor, pad_size: int, num_attention_heads: int -) -> tuple[torch.Tensor, torch.Tensor]: - cos_freq = freqs.cos() - sin_freq = freqs.sin() - - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) - sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) - cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) - sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) - - b = cos_freq.shape[0] - t = cos_freq.shape[1] - cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1) - sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1) - cos_freq = torch.swapaxes(cos_freq, 1, 2) - sin_freq = torch.swapaxes(sin_freq, 1, 2) - return cos_freq, sin_freq - - -def _ltx_interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]: - cos_freq = freqs.cos().repeat_interleave(2, dim=-1) - sin_freq = freqs.sin().repeat_interleave(2, dim=-1) - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) - sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size]) - cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) - sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) - return cos_freq, sin_freq - - -def precompute_ltx_freqs_cis( - indices_grid: torch.Tensor, - dim: int, - out_dtype: torch.dtype, - theta: float = 10000.0, - max_pos: list[int] | None = None, - use_middle_indices_grid: bool = False, - num_attention_heads: int = 32, - rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, - freq_grid_generator: Callable[[float, int, int], torch.Tensor] = generate_ltx_freq_grid_pytorch, -) -> tuple[torch.Tensor, torch.Tensor]: - """Precompute LTX-2 rotary cos/sin grids for (t, x, y) positions.""" - if max_pos is None: - max_pos = [20, 2048, 2048] - - indices = freq_grid_generator(theta, indices_grid.shape[1], dim) - freqs = _ltx_generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid) - - if rope_type == LTXRopeType.SPLIT: - expected_freqs = dim // 2 - current_freqs = freqs.shape[-1] - pad_size = expected_freqs - current_freqs - cos_freq, sin_freq = _ltx_split_freqs_cis(freqs, pad_size, num_attention_heads) - else: - n_elem = 2 * indices_grid.shape[1] - cos_freq, sin_freq = _ltx_interleaved_freqs_cis(freqs, dim % n_elem) - return cos_freq.to(out_dtype), sin_freq.to(out_dtype) - - -@dataclass(frozen=True) -class TransformerArgs: - """Pack transformer inputs for LTX-2 blocks.""" - x: torch.Tensor - context: torch.Tensor - context_mask: torch.Tensor | None - timesteps: torch.Tensor - embedded_timestep: torch.Tensor - positional_embeddings: torch.Tensor - cross_positional_embeddings: torch.Tensor | None - cross_scale_shift_timestep: torch.Tensor | None - cross_gate_timestep: torch.Tensor | None - enabled: bool - - -@dataclass(frozen=True) -class Modality: - """Lightweight modality container for LTX-2 inputs.""" - enabled: bool - latent: torch.Tensor - timesteps: torch.Tensor - positions: torch.Tensor - context: torch.Tensor - context_mask: torch.Tensor | None = None - - -class TransformerArgsPreprocessor: - """Prepare LTX-2 transformer inputs (patchify, AdaLN, rope).""" - def __init__( # noqa: PLR0913 - self, - patchify_proj: torch.nn.Linear, - adaln: AdaLayerNormSingle, - caption_projection: PixArtAlphaTextProjection, - inner_dim: int, - max_pos: list[int], - num_attention_heads: int, - use_middle_indices_grid: bool, - timestep_scale_multiplier: int, - double_precision_rope: bool, - positional_embedding_theta: float, - rope_type: LTXRopeType, - ) -> None: - self.patchify_proj = patchify_proj - self.adaln = adaln - self.caption_projection = caption_projection - self.inner_dim = inner_dim - self.max_pos = max_pos - self.num_attention_heads = num_attention_heads - self.use_middle_indices_grid = use_middle_indices_grid - self.timestep_scale_multiplier = timestep_scale_multiplier - self.double_precision_rope = double_precision_rope - self.positional_embedding_theta = positional_embedding_theta - self.rope_type = rope_type - - def _prepare_timestep( - self, timestep: torch.Tensor, batch_size: int, hidden_dtype: torch.dtype - ) -> tuple[torch.Tensor, torch.Tensor]: - timestep = timestep * self.timestep_scale_multiplier - timestep, embedded_timestep = self.adaln(timestep.flatten(), hidden_dtype=hidden_dtype) - timestep = timestep.view(batch_size, -1, timestep.shape[-1]) - embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1]) - return timestep, embedded_timestep - - def _prepare_context( - self, - context: torch.Tensor, - x: torch.Tensor, - attention_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - batch_size = x.shape[0] - context = self.caption_projection(context) - context = context.view(batch_size, -1, x.shape[-1]) - return context, attention_mask - - def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: - if attention_mask is None or torch.is_floating_point(attention_mask): - return attention_mask - return (attention_mask - 1).to(x_dtype).reshape( - (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) - ) * torch.finfo(x_dtype).max - - def _prepare_positional_embeddings( - self, - positions: torch.Tensor, - inner_dim: int, - max_pos: list[int], - use_middle_indices_grid: bool, - num_attention_heads: int, - x_dtype: torch.dtype, - ) -> torch.Tensor: - if self.double_precision_rope: - freq_grid_generator = generate_ltx_freq_grid_np - else: - freq_grid_generator = generate_ltx_freq_grid_pytorch - return precompute_ltx_freqs_cis( - positions, - dim=inner_dim, - out_dtype=x_dtype, - theta=self.positional_embedding_theta, - max_pos=max_pos, - use_middle_indices_grid=use_middle_indices_grid, - num_attention_heads=num_attention_heads, - rope_type=self.rope_type, - freq_grid_generator=freq_grid_generator, - ) - - def prepare(self, modality: Modality) -> TransformerArgs: - x = self.patchify_proj(modality.latent) - timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype) - context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask) - attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype) - pe = self._prepare_positional_embeddings( - positions=modality.positions, - inner_dim=self.inner_dim, - max_pos=self.max_pos, - use_middle_indices_grid=self.use_middle_indices_grid, - num_attention_heads=self.num_attention_heads, - x_dtype=modality.latent.dtype, - ) - return TransformerArgs( - x=x, - context=context, - context_mask=attention_mask, - timesteps=timestep, - embedded_timestep=embedded_timestep, - positional_embeddings=pe, - cross_positional_embeddings=None, - cross_scale_shift_timestep=None, - cross_gate_timestep=None, - enabled=modality.enabled, - ) - - -class MultiModalTransformerArgsPreprocessor: - """Prepare transformer args for audio/video cross-attention paths.""" - - def __init__( # noqa: PLR0913 - self, - patchify_proj: torch.nn.Linear, - adaln: AdaLayerNormSingle, - caption_projection: PixArtAlphaTextProjection, - cross_scale_shift_adaln: AdaLayerNormSingle, - cross_gate_adaln: AdaLayerNormSingle, - inner_dim: int, - max_pos: list[int], - num_attention_heads: int, - cross_pe_max_pos: int, - use_middle_indices_grid: bool, - audio_cross_attention_dim: int, - timestep_scale_multiplier: int, - double_precision_rope: bool, - positional_embedding_theta: float, - rope_type: LTXRopeType, - av_ca_timestep_scale_multiplier: int, - ) -> None: - self.simple_preprocessor = TransformerArgsPreprocessor( - patchify_proj=patchify_proj, - adaln=adaln, - caption_projection=caption_projection, - inner_dim=inner_dim, - max_pos=max_pos, - num_attention_heads=num_attention_heads, - use_middle_indices_grid=use_middle_indices_grid, - timestep_scale_multiplier=timestep_scale_multiplier, - double_precision_rope=double_precision_rope, - positional_embedding_theta=positional_embedding_theta, - rope_type=rope_type, - ) - 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 - - def prepare(self, modality: Modality) -> TransformerArgs: - transformer_args = self.simple_preprocessor.prepare(modality) - cross_pe = self.simple_preprocessor._prepare_positional_embeddings( - positions=modality.positions[:, 0:1, :], - inner_dim=self.audio_cross_attention_dim, - max_pos=[self.cross_pe_max_pos], - use_middle_indices_grid=True, - num_attention_heads=self.simple_preprocessor.num_attention_heads, - x_dtype=modality.latent.dtype, - ) - - cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( - timestep=modality.timesteps, - timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, - batch_size=transformer_args.x.shape[0], - hidden_dtype=modality.latent.dtype, - ) - - return replace( - transformer_args, - cross_positional_embeddings=cross_pe, - cross_scale_shift_timestep=cross_scale_shift_timestep, - cross_gate_timestep=cross_gate_timestep, - ) - - def _prepare_cross_attention_timestep( - self, - timestep: torch.Tensor, - timestep_scale_multiplier: int, - batch_size: int, - hidden_dtype: torch.dtype, - ) -> tuple[torch.Tensor, torch.Tensor]: - timestep = timestep * timestep_scale_multiplier - - av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier - - scale_shift_timestep, _ = self.cross_scale_shift_adaln( - 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( - timestep.flatten() * av_ca_factor, - hidden_dtype=hidden_dtype, - ) - gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) - - return scale_shift_timestep, gate_noise_timestep - - -@dataclass -class TransformerConfig: - """Attention/FFN dims for LTX-2 transformer blocks.""" - dim: int - heads: int - d_head: int - context_dim: int - - -class LTXSelfAttention(nn.Module): - """LTX-2 attention block with RMSNorm + FastVideo LocalAttention.""" - def __init__( - self, - query_dim: int, - context_dim: int | None, - heads: int, - dim_head: int, - norm_eps: float, - rope_type: LTXRopeType, - supported_attention_backends: tuple[AttentionBackendEnum, ...], - ) -> None: - super().__init__() - inner_dim = dim_head * heads - context_dim = query_dim if context_dim is None else context_dim - - self.heads = heads - self.dim_head = dim_head - self.rope_type = rope_type - - self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) - self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) - self.to_q = nn.Linear(query_dim, inner_dim, bias=True) - self.to_k = nn.Linear(context_dim, inner_dim, bias=True) - self.to_v = nn.Linear(context_dim, inner_dim, bias=True) - self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) - - self.attn = LocalAttention( - num_heads=heads, - head_size=dim_head, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=supported_attention_backends, - ) - - def forward( - self, - x: torch.Tensor, - context: torch.Tensor | None = None, - mask: torch.Tensor | None = None, - pe: tuple[torch.Tensor, torch.Tensor] | None = None, - k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> torch.Tensor: - q = self.to_q(x) - context = x if context is None else context - k = self.to_k(context) - v = self.to_v(context) - - q = self.q_norm(q) - k = self.k_norm(k) - - if pe is not None: - q = apply_ltx_rotary_emb(q, pe, self.rope_type) - k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) - - b, q_len, _ = q.shape - k_len = k.shape[1] - q = q.view(b, q_len, self.heads, self.dim_head) - k = k.view(b, k_len, self.heads, self.dim_head) - v = v.view(b, k_len, self.heads, self.dim_head) - out = self.attn(q, k, v) - out = out.view(b, q_len, -1) - return self.to_out(out) - - -class BasicAVTransformerBlock(torch.nn.Module): - """LTX-2 transformer block (audio/video + cross attention + FFN).""" - - def __init__( - self, - idx: int, - video: TransformerConfig | None = None, - audio: TransformerConfig | None = None, - rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, - norm_eps: float = 1e-6, - ): - super().__init__() - self.idx = idx - - if video is not None: - self.attn1 = LTXSelfAttention( - query_dim=video.dim, - context_dim=None, - heads=video.heads, - dim_head=video.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), - ) - self.attn2 = LTXSelfAttention( - query_dim=video.dim, - context_dim=video.context_dim, - heads=video.heads, - dim_head=video.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), - ) - self.ff = FeedForward(video.dim, dim_out=video.dim) - self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim)) - - if audio is not None: - self.audio_attn1 = LTXSelfAttention( - query_dim=audio.dim, - context_dim=None, - heads=audio.heads, - dim_head=audio.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), - ) - self.audio_attn2 = LTXSelfAttention( - query_dim=audio.dim, - context_dim=audio.context_dim, - heads=audio.heads, - dim_head=audio.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), - ) - self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) - self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim)) - - if audio is not None and video is not None: - self.audio_to_video_attn = LTXSelfAttention( - query_dim=video.dim, - context_dim=audio.dim, - heads=audio.heads, - dim_head=audio.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), - ) - self.video_to_audio_attn = LTXSelfAttention( - query_dim=audio.dim, - context_dim=video.dim, - heads=audio.heads, - dim_head=audio.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), - ) - self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim)) - self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim)) - - self.norm_eps = norm_eps - - def get_ada_values( - self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice - ) -> tuple[torch.Tensor, ...]: - num_ada_params = scale_shift_table.shape[0] - ada_values = ( - scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) - + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] - ).unbind(dim=2) - return ada_values - - def get_av_ca_ada_values( - self, - scale_shift_table: torch.Tensor, - batch_size: int, - scale_shift_timestep: torch.Tensor, - gate_timestep: torch.Tensor, - num_scale_shift_values: int = 4, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - scale_shift_ada_values = self.get_ada_values( - scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None) - ) - gate_ada_values = self.get_ada_values( - scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None) - ) - - scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values] - gate_ada_values = [t.squeeze(2) for t in gate_ada_values] - - return (*scale_shift_chunks, *gate_ada_values) - - def forward( - self, - video: TransformerArgs | None, - audio: TransformerArgs | None, - ) -> tuple[TransformerArgs | None, TransformerArgs | None]: - vx = video.x if video is not None else None - ax = audio.x if audio is not None else None - - run_vx = video is not None and video.enabled and vx.numel() > 0 - run_ax = audio is not None and audio.enabled and ax.numel() > 0 - run_a2v = run_vx and run_ax - run_v2a = run_ax and run_vx - - if run_vx: - vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( - self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3) - ) - norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa - vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa - vx = vx + self.attn2( - torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps), - context=video.context, - mask=video.context_mask, - ) - - if run_ax: - ashift_msa, ascale_msa, agate_msa = self.get_ada_values( - self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) - ) - norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa - ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa - ax = ax + self.audio_attn2( - torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps), - context=audio.context, - mask=audio.context_mask, - ) - - if run_a2v or run_v2a: - vx_norm3 = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) - ax_norm3 = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) - - ( - scale_ca_audio_hidden_states_a2v, - shift_ca_audio_hidden_states_a2v, - scale_ca_audio_hidden_states_v2a, - shift_ca_audio_hidden_states_v2a, - gate_out_v2a, - ) = self.get_av_ca_ada_values( - self.scale_shift_table_a2v_ca_audio, - ax.shape[0], - audio.cross_scale_shift_timestep, - audio.cross_gate_timestep, - ) - - ( - scale_ca_video_hidden_states_a2v, - shift_ca_video_hidden_states_a2v, - scale_ca_video_hidden_states_v2a, - shift_ca_video_hidden_states_v2a, - gate_out_a2v, - ) = self.get_av_ca_ada_values( - self.scale_shift_table_a2v_ca_video, - vx.shape[0], - video.cross_scale_shift_timestep, - video.cross_gate_timestep, - ) - - if run_a2v: - vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v - ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v - vx = vx + ( - self.audio_to_video_attn( - vx_scaled, - context=ax_scaled, - pe=video.cross_positional_embeddings, - k_pe=audio.cross_positional_embeddings, - ) - * gate_out_a2v - ) - - if run_v2a: - ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a - vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a - ax = ax + ( - self.video_to_audio_attn( - ax_scaled, - context=vx_scaled, - pe=audio.cross_positional_embeddings, - k_pe=video.cross_positional_embeddings, - ) - * gate_out_v2a - ) - - if run_vx: - vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( - self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, None) - ) - vx_scaled = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp - vx = vx + self.ff(vx_scaled) * vgate_mlp - - if run_ax: - ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( - self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, None) - ) - ax_scaled = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp - ax = ax + self.audio_ff(ax_scaled) * agate_mlp - - return ( - replace(video, x=vx) if video is not None else None, - replace(audio, x=ax) if audio is not None else None, - ) - - -class LTXModelType(Enum): - """Model type flags for LTX-2.""" - AudioVideo = "ltx av model" - VideoOnly = "ltx video only model" - AudioOnly = "ltx audio only model" - - def is_video_enabled(self) -> bool: - return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly) - - def is_audio_enabled(self) -> bool: - return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly) - - -class LTXModel(torch.nn.Module): - """Core LTX-2 transformer stack for audio/video latents.""" - - def __init__( # noqa: PLR0913 - self, - *, - model_type: LTXModelType = LTXModelType.AudioVideo, - num_attention_heads: int = 32, - attention_head_dim: int = 128, - in_channels: int = 128, - out_channels: int = 128, - num_layers: int = 48, - cross_attention_dim: int = 4096, - norm_eps: float = 1e-06, - caption_channels: int = 3840, - positional_embedding_theta: float = 10000.0, - positional_embedding_max_pos: list[int] | None = None, - timestep_scale_multiplier: int = 1000, - use_middle_indices_grid: bool = True, - audio_num_attention_heads: int = 32, - audio_attention_head_dim: int = 64, - audio_in_channels: int = 128, - audio_out_channels: int = 128, - audio_cross_attention_dim: int = 2048, - audio_positional_embedding_max_pos: list[int] | None = None, - av_ca_timestep_scale_multiplier: int = 1, - rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, - double_precision_rope: bool = False, - ): - super().__init__() - self._enable_gradient_checkpointing = False - self.use_middle_indices_grid = use_middle_indices_grid - self.rope_type = rope_type - self.double_precision_rope = double_precision_rope - self.timestep_scale_multiplier = timestep_scale_multiplier - self.positional_embedding_theta = positional_embedding_theta - self.model_type = model_type - 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] - self.positional_embedding_max_pos = positional_embedding_max_pos - self.num_attention_heads = num_attention_heads - self.inner_dim = num_attention_heads * attention_head_dim - self._init_video( - in_channels=in_channels, - out_channels=out_channels, - caption_channels=caption_channels, - norm_eps=norm_eps, - ) - - if model_type.is_audio_enabled(): - if audio_positional_embedding_max_pos is None: - audio_positional_embedding_max_pos = [20] - self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos - self.audio_num_attention_heads = audio_num_attention_heads - self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim - self._init_audio( - in_channels=audio_in_channels, - out_channels=audio_out_channels, - caption_channels=caption_channels, - norm_eps=norm_eps, - ) - - 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) - self._init_transformer_blocks( - num_layers=num_layers, - attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0, - cross_attention_dim=cross_attention_dim, - audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0, - audio_cross_attention_dim=audio_cross_attention_dim, - norm_eps=norm_eps, - ) - - def _init_video( - self, - in_channels: int, - out_channels: int, - caption_channels: int, - norm_eps: float, - ) -> None: - self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) - self.adaln_single = AdaLayerNormSingle(self.inner_dim) - self.caption_projection = PixArtAlphaTextProjection( - in_features=caption_channels, - hidden_size=self.inner_dim, - ) - self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim)) - self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps) - self.proj_out = torch.nn.Linear(self.inner_dim, out_channels) - - def _init_audio( - self, - in_channels: int, - out_channels: int, - caption_channels: int, - norm_eps: float, - ) -> None: - self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) - self.audio_adaln_single = AdaLayerNormSingle(self.audio_inner_dim) - self.audio_caption_projection = PixArtAlphaTextProjection( - in_features=caption_channels, - hidden_size=self.audio_inner_dim, - ) - self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim)) - self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps) - self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels) - - def _init_audio_video(self, num_scale_shift_values: int) -> None: - self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle( - self.inner_dim, - embedding_coefficient=num_scale_shift_values, - ) - self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, - embedding_coefficient=num_scale_shift_values, - ) - self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle( - self.inner_dim, - embedding_coefficient=1, - ) - self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, - embedding_coefficient=1, - ) - - def _init_preprocessors(self, cross_pe_max_pos: int | None = None) -> None: - if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): - self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor( - patchify_proj=self.patchify_proj, - adaln=self.adaln_single, - caption_projection=self.caption_projection, - cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single, - cross_gate_adaln=self.av_ca_a2v_gate_adaln_single, - inner_dim=self.inner_dim, - max_pos=self.positional_embedding_max_pos, - num_attention_heads=self.num_attention_heads, - cross_pe_max_pos=cross_pe_max_pos, - use_middle_indices_grid=self.use_middle_indices_grid, - audio_cross_attention_dim=self.audio_cross_attention_dim, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, - ) - self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( - patchify_proj=self.audio_patchify_proj, - adaln=self.audio_adaln_single, - caption_projection=self.audio_caption_projection, - cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single, - cross_gate_adaln=self.av_ca_v2a_gate_adaln_single, - inner_dim=self.audio_inner_dim, - max_pos=self.audio_positional_embedding_max_pos, - num_attention_heads=self.audio_num_attention_heads, - cross_pe_max_pos=cross_pe_max_pos, - use_middle_indices_grid=self.use_middle_indices_grid, - audio_cross_attention_dim=self.audio_cross_attention_dim, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, - ) - elif self.model_type.is_video_enabled(): - self.video_args_preprocessor = TransformerArgsPreprocessor( - patchify_proj=self.patchify_proj, - adaln=self.adaln_single, - caption_projection=self.caption_projection, - inner_dim=self.inner_dim, - max_pos=self.positional_embedding_max_pos, - num_attention_heads=self.num_attention_heads, - use_middle_indices_grid=self.use_middle_indices_grid, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - ) - elif self.model_type.is_audio_enabled(): - self.audio_args_preprocessor = TransformerArgsPreprocessor( - patchify_proj=self.audio_patchify_proj, - adaln=self.audio_adaln_single, - caption_projection=self.audio_caption_projection, - inner_dim=self.audio_inner_dim, - max_pos=self.audio_positional_embedding_max_pos, - num_attention_heads=self.audio_num_attention_heads, - use_middle_indices_grid=self.use_middle_indices_grid, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - ) - - def _init_transformer_blocks( - self, - num_layers: int, - attention_head_dim: int, - cross_attention_dim: int, - audio_attention_head_dim: int, - audio_cross_attention_dim: int, - norm_eps: float, - ) -> None: - video_config = ( - TransformerConfig( - dim=self.inner_dim, - heads=self.num_attention_heads, - d_head=attention_head_dim, - context_dim=cross_attention_dim, - ) - if self.model_type.is_video_enabled() - else None - ) - audio_config = ( - TransformerConfig( - dim=self.audio_inner_dim, - heads=self.audio_num_attention_heads, - d_head=audio_attention_head_dim, - context_dim=audio_cross_attention_dim, - ) - if self.model_type.is_audio_enabled() - else None - ) - self.transformer_blocks = torch.nn.ModuleList( - [ - BasicAVTransformerBlock( - idx=idx, - video=video_config, - audio=audio_config, - rope_type=self.rope_type, - norm_eps=norm_eps, - ) - for idx in range(num_layers) - ] - ) - - def _process_transformer_blocks( - self, - video: TransformerArgs | None, - audio: TransformerArgs | None, - ) -> tuple[TransformerArgs | None, TransformerArgs | None]: - for block in self.transformer_blocks: - video, audio = block(video=video, audio=audio) - return video, audio - - def _process_output( - self, - scale_shift_table: torch.Tensor, - norm_out: torch.nn.LayerNorm, - proj_out: torch.nn.Linear, - x: torch.Tensor, - embedded_timestep: torch.Tensor, - ) -> torch.Tensor: - scale_shift_values = ( - scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None] - ) - shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] - x = norm_out(x) - x = x * (1 + scale) + shift - x = proj_out(x) - return x - - def forward( - self, - video: Modality | None, - audio: Modality | None, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - if not self.model_type.is_video_enabled() and video is not None: - raise ValueError("Video is not enabled for this model") - if not self.model_type.is_audio_enabled() and audio is not None: - raise ValueError("Audio is not enabled for this model") - - video_args = self.video_args_preprocessor.prepare(video) if video is not None else None - audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None - video_out, audio_out = self._process_transformer_blocks(video_args, audio_args) - - vx = ( - self._process_output( - self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep - ) - if video_out is not None - else None - ) - ax = ( - self._process_output( - self.audio_scale_shift_table, - self.audio_norm_out, - self.audio_proj_out, - audio_out.x, - audio_out.embedded_timestep, - ) - if audio_out is not None - else None - ) - return vx, ax - - -class LTX2Transformer3DModel(CachableDiT): - """ - LTX-2 transformer using native FastVideo LTX-2 modules. - """ - - param_names_mapping = LTX2VideoConfig().param_names_mapping - reverse_param_names_mapping = LTX2VideoConfig().reverse_param_names_mapping - lora_param_names_mapping = LTX2VideoConfig().lora_param_names_mapping - - def __init__(self, config: LTX2VideoConfig, hf_config: dict[str, Any]): - super().__init__(config=config, hf_config=hf_config) - - arch = config.arch_config - - model_type = LTXModelType.AudioVideo - self.model = LTXModel( - model_type=model_type, - num_attention_heads=arch.num_attention_heads, - attention_head_dim=arch.attention_head_dim, - in_channels=arch.in_channels, - out_channels=arch.out_channels, - num_layers=arch.num_layers, - cross_attention_dim=arch.cross_attention_dim, - norm_eps=arch.norm_eps, - caption_channels=arch.caption_channels, - positional_embedding_theta=arch.positional_embedding_theta, - positional_embedding_max_pos=arch.positional_embedding_max_pos, - timestep_scale_multiplier=arch.timestep_scale_multiplier, - use_middle_indices_grid=arch.use_middle_indices_grid, - rope_type=LTXRopeType(arch.rope_type), - double_precision_rope=arch.double_precision_rope, - audio_num_attention_heads=arch.audio_num_attention_heads, - audio_attention_head_dim=arch.audio_attention_head_dim, - audio_in_channels=arch.audio_in_channels, - audio_out_channels=arch.audio_out_channels, - audio_cross_attention_dim=arch.audio_cross_attention_dim, - audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos, - av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier, - ) - - self.patchifier = VideoLatentPatchifier(patch_size=arch.patch_size[1]) - - self.hidden_size = arch.num_attention_heads * arch.attention_head_dim - self.num_attention_heads = arch.num_attention_heads - self.num_channels_latents = arch.num_channels_latents - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, - guidance=None, - **kwargs, - ) -> torch.Tensor: - if isinstance(encoder_hidden_states, list): - encoder_hidden_states = encoder_hidden_states[0] - - video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) - positions = self.patchifier.get_patch_grid_bounds( - video_shape, device=hidden_states.device - ) - latents = self.patchifier.patchify(hidden_states) - video_modality = Modality( - enabled=True, - latent=latents, - timesteps=timestep, - positions=positions, - context=encoder_hidden_states, - context_mask=None, - ) - video_out, _ = self.model( - video=video_modality, - audio=None, - ) - return self.patchifier.unpatchify(video_out, output_shape=video_shape) diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 2458eda27..cbe2685a6 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -32,7 +32,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), - "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"), } _IMAGE_TO_VIDEO_DIT_MODELS = { @@ -64,8 +63,7 @@ _VAE_MODELS = { ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"), "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"), - "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"), - "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"), + "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo") } _SCHEDULERS = { diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py deleted file mode 100644 index 25d02b396..000000000 --- a/fastvideo/models/vaes/ltx2vae.py +++ /dev/null @@ -1,90 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LTX-2 video VAE wrappers using the official ltx-core implementation. -""" - -from pathlib import Path -import sys -from typing import Any - -import torch -import torch.nn as nn - -from fastvideo.models.vaes.common import DiagonalGaussianDistribution - - -def _require_ltx2(): - repo_root = Path(__file__).resolve().parents[3] - local_core = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" - if local_core.exists() and str(local_core) not in sys.path: - sys.path.insert(0, str(local_core)) - - try: - from ltx_core.model.video_vae import ( # type: ignore - VideoDecoder, - VideoDecoderConfigurator, - VideoEncoder, - VideoEncoderConfigurator, - ) - except ImportError as exc: - raise ImportError( - "LTX-2 sources are required. Ensure FastVideo/LTX-2 is present." - ) from exc - - return VideoDecoder, VideoDecoderConfigurator, VideoEncoder, VideoEncoderConfigurator - - -class LTX2VideoEncoder(nn.Module): - def __init__(self, config: dict[str, Any]): - super().__init__() - VideoDecoder, VideoDecoderConfigurator, VideoEncoder, VideoEncoderConfigurator = _require_ltx2() - self.model: VideoEncoder = VideoEncoderConfigurator.from_config(config) - - def forward(self, sample: torch.Tensor) -> torch.Tensor: - return self.model(sample) - - -class LTX2VideoDecoder(nn.Module): - def __init__(self, config: dict[str, Any]): - super().__init__() - VideoDecoder, VideoDecoderConfigurator, VideoEncoder, VideoEncoderConfigurator = _require_ltx2() - self.model: VideoDecoder = VideoDecoderConfigurator.from_config(config) - - def forward( - self, - sample: torch.Tensor, - timestep: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - return self.model(sample, timestep=timestep, generator=generator) - - -class LTX2CausalVideoAutoencoder(nn.Module): - """ - LTX-2 VAE wrapper that exposes FastVideo's VAE encode/decode interface. - """ - - def __init__(self, config: dict[str, Any]): - super().__init__() - self.config = config - VideoDecoder, VideoDecoderConfigurator, VideoEncoder, VideoEncoderConfigurator = _require_ltx2() - - self.encoder = VideoEncoderConfigurator.from_config(config) - self.decoder = VideoDecoderConfigurator.from_config(config) - - def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: - means = self.encoder(x) - zeros = torch.zeros_like(means) - return DiagonalGaussianDistribution(torch.cat([means, zeros], dim=1), deterministic=True) - - def decode( - self, - z: torch.Tensor, - timestep: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - return self.decoder(z, timestep=timestep, generator=generator) - - def enable_tiling(self) -> None: - # LTX-2 VAE does not implement FastVideo tiling; no-op to satisfy pipeline calls. - return None diff --git a/fastvideo/tests/conftest.py b/fastvideo/tests/conftest.py index 250848e0b..35657b7c8 100644 --- a/fastvideo/tests/conftest.py +++ b/fastvideo/tests/conftest.py @@ -1,12 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -import os import pytest import torch import numpy as np -os.environ["FASTVIDEO_LIGHT_IMPORT"] = "1" - from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, cleanup_dist_env_and_memory) diff --git a/fastvideo/tests/transformers/test_ltx2.py b/fastvideo/tests/transformers/test_ltx2.py deleted file mode 100644 index fd543d036..000000000 --- a/fastvideo/tests/transformers/test_ltx2.py +++ /dev/null @@ -1,458 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import os -import re -from pathlib import Path -import sys - -import pytest -import torch -from safetensors.torch import load_file -from torch.testing import assert_close - -os.environ.setdefault("FASTVIDEO_LIGHT_IMPORT", "1") -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29513") - -repo_root = Path(__file__).resolve().parents[3] -ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" -if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: - sys.path.insert(0, str(ltx_core_path)) - -from fastvideo.configs.models.dits import LTX2VideoConfig -from fastvideo.configs.pipelines import PipelineConfig -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.models.loader.component_loader import TransformerLoader - - -def _load_safetensors(path: Path) -> dict[str, torch.Tensor]: - if path.is_file(): - print(f"[LTX2 TEST] Loading weights from file: {path}") - return load_file(str(path)) - if not path.is_dir(): - raise FileNotFoundError(f"LTX-2 weights not found at {path}") - - model_file = path / "model.safetensors" - if model_file.exists(): - print(f"[LTX2 TEST] Loading weights from file: {model_file}") - return load_file(str(model_file)) - - index_files = list(path.glob("*.safetensors.index.json")) - if index_files: - print(f"[LTX2 TEST] Loading weights from index: {index_files[0]}") - index = index_files[0].read_text(encoding="utf-8") - weight_map = __import__("json").loads(index)["weight_map"] - shards = sorted({path / shard for shard in weight_map.values()}) - else: - shards = sorted(Path(p) for p in path.glob("*.safetensors")) - print(f"[LTX2 TEST] Loading {len(shards)} shard(s) from {path}") - if not shards: - raise FileNotFoundError(f"No safetensors found in {path}") - weights: dict[str, torch.Tensor] = {} - for shard in shards: - print(f"[LTX2 TEST] Loading shard: {shard}") - weights.update(load_file(str(shard))) - print(f"[LTX2 TEST] Loaded {len(weights)} total tensors") - return weights - - -def _normalize_keys(weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - if any(k.startswith("model.diffusion_model.") for k in weights): - normalized = { - k.replace("model.diffusion_model.", ""): v - for k, v in weights.items() - if k.startswith("model.diffusion_model.") - } - print(f"[LTX2 TEST] Normalized model.diffusion_model.* -> {len(normalized)} tensors") - return normalized - if any(k.startswith("diffusion_model.") for k in weights): - normalized = { - k.replace("diffusion_model.", ""): v - for k, v in weights.items() - if k.startswith("diffusion_model.") - } - print(f"[LTX2 TEST] Normalized diffusion_model.* -> {len(normalized)} tensors") - return normalized - if any(k.startswith("model.") for k in weights): - normalized = {k.replace("model.", ""): v for k, v in weights.items()} - print(f"[LTX2 TEST] Normalized model.* -> {len(normalized)} tensors") - return normalized - print(f"[LTX2 TEST] No normalization applied; {len(weights)} tensors") - return weights - - -def _select_transformer_weights(weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - allowed_prefixes = ( - "patchify_proj.", - "adaln_single.", - "caption_projection.", - "audio_patchify_proj.", - "audio_adaln_single.", - "audio_caption_projection.", - "transformer_blocks.", - "scale_shift_table", - "audio_scale_shift_table", - "av_ca_", - "norm_out.", - "audio_norm_out.", - "proj_out.", - "audio_proj_out.", - "video_args_preprocessor.", - ) - filtered = { - k: v - for k, v in weights.items() - if k.startswith(allowed_prefixes) - and not k.startswith("audio_embeddings_connector.") - and not k.startswith("video_embeddings_connector.") - } - print(f"[LTX2 TEST] Selected transformer tensors: {len(filtered)}") - return filtered - - -def _infer_arch(weights: dict[str, torch.Tensor]) -> dict[str, int]: - patch_key = None - if "patchify_proj.weight" in weights: - patch_key = "patchify_proj.weight" - else: - patch_key = next( - (k for k in weights if k.endswith(".patchify_proj.weight") and "audio_" not in k), - None, - ) - if patch_key is None: - raise KeyError("patchify_proj.weight not found in LTX-2 weights.") - in_channels = weights[patch_key].shape[1] - inner_dim = weights[patch_key].shape[0] - - caption_key = next((k for k in weights if k.endswith("caption_projection.linear_1.weight")), None) - caption_channels = weights[caption_key].shape[1] if caption_key else inner_dim - - block_indices = [ - int(m.group(1)) - for k in weights - if (m := re.match(r"transformer_blocks\.(\d+)\.", k)) - ] - num_layers = max(block_indices) + 1 if block_indices else 0 - - num_heads_candidates = [32, 16, 64, 8] - num_attention_heads = next((h for h in num_heads_candidates if inner_dim % h == 0), 32) - attention_head_dim = inner_dim // num_attention_heads - - patch_size = 2 - num_channels_latents = 16 - for candidate in (16, 32, 64, 128): - if in_channels % candidate != 0: - continue - patch_volume = in_channels // candidate - root = int(round(patch_volume**0.5)) - if root * root == patch_volume: - patch_size = root - num_channels_latents = candidate - break - - arch = { - "in_channels": in_channels, - "out_channels": in_channels, - "inner_dim": inner_dim, - "caption_channels": caption_channels, - "num_layers": num_layers, - "num_attention_heads": num_attention_heads, - "attention_head_dim": attention_head_dim, - "num_channels_latents": num_channels_latents, - "patch_size": patch_size, - } - print(f"[LTX2 TEST] Inferred arch: {arch}") - return arch - - -def _load_into_model(model: torch.nn.Module, weights: dict[str, torch.Tensor]) -> int: - model_state = model.state_dict() - filtered = { - k: v - for k, v in weights.items() - if k in model_state and model_state[k].shape == v.shape - } - print( - f"[LTX2 TEST] Loading {len(filtered)} / {len(model_state)} tensors " - f"from {len(weights)} available weights" - ) - if not filtered: - return 0 - model.load_state_dict(filtered, strict=False) - return len(filtered) - - -def _attach_block_sum_logging( - model: torch.nn.Module, - log_path: Path, - label: str, - enabled: bool, -) -> None: - if not enabled: - return - - log_path.parent.mkdir(parents=True, exist_ok=True) - if log_path.exists(): - log_path.unlink() - - def _format_sum(tensor: torch.Tensor | None) -> str: - if tensor is None: - return "None" - return f"{tensor.float().sum().item():.6f}" - - def _hook(module, inputs, outputs): # noqa: ANN001 - if isinstance(outputs, tuple): - video_args, audio_args = outputs - video_sum = _format_sum(video_args.x if video_args is not None else None) - audio_sum = _format_sum(audio_args.x if audio_args is not None else None) - else: - video_sum = _format_sum(outputs) - audio_sum = "None" - with log_path.open("a", encoding="utf-8") as f: - f.write(f"{label}:{module.idx}:video_sum={video_sum},audio_sum={audio_sum}\n") - - for block in model.transformer_blocks: - block.register_forward_hook(_hook) - - -def _attach_block_detail_logging( - model: torch.nn.Module, - log_path: Path, - label: str, - enabled: bool, -) -> None: - if not enabled: - return - - log_path.parent.mkdir(parents=True, exist_ok=True) - if log_path.exists(): - log_path.unlink() - - def _format_sum(tensor: torch.Tensor | None) -> str: - if tensor is None: - return "None" - return f"{tensor.float().sum().item():.6f}" - - def _hook_factory(block_idx: int, name: str): - def _hook(_module, _inputs, outputs): # noqa: ANN001 - if isinstance(outputs, tuple): - out = outputs[0] - else: - out = outputs - out_sum = _format_sum(out if torch.is_tensor(out) else None) - with log_path.open("a", encoding="utf-8") as f: - f.write(f"{label}:{block_idx}:{name}:out_sum={out_sum}\n") - return _hook - - for block in model.transformer_blocks: - idx = block.idx - for name in ( - "attn1", - "attn2", - "ff", - "audio_attn1", - "audio_attn2", - "audio_ff", - "audio_to_video_attn", - "video_to_audio_attn", - ): - if hasattr(block, name): - getattr(block, name).register_forward_hook(_hook_factory(idx, name)) - - -def test_ltx2_transformer_parity(): - torch.manual_seed(42) - official_path = Path( - os.getenv( - "LTX2_OFFICIAL_PATH", - "official_ltx_weights/ltx-2-19b-distilled.safetensors", - ) - ) - fastvideo_path = Path( - os.getenv( - "LTX2_FASTVIDEO_PATH", - "converted/ltx2/transformer", - ) - ) - if not official_path.exists(): - pytest.skip(f"LTX-2 official weights not found at {official_path}") - if not fastvideo_path.exists(): - pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}") - - try: - from ltx_core.components.patchifiers import VideoLatentPatchifier - from ltx_core.guidance.perturbations import BatchedPerturbationConfig - from ltx_core.model.transformer import LTXModel - from ltx_core.model.transformer.model import LTXModelType - from ltx_core.model.transformer.modality import Modality - from ltx_core.model.transformer.attention import AttentionFunction - from ltx_core.model.transformer.rope import LTXRopeType - from ltx_core.types import VideoLatentShape - except ImportError as exc: - pytest.skip(f"LTX-2 import failed: {exc}") - - ref_raw_weights = _load_safetensors(official_path) - ref_weights = _normalize_keys(ref_raw_weights) - ref_weights = _select_transformer_weights(ref_weights) - if not ref_weights: - pytest.skip("No transformer weights found in safetensors file.") - arch = _infer_arch(ref_weights) - if arch["num_layers"] == 0: - pytest.skip("Could not infer transformer block count from LTX-2 weights.") - - config = LTX2VideoConfig() - cfg = config.arch_config - cfg.in_channels = arch["in_channels"] - cfg.out_channels = arch["out_channels"] - cfg.num_layers = arch["num_layers"] - cfg.num_attention_heads = arch["num_attention_heads"] - cfg.attention_head_dim = arch["attention_head_dim"] - cfg.cross_attention_dim = arch["inner_dim"] - cfg.caption_channels = arch["caption_channels"] - cfg.num_channels_latents = arch["num_channels_latents"] - cfg.patch_size = (1, arch["patch_size"], arch["patch_size"]) - - if not torch.cuda.is_available(): - pytest.skip("LTX-2 transformer parity test requires CUDA for attention backends.") - - device = torch.device("cuda:0") - precision = torch.bfloat16 - precision_str = "bf16" - - args = FastVideoArgs( - model_path=str(fastvideo_path), - dit_cpu_offload=True, - use_fsdp_inference=False, - pipeline_config=PipelineConfig(dit_config=config, dit_precision=precision_str), - ) - args.device = device - - loader = TransformerLoader() - fastvideo_model = loader.load(str(fastvideo_path), args).to(device=device, dtype=precision) - - reference_attention = AttentionFunction.PYTORCH - reference_attn_env = os.getenv("LTX2_REFERENCE_ATTN") - if reference_attn_env: - reference_attention = AttentionFunction(reference_attn_env) - else: - fastvideo_attn_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND") - if fastvideo_attn_backend == "FLASH_ATTN": - reference_attention = AttentionFunction.FLASH_ATTENTION_3 - - reference_model = LTXModel( - model_type=LTXModelType.AudioVideo, - num_attention_heads=cfg.num_attention_heads, - attention_head_dim=cfg.attention_head_dim, - in_channels=cfg.in_channels, - out_channels=cfg.out_channels, - num_layers=cfg.num_layers, - cross_attention_dim=cfg.cross_attention_dim, - norm_eps=cfg.norm_eps, - caption_channels=cfg.caption_channels, - positional_embedding_theta=cfg.positional_embedding_theta, - positional_embedding_max_pos=cfg.positional_embedding_max_pos, - timestep_scale_multiplier=cfg.timestep_scale_multiplier, - use_middle_indices_grid=cfg.use_middle_indices_grid, - attention_type=reference_attention, - rope_type=LTXRopeType(cfg.rope_type), - double_precision_rope=cfg.double_precision_rope, - audio_num_attention_heads=cfg.audio_num_attention_heads, - audio_attention_head_dim=cfg.audio_attention_head_dim, - audio_in_channels=cfg.audio_in_channels, - audio_out_channels=cfg.audio_out_channels, - audio_cross_attention_dim=cfg.audio_cross_attention_dim, - audio_positional_embedding_max_pos=cfg.audio_positional_embedding_max_pos, - av_ca_timestep_scale_multiplier=cfg.av_ca_timestep_scale_multiplier, - ).to(device=device, dtype=precision) - - loaded_reference = _load_into_model(reference_model, ref_weights) - if loaded_reference == 0: - pytest.skip("No matching keys loaded into LTX-2 transformer models.") - - fastvideo_model.eval() - reference_model.eval() - - debug_logs = os.getenv("LTX2_DEBUG_LOGS", "0") == "1" - _attach_block_sum_logging( - fastvideo_model.model, - repo_root / "ltx2_debug" / "fastvideo.log", - "fastvideo", - debug_logs, - ) - _attach_block_sum_logging( - reference_model, - repo_root / "ltx2_debug" / "reference.log", - "reference", - debug_logs, - ) - _attach_block_detail_logging( - fastvideo_model.model, - repo_root / "ltx2_debug" / "fastvideo_detail.log", - "fastvideo", - os.getenv("LTX2_DEBUG_DETAIL", "0") == "1", - ) - _attach_block_detail_logging( - reference_model, - repo_root / "ltx2_debug" / "reference_detail.log", - "reference", - os.getenv("LTX2_DEBUG_DETAIL", "0") == "1", - ) - - patchifier = VideoLatentPatchifier(patch_size=cfg.patch_size[1]) - batch_size = 1 - frames = 4 - height = cfg.patch_size[1] * 4 - width = cfg.patch_size[2] * 4 - hidden_states = torch.randn( - batch_size, - cfg.num_channels_latents, - frames, - height, - width, - device=device, - dtype=precision, - ) - encoder_hidden_states = torch.randn( - batch_size, - 16, - cfg.caption_channels, - device=device, - dtype=precision, - ) - timestep = torch.tensor([500], device=device, dtype=precision) - - video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) - positions = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) - latents = patchifier.patchify(hidden_states) - - video = Modality( - enabled=True, - latent=latents, - timesteps=timestep, - positions=positions, - context=encoder_hidden_states, - context_mask=None, - ) - - with torch.no_grad(): - ref_out, _ = reference_model( - video=video, - audio=None, - perturbations=BatchedPerturbationConfig.empty(batch_size), - ) - ref_out = patchifier.unpatchify(ref_out, output_shape=video_shape) - print(f"[LTX2 TEST] Reference model output shape: {ref_out.shape}") - with set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=None, - ): - fastvideo_out = fastvideo_model( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - ) - print(f"[LTX2 TEST] FastVideo model output shape: {fastvideo_out.shape}") - assert ref_out.shape == fastvideo_out.shape - assert ref_out.dtype == fastvideo_out.dtype - assert_close(ref_out, fastvideo_out, atol=1e-2, rtol=1e-2) diff --git a/fastvideo/tests/vaes/test_ltx2_vae.py b/fastvideo/tests/vaes/test_ltx2_vae.py deleted file mode 100644 index ebd9e8cf9..000000000 --- a/fastvideo/tests/vaes/test_ltx2_vae.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import json -import os -from pathlib import Path -import sys - -import pytest -import torch -from safetensors import safe_open -from safetensors.torch import load_file -from torch.testing import assert_close - -os.environ.setdefault("FASTVIDEO_LIGHT_IMPORT", "1") - -repo_root = Path(__file__).resolve().parents[3] -ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" -if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: - sys.path.insert(0, str(ltx_core_path)) - -from fastvideo.models.vaes.ltx2vae import LTX2VideoDecoder, LTX2VideoEncoder - - -def _load_metadata(path: Path) -> dict: - with safe_open(str(path), framework="pt") as f: - meta = f.metadata() - if not meta or "config" not in meta: - raise KeyError("Missing config metadata in safetensors file.") - return json.loads(meta["config"]) - - -def _load_weights(path: Path) -> dict[str, torch.Tensor]: - print(f"[LTX2 VAE TEST] Loading weights from {path}") - return load_file(str(path)) - - -def _select_vae_weights(weights: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]: - filtered = { - k.replace(prefix, ""): v - for k, v in weights.items() - if k.startswith(prefix) or k.startswith("vae.per_channel_statistics.") - } - print(f"[LTX2 VAE TEST] Selected {len(filtered)} tensors for {prefix}") - return filtered - - -def _load_into_model(model: torch.nn.Module, weights: dict[str, torch.Tensor]) -> tuple[int, list[str]]: - model_state = model.state_dict() - filtered = { - k: v - for k, v in weights.items() - if k in model_state and model_state[k].shape == v.shape - } - missing = [k for k in model_state.keys() if k not in filtered] - print( - f"[LTX2 VAE TEST] Loading {len(filtered)} / {len(model_state)} tensors " - f"from {len(weights)} available" - ) - if not filtered: - return 0, missing - model.load_state_dict(filtered, strict=False) - return len(filtered), missing - - -def test_ltx2_vae_parity(): - model_path = Path( - os.getenv( - "LTX2_MODEL_PATH", - "official_ltx_weights/ltx-2-19b-distilled.safetensors", - ) - ) - if not model_path.exists(): - pytest.skip(f"LTX-2 weights not found at {model_path}") - - config = _load_metadata(model_path) - if "vae" not in config: - pytest.skip("VAE config not found in safetensors metadata.") - - try: - from ltx_core.model.video_vae import VideoDecoderConfigurator, VideoEncoderConfigurator - except ImportError as exc: - pytest.skip(f"LTX-2 import failed: {exc}") - - weights = _load_weights(model_path) - encoder_weights = _select_vae_weights(weights, "vae.encoder.") - decoder_weights = _select_vae_weights(weights, "vae.decoder.") - if not encoder_weights or not decoder_weights: - pytest.skip("VAE weights not found in safetensors file.") - - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - precision = torch.bfloat16 if torch.cuda.is_available() else torch.float32 - - fastvideo_encoder = LTX2VideoEncoder(config).to(device=device, dtype=precision) - fastvideo_decoder = LTX2VideoDecoder(config).to(device=device, dtype=precision) - ref_encoder = VideoEncoderConfigurator.from_config(config).to(device=device, dtype=precision) - ref_decoder = VideoDecoderConfigurator.from_config(config).to(device=device, dtype=precision) - - loaded_fastvideo_encoder, missing_fastvideo_encoder = _load_into_model( - fastvideo_encoder.model, encoder_weights - ) - loaded_ref_encoder, missing_ref_encoder = _load_into_model(ref_encoder, encoder_weights) - loaded_fastvideo_decoder, missing_fastvideo_decoder = _load_into_model( - fastvideo_decoder.model, decoder_weights - ) - loaded_ref_decoder, missing_ref_decoder = _load_into_model(ref_decoder, decoder_weights) - - if min( - loaded_fastvideo_encoder, - loaded_ref_encoder, - loaded_fastvideo_decoder, - loaded_ref_decoder, - ) == 0: - pytest.skip("Failed to load VAE weights into one or more models.") - if ( - missing_fastvideo_encoder - or missing_ref_encoder - or missing_fastvideo_decoder - or missing_ref_decoder - ): - print(f"[LTX2 VAE TEST] Missing encoder keys: {len(missing_fastvideo_encoder)}") - print(f"[LTX2 VAE TEST] Missing decoder keys: {len(missing_fastvideo_decoder)}") - pytest.skip("Missing VAE keys; cannot ensure parity.") - - fastvideo_encoder.model.eval() - fastvideo_decoder.model.eval() - ref_encoder.eval() - ref_decoder.eval() - - fastvideo_decoder.model.decode_noise_scale = 0.0 - ref_decoder.decode_noise_scale = 0.0 - - batch_size = 1 - frames = 9 - height = 64 - width = 64 - video = torch.randn( - batch_size, - 3, - frames, - height, - width, - device=device, - dtype=precision, - ) - - with torch.no_grad(): - ref_latents = ref_encoder(video) - fast_latents = fastvideo_encoder(video) - - assert ref_latents.shape == fast_latents.shape - assert ref_latents.dtype == fast_latents.dtype - assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents." - assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents." - assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2) - - timestep = torch.tensor([0.05], device=device, dtype=precision) - with torch.no_grad(): - ref_decoded = ref_decoder(ref_latents, timestep=timestep) - fast_decoded = fastvideo_decoder(fast_latents, timestep=timestep) - - assert ref_decoded.shape == fast_decoded.shape - assert ref_decoded.dtype == fast_decoded.dtype - assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output." - assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output." - assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2) diff --git a/scripts/checkpoint_conversion/convert_ltx2_weights.py b/scripts/checkpoint_conversion/convert_ltx2_weights.py deleted file mode 100644 index cd33100a1..000000000 --- a/scripts/checkpoint_conversion/convert_ltx2_weights.py +++ /dev/null @@ -1,264 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Convert LTX-2 weights to FastVideo naming conventions and split by component. -""" - -from __future__ import annotations - -import argparse -import glob -import json -import os -import re -from collections import OrderedDict -from pathlib import Path - -import torch -from safetensors import safe_open -from safetensors.torch import load_file, save_file - -try: - from huggingface_hub import snapshot_download -except ImportError: # pragma: no cover - optional dependency - snapshot_download = None - - -PARAM_NAME_MAP: dict[str, str] = { - r"^model\.diffusion_model\.(.*)$": r"\1", -} - -COMPONENT_PREFIXES: dict[str, tuple[str, ...]] = { - "transformer": ("model.diffusion_model.",), - "vae": ("vae.",), - "audio_vae": ("audio_vae.",), - "vocoder": ("vocoder.",), - "text_embedding_projection": ("text_embedding_projection.", "model.text_embedding_projection."), -} - - -def _find_shards(model_path: Path) -> list[Path]: - if model_path.is_file(): - return [model_path] - - index_files = list(model_path.glob("*.safetensors.index.json")) - if index_files: - with index_files[0].open("r", encoding="utf-8") as f: - index = json.load(f) - return sorted({model_path / shard for shard in index["weight_map"].values()}) - return sorted(Path(p) for p in glob.glob(str(model_path / "*.safetensors"))) - - -def _apply_mapping(key: str) -> str: - for pattern, replacement in PARAM_NAME_MAP.items(): - if re.match(pattern, key): - return re.sub(pattern, replacement, key) - return key - - -def _load_weights(shards: list[Path]) -> dict[str, torch.Tensor]: - weights: dict[str, torch.Tensor] = {} - for shard in shards: - weights.update(load_file(str(shard))) - return weights - - -def _read_metadata_config(path: Path) -> dict: - with safe_open(str(path), framework="pt") as f: - metadata = f.metadata() - if not metadata or "config" not in metadata: - return {} - return json.loads(metadata["config"]) - - -def _filter_transformer_config(config: dict) -> dict: - transformer = config.get("transformer", {}) - allowed = { - "num_attention_heads", - "attention_head_dim", - "num_layers", - "cross_attention_dim", - "caption_channels", - "norm_eps", - "attention_type", - "positional_embedding_theta", - "positional_embedding_max_pos", - "timestep_scale_multiplier", - "use_middle_indices_grid", - "rope_type", - "frequencies_precision", - "in_channels", - "out_channels", - "audio_num_attention_heads", - "audio_attention_head_dim", - "audio_in_channels", - "audio_out_channels", - "audio_cross_attention_dim", - "audio_positional_embedding_max_pos", - "av_ca_timestep_scale_multiplier", - } - filtered = {k: v for k, v in transformer.items() if k in allowed} - if "frequencies_precision" in filtered: - filtered["double_precision_rope"] = filtered["frequencies_precision"] == "float64" - del filtered["frequencies_precision"] - return filtered - - -def _split_component_weights(weights: dict[str, torch.Tensor]) -> dict[str, OrderedDict]: - components: dict[str, OrderedDict] = {name: OrderedDict() for name in COMPONENT_PREFIXES} - for key, value in weights.items(): - if key.startswith("model.diffusion_model.audio_embeddings_connector."): - new_key = key.replace("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.") - components["text_embedding_projection"][new_key] = value - continue - if key.startswith("model.diffusion_model.video_embeddings_connector."): - new_key = key.replace("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.") - components["text_embedding_projection"][new_key] = value - continue - - matched = False - for component, prefixes in COMPONENT_PREFIXES.items(): - for prefix in prefixes: - if key.startswith(prefix): - new_key = key[len(prefix):] - components[component][new_key] = value - matched = True - break - if matched: - break - return {name: weights for name, weights in components.items() if weights} - - -def _write_component(output_dir: Path, name: str, weights: OrderedDict, config: dict | None) -> None: - component_dir = output_dir / name - component_dir.mkdir(parents=True, exist_ok=True) - output_file = component_dir / "model.safetensors" - save_file(weights, str(output_file)) - print(f"Saved {name} weights to {output_file}") - - if config is not None: - config_path = component_dir / "config.json" - with config_path.open("w", encoding="utf-8") as f: - json.dump(config, f, indent=2) - f.write("\n") - print(f"Saved {name} config to {config_path}") - - -def convert_components( - source_path: Path, - output_dir: Path, - metadata_config: dict, - transformer_class_name: str, - components_to_write: set[str] | None = None, -) -> None: - shards = _find_shards(source_path) - if not shards: - raise FileNotFoundError(f"No safetensors found in {source_path}") - - weights = _load_weights(shards) - split_weights = _split_component_weights(weights) - if components_to_write is not None: - split_weights = {name: weights for name, weights in split_weights.items() if name in components_to_write} - - transformer_weights = split_weights.get("transformer", OrderedDict()) - converted_transformer = OrderedDict() - for key, value in transformer_weights.items(): - new_key = _apply_mapping(f"model.diffusion_model.{key}") - converted_transformer[new_key] = value - split_weights["transformer"] = converted_transformer - - transformer_config = _filter_transformer_config(metadata_config) - if transformer_config: - transformer_config["_class_name"] = transformer_class_name - - component_configs: dict[str, dict | None] = { - "transformer": transformer_config or None, - "vae": metadata_config.get("vae"), - "audio_vae": metadata_config.get("audio_vae"), - "vocoder": metadata_config.get("vocoder"), - "text_embedding_projection": metadata_config.get("text_embedding_projection"), - } - - output_dir.mkdir(parents=True, exist_ok=True) - for name, component_weights in split_weights.items(): - _write_component(output_dir, name, component_weights, component_configs.get(name)) - - -def update_transformer_config(config_path: Path, class_name: str) -> None: - if not config_path.exists(): - print(f"Config file not found: {config_path}") - return - - with config_path.open("r", encoding="utf-8") as f: - config = json.load(f) - - config["_class_name"] = class_name - with config_path.open("w", encoding="utf-8") as f: - json.dump(config, f, indent=2) - f.write("\n") - print(f"Updated _class_name in {config_path} -> {class_name}") - - -def maybe_download(repo_id: str, target_dir: Path, token: str | None, allow_patterns: str | None) -> Path: - if snapshot_download is None: - raise RuntimeError("huggingface_hub is required for --download") - target_dir.mkdir(parents=True, exist_ok=True) - snapshot_download( - repo_id=repo_id, - local_dir=str(target_dir), - local_dir_use_symlinks=False, - token=token, - allow_patterns=allow_patterns, - ) - return target_dir - - -def main() -> None: - parser = argparse.ArgumentParser(description="Convert LTX-2 weights to FastVideo format") - parser.add_argument("--source", type=str, help="Path to transformer weights directory") - parser.add_argument("--output", type=str, required=True, help="Output directory for converted weights") - parser.add_argument("--download", type=str, help="HF repo id to download before conversion") - parser.add_argument("--allow-patterns", type=str, help="Limit HF download to matching files") - parser.add_argument("--token", type=str, default=os.getenv("HF_TOKEN"), help="HF token (or set HF_TOKEN)") - parser.add_argument("--update-config", action="store_true", help="Update source config.json _class_name") - parser.add_argument("--class-name", type=str, default="LTX2Transformer3DModel") - parser.add_argument( - "--transformer-only", - action="store_true", - help="Only convert transformer weights (no component split).", - ) - - args = parser.parse_args() - - if args.download: - if args.source: - raise ValueError("Use either --download or --source, not both.") - source_dir = maybe_download(args.download, Path(args.output) / "download", args.token, args.allow_patterns) - else: - if not args.source: - raise ValueError("--source is required when not using --download") - source_dir = Path(args.source) - - output_dir = Path(args.output) - shards = _find_shards(source_dir) - if not shards: - raise FileNotFoundError(f"No safetensors found in {source_dir}") - metadata_path = shards[0] - metadata_config = _read_metadata_config(metadata_path) - if args.transformer_only: - convert_components( - source_dir, - output_dir, - metadata_config, - args.class_name, - components_to_write={"transformer"}, - ) - else: - convert_components(source_dir, output_dir, metadata_config, args.class_name) - - if args.update_config: - if source_dir.is_dir(): - update_transformer_config(source_dir / "config.json", args.class_name) - - -if __name__ == "__main__": - main()