Will Lin
+4
-14
@@ -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__"]
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
|
||||
@@ -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"
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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 = {
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user