110 lines
5.5 KiB
Python
110 lines
5.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Regression tests for ``LTX2VideoArchConfig.param_names_mapping``.
|
|
|
|
The ``to_gate_compress`` -> ``to_gate_logits`` rename is the LTX-2.3 gated
|
|
attention loader rule. It must only fire when ``apply_gated_attention=True``;
|
|
otherwise it would silently retarget:
|
|
|
|
- LTX-2.0 ``VIDEO_SPARSE_ATTN`` checkpoints, whose attention modules
|
|
legitimately carry a ``to_gate_compress`` VSA-QAT gate (a sibling of
|
|
``attn_masked`` in ``fastvideo/models/dits/ltx2.py``).
|
|
- LoRAs trained with the default ``lora_target_modules`` list (which
|
|
includes ``to_gate_compress``; see ``fastvideo/train/utils/lora.py:36``
|
|
and ``fastvideo/pipelines/lora_pipeline.py:171``).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from fastvideo.configs.models.dits.ltx2 import LTX2VideoArchConfig
|
|
from fastvideo.models.loader.utils import get_param_names_mapping
|
|
|
|
|
|
def _map(name: str, *, apply_gated_attention: bool) -> str:
|
|
"""Run ``name`` through a fresh config's mapping function."""
|
|
cfg = LTX2VideoArchConfig(apply_gated_attention=apply_gated_attention)
|
|
mapper = get_param_names_mapping(cfg.param_names_mapping)
|
|
target, _, _ = mapper(name)
|
|
return target
|
|
|
|
|
|
class TestLTX20ParamMappingDefault:
|
|
"""``apply_gated_attention=False`` (LTX-2.0 default): ``to_gate_compress``
|
|
must pass through unmodified except for prefix normalization."""
|
|
|
|
@pytest.mark.parametrize("source", [
|
|
"transformer_blocks.0.attn1.to_gate_compress.weight",
|
|
"model.transformer_blocks.5.attn1.to_gate_compress.bias",
|
|
"diffusion_model.transformer_blocks.10.attn1.to_gate_compress.weight",
|
|
"model.diffusion_model.transformer_blocks.2.attn1.to_gate_compress.weight",
|
|
])
|
|
def test_vsa_gate_weight_not_renamed(self, source):
|
|
target = _map(source, apply_gated_attention=False)
|
|
assert "to_gate_logits" not in target, (f"LTX-2.0 VSA gate {source!r} was rewritten to {target!r}; "
|
|
"expected pass-through with only prefix normalization.")
|
|
assert ".to_gate_compress." in target, (f"LTX-2.0 VSA gate {source!r} lost its identity in {target!r}.")
|
|
|
|
@pytest.mark.parametrize("source", [
|
|
"transformer_blocks.0.attn1.to_gate_compress.lora_A.weight",
|
|
"transformer_blocks.0.attn1.to_gate_compress.lora_B.weight",
|
|
"model.transformer_blocks.3.attn1.to_gate_compress.lora_A.weight",
|
|
])
|
|
def test_default_lora_target_not_renamed(self, source):
|
|
"""``to_gate_compress`` is in ``DEFAULT_LORA_TARGET_MODULES``, so LoRAs
|
|
trained with default targets ship these keys."""
|
|
target = _map(source, apply_gated_attention=False)
|
|
assert "to_gate_logits" not in target, (f"Default-target LoRA key {source!r} was rewritten to {target!r}; "
|
|
"this would break every LTX-2.0 LoRA in the wild.")
|
|
|
|
def test_unrelated_param_still_prefix_stripped(self):
|
|
"""Generic prefix-strip behavior is unchanged in LTX-2.0 mode."""
|
|
assert (_map("diffusion_model.transformer_blocks.0.attn1.to_q.weight",
|
|
apply_gated_attention=False) == "model.transformer_blocks.0.attn1.to_q.weight")
|
|
|
|
|
|
class TestLTX23ParamMappingGated:
|
|
"""``apply_gated_attention=True`` (LTX-2.3): the gated-attention
|
|
``to_gate_compress`` upstream key is renamed to ``to_gate_logits``."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"source,expected",
|
|
[
|
|
# No upstream prefix
|
|
("transformer_blocks.0.attn1.to_gate_compress.weight",
|
|
"model.transformer_blocks.0.attn1.to_gate_logits.weight"),
|
|
# ``model.`` prefix
|
|
("model.transformer_blocks.0.attn1.to_gate_compress.weight",
|
|
"model.transformer_blocks.0.attn1.to_gate_logits.weight"),
|
|
# ``diffusion_model.`` prefix
|
|
("diffusion_model.transformer_blocks.0.attn1.to_gate_compress.weight",
|
|
"model.transformer_blocks.0.attn1.to_gate_logits.weight"),
|
|
# ``model.diffusion_model.`` prefix
|
|
("model.diffusion_model.transformer_blocks.0.attn1.to_gate_compress.weight",
|
|
"model.transformer_blocks.0.attn1.to_gate_logits.weight"),
|
|
])
|
|
def test_gate_rename_fires(self, source, expected):
|
|
assert _map(source, apply_gated_attention=True) == expected
|
|
|
|
def test_non_gate_param_still_prefix_stripped(self):
|
|
"""Non-gate params still get the prefix-strip behavior."""
|
|
assert (_map("diffusion_model.transformer_blocks.0.attn1.to_q.weight",
|
|
apply_gated_attention=True) == "model.transformer_blocks.0.attn1.to_q.weight")
|
|
|
|
|
|
class TestParamMappingRuleOrdering:
|
|
"""The gate rules must be inserted *before* the generic prefix-strip rules
|
|
so first-match-wins matching fires the rename first."""
|
|
|
|
def test_gate_rule_wins_over_prefix_strip_when_enabled(self):
|
|
cfg = LTX2VideoArchConfig(apply_gated_attention=True)
|
|
# The first key in the dict iteration order should be a gate rule
|
|
# (not a generic prefix-strip rule), so that ``transformer_blocks
|
|
# .0.attn1.to_gate_compress.weight`` is renamed rather than just
|
|
# prefix-normalized to ``model.transformer_blocks.0.attn1
|
|
# .to_gate_compress.weight``.
|
|
first_pattern = next(iter(cfg.param_names_mapping))
|
|
assert "to_gate_compress" in first_pattern, ("gate rename rule must be inserted at the front of "
|
|
"param_names_mapping; got first pattern: "
|
|
f"{first_pattern!r}")
|