325 lines
12 KiB
Python
325 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
from torch.testing import assert_close
|
|
|
|
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 _read_transformer_config(config: dict) -> dict:
|
|
transformer_config = config.get("transformer", {})
|
|
if not transformer_config:
|
|
raise ValueError("Missing transformer config in LTX-2 metadata.")
|
|
return transformer_config
|
|
|
|
|
|
def _infer_patch_params(in_channels: int) -> tuple[int, int]:
|
|
patch_size = 1
|
|
num_channels_latents = 128
|
|
for candidate in (8, 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
|
|
print(f"[LTX2 TEST] Inferred patch_size={patch_size}, "
|
|
f"num_channels_latents={num_channels_latents}")
|
|
return patch_size, num_channels_latents
|
|
|
|
|
|
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 _output_hook(name: str):
|
|
|
|
def _hook(_module, _inputs, outputs): # noqa: ANN001
|
|
out = outputs[0] if isinstance(outputs, tuple) else 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}:output:{name}:out_sum={out_sum}\n")
|
|
|
|
return _hook
|
|
|
|
for name in ("proj_out", "audio_proj_out"):
|
|
if hasattr(model, name):
|
|
getattr(model, name).register_forward_hook(_output_hook(name))
|
|
|
|
|
|
def test_ltx2_transformer_parity():
|
|
torch.manual_seed(42)
|
|
diffusers_root = Path(os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers"))
|
|
official_path = Path(os.getenv(
|
|
"LTX2_OFFICIAL_PATH",
|
|
"official_ltx_weights/ltx-2-19b-distilled.safetensors",
|
|
))
|
|
fastvideo_path = Path(os.getenv(
|
|
"LTX2_FASTVIDEO_PATH",
|
|
str(diffusers_root / "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.loader.sft_loader import SafetensorsModelStateDictLoader
|
|
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
|
|
from ltx_core.model.transformer import (LTXModelConfigurator, LTXV_MODEL_COMFY_RENAMING_MAP)
|
|
from ltx_core.model.transformer.modality import Modality
|
|
from ltx_core.types import VideoLatentShape
|
|
except ImportError as exc:
|
|
pytest.skip(f"LTX-2 import failed: {exc}")
|
|
|
|
config_loader = SafetensorsModelStateDictLoader()
|
|
metadata = config_loader.metadata(str(official_path))
|
|
transformer_config = _read_transformer_config(metadata)
|
|
|
|
config = LTX2VideoConfig()
|
|
cfg = config.arch_config
|
|
cfg.num_attention_heads = transformer_config.get("num_attention_heads", cfg.num_attention_heads)
|
|
cfg.attention_head_dim = transformer_config.get("attention_head_dim", cfg.attention_head_dim)
|
|
cfg.num_layers = transformer_config.get("num_layers", cfg.num_layers)
|
|
cfg.cross_attention_dim = transformer_config.get("cross_attention_dim", cfg.cross_attention_dim)
|
|
cfg.caption_channels = transformer_config.get("caption_channels", cfg.caption_channels)
|
|
cfg.norm_eps = transformer_config.get("norm_eps", cfg.norm_eps)
|
|
cfg.attention_type = transformer_config.get("attention_type", cfg.attention_type)
|
|
cfg.positional_embedding_theta = transformer_config.get("positional_embedding_theta",
|
|
cfg.positional_embedding_theta)
|
|
cfg.positional_embedding_max_pos = transformer_config.get("positional_embedding_max_pos",
|
|
cfg.positional_embedding_max_pos)
|
|
cfg.timestep_scale_multiplier = transformer_config.get("timestep_scale_multiplier", cfg.timestep_scale_multiplier)
|
|
cfg.use_middle_indices_grid = transformer_config.get("use_middle_indices_grid", cfg.use_middle_indices_grid)
|
|
cfg.rope_type = transformer_config.get("rope_type", cfg.rope_type)
|
|
cfg.double_precision_rope = transformer_config.get(
|
|
"double_precision_rope",
|
|
transformer_config.get("frequencies_precision", "") == "float64",
|
|
)
|
|
cfg.audio_num_attention_heads = transformer_config.get("audio_num_attention_heads", cfg.audio_num_attention_heads)
|
|
cfg.audio_attention_head_dim = transformer_config.get("audio_attention_head_dim", cfg.audio_attention_head_dim)
|
|
cfg.audio_in_channels = transformer_config.get("audio_in_channels", cfg.audio_in_channels)
|
|
cfg.audio_out_channels = transformer_config.get("audio_out_channels", cfg.audio_out_channels)
|
|
cfg.audio_cross_attention_dim = transformer_config.get("audio_cross_attention_dim", cfg.audio_cross_attention_dim)
|
|
cfg.audio_positional_embedding_max_pos = transformer_config.get(
|
|
"audio_positional_embedding_max_pos",
|
|
cfg.audio_positional_embedding_max_pos,
|
|
)
|
|
cfg.av_ca_timestep_scale_multiplier = transformer_config.get("av_ca_timestep_scale_multiplier",
|
|
cfg.av_ca_timestep_scale_multiplier)
|
|
cfg.in_channels = transformer_config.get("in_channels", cfg.in_channels)
|
|
cfg.out_channels = transformer_config.get("out_channels", cfg.out_channels)
|
|
|
|
patch_size, num_channels_latents = _infer_patch_params(cfg.in_channels)
|
|
cfg.patch_size = (1, patch_size, patch_size)
|
|
cfg.num_channels_latents = num_channels_latents
|
|
|
|
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_builder = SingleGPUModelBuilder(
|
|
model_class_configurator=LTXModelConfigurator,
|
|
model_path=str(official_path),
|
|
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
|
|
)
|
|
reference_model = reference_builder.build(device=device, dtype=precision).to(device=device, dtype=precision)
|
|
reference_model.set_gradient_checkpointing(False)
|
|
|
|
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-4, rtol=1e-4)
|