anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feat] LTX Standard Two-Stage Pipeline (#5500)

Browse Source 
Signed-off-by: mglyn <1203789601@qq.com>
This commit is contained in:
Mu GuanLin authored and GitHub committed 2026-08-14 00:13:14 +08:00
1 parent 4b66f62e10
commit 83d1f11449
29 files changed
+2644 -371
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/ltx2/test_ltx2_latents.py
+41 -2
View File
@@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
def test_prepare_video_latents_samples_directly_in_packed_token_space():
def test_prepare_video_latents_matches_official_values_and_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
@@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
@@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
generator_sp1 = torch.Generator().manual_seed(42)
generator_sp4 = torch.Generator().manual_seed(42)
latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp1,
)
latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp4,
)
torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
torch.testing.assert_close(
torch.randn(8, generator=generator_sp4),
torch.randn(8, generator=generator_sp1),
rtol=0,
atol=0,
)
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
@@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
packed = latent_ops.pack_audio_latents(latents)
torch.testing.assert_close(padded[:, :10], packed[:, :10])
torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
tests/diffusion/models/ltx2/test_ltx2_parallel.py
+14
View File
@@ -12,6 +12,19 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
sigma = mocker.MagicMock()
sigma.to.return_value.item.return_value = 0.5
actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
sigma.to.assert_called_once_with(torch.float32)
sigma.to.return_value.item.assert_called_once_with()
torch.testing.assert_close(actual, torch.tensor([2.0]))
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
@@ -199,6 +212,7 @@ class TestCFGParallelHelpers:
),
attention_kwargs=None,
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
original_audio_num_frames=state.audio.shape[1],
)
actual_video, actual_audio = executor.predict_parallel_guidance(
tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py
+483
View File
@@ -0,0 +1,483 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the fixed LTX refinement-phase adapter."""
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
AdapterTarget,
parse_ltx_adapter,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_TWO_STAGE_RECIPE,
LTX_DISTILLED_ADAPTER_SLOT,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _target(
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
):
return AdapterTarget(
source_module=source_module,
module=module,
slice=target_slice,
a_key="a",
b_key="b",
a_shape=tuple(lora_a.shape),
b_shape=tuple(lora_b.shape),
)
def _wrap(
layer,
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
layer_fused=False,
):
target = _target(
lora_a,
lora_b,
source_module=source_module,
module=module,
target_slice=target_slice,
)
wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
layer,
[target],
adapter_dtype=lora_a.dtype,
layer_fused=layer_fused,
)
wrapper.load_target(target, lora_a, lora_b)
wrapper.set_enabled(True)
return wrapper
class _LoRAAttention(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.to_qkv = torch.nn.Identity()
class _LoRATransformer(torch.nn.Module):
stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
def __init__(self) -> None:
super().__init__()
self.proj_in = torch.nn.Linear(4, 6, bias=False)
block = torch.nn.Module()
block.attn1 = _LoRAAttention()
self.transformer_blocks = torch.nn.ModuleList([block])
def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
},
adapter_path,
)
manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
("proj_in", "proj_in", None),
("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
}
def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
adapter_path,
)
with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
def _pipeline(*, dtype, quantization_config):
return SimpleNamespace(
pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
transformer=object(),
od_config=SimpleNamespace(
model="unused",
lora_path=None,
dtype=dtype,
quantization_config=quantization_config,
),
)
@pytest.mark.parametrize(
("dtype", "quantization_config", "layer_fused"),
[
(torch.float32, object(), False),
(torch.bfloat16, None, True),
],
)
def test_ltx_phase_adapter_selects_execution_from_quantization(
monkeypatch,
dtype,
quantization_config,
layer_fused,
):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object()
captured = {}
def resolve_artifact(*_args, **kwargs):
assert not kwargs
return "adapter"
def parse(transformer, path):
assert transformer is pipeline.transformer
assert path == "adapter"
return manifest
class Runtime:
def __init__(self, transformer, parsed_manifest, **kwargs):
captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
assert isinstance(runtime, Runtime)
assert captured == {
"transformer": pipeline.transformer,
"manifest": manifest,
"dtype": dtype,
"layer_fused": layer_fused,
}
def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
with pytest.raises(ValueError, match="bfloat16"):
ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
class TinyTransformer(torch.nn.Module):
stacked_params_mapping = ()
def __init__(self) -> None:
super().__init__()
self.proj = torch.nn.Linear(2, 2, bias=False)
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
},
adapter_path,
)
transformer = TinyTransformer()
runtime = LTXPhaseAdapterRuntime(
transformer,
parse_ltx_adapter(transformer, str(adapter_path)),
dtype=torch.float32,
)
assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
with torch.no_grad():
transformer.proj.base_layer.weight.copy_(torch.eye(2))
runtime.finalize()
x = torch.tensor([[1.0, 1.0]])
runtime.activate(None)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
base_weight = torch.tensor(
[
[0.5, -0.25, 0.125, 0.75],
[-0.375, 0.625, 0.25, -0.5],
[0.875, -0.125, -0.75, 0.375],
],
dtype=torch.bfloat16,
)
lora_a = torch.tensor(
[
[0.1015625, -0.203125, 0.3046875, -0.40625],
[0.5078125, -0.609375, 0.7109375, -0.8125],
],
dtype=torch.bfloat16,
)
lora_b = torch.tensor(
[[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
dtype=torch.bfloat16,
)
with torch.no_grad():
layer.weight.copy_(base_weight)
layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeRow(torch.nn.Module):
input_size = 4
output_size = 2
input_size_per_partition = 2
input_is_parallel = True
tp_rank = 0
tp_size = 2
reduce_results = True
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
reductions = []
def all_reduce(value):
reductions.append(value.clone())
return value + 1
monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
layer = FakeRow()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a[:, :2])
expected_weight.add_(layer.weight)
x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
actual = wrapper(x)
assert len(reductions) == 1
assert torch.equal(reductions[0], expected_local)
assert torch.equal(actual, expected_local + 1)
def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
input_size = 3
output_size = 4
output_size_per_partition = 2
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
)
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
layer = FakeColumn()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b[2:], lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
pass
class FakeQKV(FakeColumn):
input_size = 4
head_size = 2
v_head_size = 2
total_num_heads = 2
total_num_kv_heads = 2
num_heads = 1
num_kv_heads = 1
num_kv_head_replicas = 1
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = True
return_bias = True
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
)
self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
layer = FakeQKV()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(
layer,
lora_a,
lora_b,
source_module="to_k",
module="to_qkv",
target_slice="k",
layer_fused=True,
)
expected_weight = base_weight.clone()
expected_k = torch.matmul(lora_b[2:], lora_a)
expected_k.add_(base_weight[2:4])
expected_weight[2:4].copy_(expected_k)
x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
expected_output = torch.nn.functional.linear(x, expected_weight)
actual_output, returned_bias = wrapper(x)
assert torch.equal(actual_output, expected_output)
assert returned_bias is layer.bias
assert torch.equal(layer.weight, base_weight)
@pytest.mark.parametrize(
("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
[
(
lambda: type(
"FakeColumn",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 4,
"output_size_per_partition": 2,
"tp_rank": 1,
"gather_output": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[12963.0, 17284.0]]),
),
(
lambda: type(
"FakeRow",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 2,
"input_size_per_partition": 2,
"tp_rank": 1,
"reduce_results": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[2.0], [3.0]]),
torch.tensor([[8600.0, 12900.0]]),
),
(
lambda: type(
"FakeQKV",
(torch.nn.Module,),
{
"input_size": 4,
"head_size": 2,
"v_head_size": 2,
"total_num_heads": 2,
"total_num_kv_heads": 2,
"num_heads": 1,
"num_kv_heads": 1,
"num_kv_head_replicas": 1,
"tp_rank": 1,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
},
)(),
"k",
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
),
],
)
def test_ltx_phase_adapter_uses_local_tp_lora_slices(
monkeypatch,
layer_factory,
target_slice,
expected_a,
expected_b,
expected_output,
):
layer = layer_factory()
layer.register_buffer("device_anchor", torch.empty(0))
layer_type = type(layer)
if target_slice == "k":
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
elif "Column" in layer_type.__name__:
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
else:
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
tests/diffusion/models/ltx2/test_ltx2_pipeline.py
+362 -10
View File
@@ -13,13 +13,21 @@ import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.ltx2 import ltx2_components
from vllm_omni.diffusion.models.ltx2 import (
ltx2_components,
ltx2_transformer,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
LTX23_DISTILLED_COMPONENT_PROFILE,
LTX23_TWO_STAGE_COMPONENT_PROFILE,
_install_connector_attention,
detect_ltx_model_version,
resolve_ltx_artifact,
resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
@@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
_repeat_batch,
build_perturbation_kwargs,
combine_guided_x0,
)
@@ -40,18 +49,29 @@ from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
LTX2_TWO_STAGE_RECIPE,
LTX23_DISTILLED_TWO_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
LTX23_TWO_STAGE_RECIPE,
LTX_DEFAULT_NEGATIVE_PROMPT,
LTX_POSITIVE_ONLY_RECIPE,
resolve_ltx_pipeline_recipe,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import (
LTXRequestInputs,
validate_ltx_checkpoint,
validate_pipeline_request,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2DistilledPipeline,
LTX2TwoStagePipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
)
assert issubclass(LTX2Pipeline, LTXRuntime)
assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
for pipeline_cls, profile in (
(LTX2Pipeline, LTX2_COMPONENT_PROFILE),
(LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
(LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
):
assert pipeline_cls._dit_modules == list(profile.dit_modules)
assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
assert pipeline_cls._vae_modules == list(profile.vae_modules)
assert pipeline_cls._resident_modules == list(profile.resident_modules)
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
@@ -123,6 +155,72 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
calls = []
def fake_download(**kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
)
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
expected = tmp_path / filename
expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config"),
[
("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
],
)
def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config", "error"),
[
("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
],
)
def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
with pytest.raises(ValueError, match=error):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
def test_ltx_rejects_advanced_uaa_before_component_initialization():
od_config = SimpleNamespace(
model="unused",
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
)
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
LTX2Pipeline(od_config=od_config)
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
@@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
assert pipe.reports_stage_durations
def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
initialized = False
def stub_components(*_args, **_kwargs):
nonlocal initialized
initialized = True
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
cache_backend="cache_dit",
)
)
assert not initialized
def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
pipe = LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
lora_path=None,
model_config={},
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = LTX2DistilledPipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert pipe._phase_adapter is None
assert pipe.preserve_sp_padded_audio_duration
assert pipe.reports_stage_durations
def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
assert LTX2Pipeline.supports_request_batch
assert LTX2Pipeline.support_image_input
assert LTX2Pipeline.unified_text_image_entry
assert LTX2Pipeline.distributed_video_decode
assert not LTX2TwoStagePipeline.supports_request_batch
assert LTX2TwoStagePipeline.support_image_input
assert LTX2TwoStagePipeline.unified_text_image_entry
assert not LTX2DistilledPipeline.supports_request_batch
assert LTX2DistilledPipeline.support_image_input
assert LTX2DistilledPipeline.unified_text_image_entry
def test_ltx_two_stage_resolves_request_embedded_images():
pipe = object.__new__(LTX2TwoStagePipeline)
request_inputs = SimpleNamespace(latents=None)
image = torch.zeros(3, 8, 8)
req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
assert pipe._resolve_request_image(req, None, request_inputs) is image
def test_ltx_dmd2_entries_retain_their_task_boundaries():
assert not LTX2T2VDMD2Pipeline.support_image_input
assert LTX2I2VDMD2Pipeline.support_image_input
@@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
def test_ltx_positive_only_guidance_preserves_token_major_layout():
tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
repeated = _repeat_batch(tensor, 1)
assert repeated is tensor
assert repeated.stride() == tensor.stride()
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
cond = torch.tensor(
[
@@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
torch.testing.assert_close(actual, torch.zeros_like(cond))
def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
expected = combine_guided_x0(
cond=logical_cond,
uncond_text=logical_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
)
padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
actual = combine_guided_x0(
cond=padded_cond,
uncond_text=padded_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
rescale_token_count=2,
)
torch.testing.assert_close(actual[:, :2], expected)
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
@@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
assert "load_weights" not in LTX2TwoStagePipeline.__dict__
assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
@@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
torch.testing.assert_close(state.audio, torch.tensor(11.0))
def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
@@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
activated_slots = []
if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
@@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
if pipeline_cls is LTX2TwoStagePipeline:
assert activated_slots == [None, "ltx_distilled"]
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
@@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
"ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
)
@pytest.mark.parametrize(
("request_inputs", "error"),
[
@@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
@pytest.mark.parametrize(
("recipe", "steps", "stg_block"),
[
(LTX2_TWO_STAGE_RECIPE, 40, 29),
(LTX23_TWO_STAGE_RECIPE, 30, 28),
],
)
def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
stage1, stage2 = recipe.phases
assert (recipe.width, recipe.height) == (1536, 1024)
assert recipe.num_inference_steps == steps
assert stage1.spatial_downscale == 2
assert stage1.adapter_slot is None
assert stage1.guidance.video.stg_blocks == (stg_block,)
assert stage1.guidance.audio.stg_blocks == (stg_block,)
assert stage2.input_transform == "spatial_upsample"
assert stage2.adapter_slot == "ltx_distilled"
assert stage2.guidance == LTXGuidanceSpec.positive_only()
assert stage2.num_inference_steps == 3
assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
# Official Stage 2 refines video only; its audio result is discarded.
assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
phase_switches = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases:
ordinary_pipeline._enter_phase(phase)
assert phase_switches == [None, "ltx_distilled"]
distilled_pipeline = object.__new__(LTX2DistilledPipeline)
for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
distilled_pipeline._enter_phase(phase)
def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
class FakeColumn(torch.nn.Module):
def __init__(
self,
input_size,
output_size,
*,
bias,
gather_output,
return_bias,
quant_config,
prefix,
):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.bias_enabled = bias
self.gather_output = gather_output
self.return_bias = return_bias
self.quant_config = quant_config
self.prefix = prefix
class FakeRow(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
class FakeAttention(torch.nn.Module):
def __init__(self, **kwargs):
super().__init__()
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
attention = ltx2_transformer.LTX2Attention(
query_dim=16,
heads=4,
dim_head=4,
qk_norm="rms_norm",
apply_gated_attention=True,
pack_qkv=False,
prefix="transformer_blocks.0.attn1",
)
gate = attention.to_gate_logits
assert isinstance(gate, FakeColumn)
assert gate.input_size == 16
assert gate.output_size == 4
assert gate.gather_output is False
assert gate.return_bias is False
assert gate.quant_config is None
assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
class TestLTXRequestParsing:
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1432,6 +1773,7 @@ class TestPipelineComponents:
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
class TestLTX23DecodeConditioning:
@@ -1477,6 +1819,11 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2TwoStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1488,6 +1835,9 @@ class TestRegistryIntegration:
"LTX23ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2ImageToVideoTwoStagePipeline",
"LTX23TwoStagePipeline",
"LTX23ImageToVideoTwoStagePipeline",
}
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
@@ -1497,6 +1847,7 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1509,6 +1860,7 @@ class TestRegistryIntegration:
"model_class_name",
[
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
tests/diffusion/models/ltx2/test_ltx2_transformer.py
+90 -3
View File
@@ -1,13 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
LTX2AudioVideoAttnProcessor,
LTX2VideoTransformer3DModel,
_LTX2ParallelAttention,
_make_rms_norm,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
processor = LTX2AudioVideoAttnProcessor()
hidden_states = torch.zeros(1, 3, 4)
encoder_hidden_states = torch.zeros(1, 4, 4)
key_padding_mask = torch.tensor([[True, True, True, False]])
actual = processor._prepare_attention_mask(
SimpleNamespace(),
hidden_states,
encoder_hidden_states,
key_padding_mask,
batch_size=1,
sequence_length=4,
)
torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
def _fake_ltx_parallel_attention():
calls = []
class Backend:
@staticmethod
def get_name():
return "FLASH_ATTN"
@staticmethod
def supports_attention_mask():
return True
def run(name):
def forward(query, _key, _value, _metadata):
calls.append(name)
return torch.zeros_like(query)
return forward
layer = object.__new__(_LTX2ParallelAttention)
nn.Module.__init__(layer)
layer.attn_backend = Backend
layer.attention = SimpleNamespace(forward=run("native"))
layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
layer.backend_pref = "FLASH_ATTN"
return layer, calls
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
layer, calls = _fake_ltx_parallel_attention()
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_masked_self_attention_uses_sdpa():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True, True, True, False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_unmasked_cross_attention_keeps_native_backend():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
layer._run_local_attention(query, key, value, None)
assert calls == ["native"]