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

[Feature] Add official LTX multimodal guidance and precision guard (#5148)

Browse Source 
Signed-off-by: mglyn <1203789601@qq.com>
This commit is contained in:
Mu GuanLin authored and GitHub committed 2026-07-28 10:44:25 +00:00
1 parent 66a9b0c84f
commit 12d75cf63a
51 files changed
+4547 -3062
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
+1 -2
View File
@@ -8,7 +8,6 @@ import torch
from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
@@ -25,7 +24,7 @@ _DMD2_BASE = {
WanT2VDMD2Pipeline: Wan22Pipeline,
WanI2VDMD2Pipeline: Wan22I2VPipeline,
LTX2T2VDMD2Pipeline: LTX2Pipeline,
LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
LTX2I2VDMD2Pipeline: LTX2Pipeline,
FluxDMD2Pipeline: FluxPipeline,
QwenImageDMD2Pipeline: QwenImagePipeline,
}
tests/diffusion/models/dmd2/test_dmd2_scheduler.py
+1 -2
View File
@@ -8,7 +8,6 @@ import torch
from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
@@ -27,7 +26,7 @@ _DMD2_BASE = {
WanT2VDMD2Pipeline: Wan22Pipeline,
WanI2VDMD2Pipeline: Wan22I2VPipeline,
LTX2T2VDMD2Pipeline: LTX2Pipeline,
LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
LTX2I2VDMD2Pipeline: LTX2Pipeline,
FluxDMD2Pipeline: FluxPipeline,
QwenImageDMD2Pipeline: QwenImagePipeline,
}
tests/diffusion/models/ltx2/test_ltx2_conditioning.py
+214 -30
View File
@@ -7,6 +7,7 @@ from types import SimpleNamespace
import pytest
import torch
from PIL import Image
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -16,17 +17,31 @@ def _make_ltx23_request_pipe(cls):
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 99
pipe.vae_spatial_compression_ratio = 32
pipe.vae_temporal_compression_ratio = 8
return pipe
class TestLTXImageToVideoForwardStages:
def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
def test_i2v_pil_preprocessing_matches_official_center_crop(self):
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
image = Image.fromarray(pixels.numpy())
actual = _preprocess_i2v_pil_images(image, height=4, width=4)
expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
torch.testing.assert_close(actual, expected)
def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
pipe = _make_ltx23_request_pipe(LTX2Pipeline)
image = torch.zeros(3, 8, 8)
req = DiffusionRequestBatch(
[
@@ -48,27 +63,45 @@ class TestLTXImageToVideoForwardStages:
)
seen = {}
def fake_forward_impl(req_arg, request_inputs, **kwargs):
def fake_run_recipe(req_arg, request_inputs, **kwargs):
seen["req"] = req_arg
seen["request_inputs"] = request_inputs
seen["kwargs"] = kwargs
return ["i2v-delegated"]
object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
object.__setattr__(pipe, "_run_recipe", fake_run_recipe)
output = pipe.forward(req, noise_scale=0.5)
output = pipe.forward(req)
assert output == ["i2v-delegated"]
assert seen["req"] is req
assert seen["request_inputs"].prompt == ["make the image move"]
assert seen["request_inputs"].negative_prompt == ["jitter"]
assert seen["kwargs"]["image"] is image
assert seen["kwargs"]["noise_scale"] == 0.5
assert seen["kwargs"]["request_sigmas"] is None
def test_unified_entry_selects_t2v_without_images_and_rejects_mixed_batches(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
request_inputs = SimpleNamespace(latents=None)
t2v_req = SimpleNamespace(prompts=["text only", {"prompt": "also text only"}])
assert pipe._resolve_request_image(t2v_req, None, request_inputs) is None
mixed_req = SimpleNamespace(
prompts=[
"text only",
{"prompt": "image conditioned", "multi_modal_data": {"image": torch.zeros(3, 8, 8)}},
]
)
with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
pipe._resolve_request_image(mixed_req, None, request_inputs)
def test_denoise_timestep_kwargs_masks_video_only(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
ts = torch.tensor([2.0, 4.0])
denoise_ctx = SimpleNamespace(
conditioning_mask=torch.tensor([[1.0, 0.0]]),
@@ -79,30 +112,32 @@ class TestLTXImageToVideoForwardStages:
ts,
SimpleNamespace(cfg_parallel_ready=False),
denoise_ctx,
video_token_count=2,
audio_token_count=1,
)
torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
torch.testing.assert_close(kwargs["audio_timestep"], ts)
torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None])
torch.testing.assert_close(kwargs["sigma"], ts)
class TestLTXImageToVideoConditioning:
def test_ltx23_i2v_supports_image_input(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
assert LTX23ImageToVideoPipeline.support_image_input is True
assert LTX2Pipeline.support_image_input is True
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
image = object()
assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
assert LTX2Pipeline._resolve_single_prompt_image([image]) is image
with pytest.raises(ValueError, match="exactly one image per prompt"):
LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
LTX2Pipeline._resolve_single_prompt_image([object(), object()])
def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
image = torch.zeros(1, 3, 4, 4)
additional = {
@@ -111,19 +146,24 @@ class TestLTXImageToVideoConditioning:
"image": torch.ones_like(image),
}
assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
assert LTX2Pipeline._resolve_additional_image(additional) is image
def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
pipe.transformer_temporal_patch_size = 1
object.__setattr__(
pipe,
"_encode_i2v_image_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
)
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
raise AssertionError("packed I2V latents should not be noised")
@@ -134,7 +174,7 @@ class TestLTXImageToVideoConditioning:
latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
out, conditioning_mask = pipe.prepare_latents(
image=None,
image=torch.zeros(1, 3, 1, 1),
batch_size=1,
num_channels_latents=2,
height=1,
@@ -147,14 +187,14 @@ class TestLTXImageToVideoConditioning:
)
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, latents)
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
@@ -165,8 +205,16 @@ class TestLTXImageToVideoConditioning:
latents_std=torch.ones(2),
config=SimpleNamespace(scaling_factor=1.0),
)
object.__setattr__(
pipe,
"_encode_i2v_image_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
)
sampled_shapes = []
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
sampled_shapes.append(tuple(shape))
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
@@ -175,7 +223,7 @@ class TestLTXImageToVideoConditioning:
latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
out, conditioning_mask = pipe.prepare_latents(
image=None,
image=torch.zeros(1, 3, 1, 1),
batch_size=1,
num_channels_latents=2,
height=1,
@@ -187,20 +235,64 @@ class TestLTXImageToVideoConditioning:
latents=latents,
)
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
assert sampled_shapes == [(1, 3, 2)]
def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
pipe.transformer_temporal_patch_size = 1
pipe.vae = SimpleNamespace(
encode=lambda image: image,
latents_mean=torch.zeros(2),
latents_std=torch.ones(2),
)
monkeypatch.setattr(
ltx2_conditioning,
"retrieve_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[10.0]]], [[[11.0]]]]]),
)
sampled_shapes = []
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
sampled_shapes.append(tuple(shape))
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
out, conditioning_mask = pipe.prepare_latents(
image=torch.zeros(1, 3, 1, 1),
batch_size=1,
num_channels_latents=2,
height=1,
width=1,
num_frames=3,
noise_scale=1.0,
dtype=torch.float32,
device=torch.device("cpu"),
)
assert sampled_shapes == [(1, 3, 2)]
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.transformer_spatial_patch_size = 1
pipe.transformer_temporal_patch_size = 1
class FakeScheduler:
def step(self, noise_pred, t, latents, return_dict=False):
return (latents + noise_pred + t,)
sigmas = torch.tensor([1.0, 0.5])
pipe.scheduler = FakeScheduler()
latents = torch.tensor([[[1.0], [2.0], [3.0]]])
@@ -209,11 +301,103 @@ class TestLTXImageToVideoConditioning:
out = pipe._step_video_latents_i2v(
noise_pred,
latents,
torch.tensor(0.5),
0,
latent_num_frames=3,
latent_height=1,
latent_width=1,
)
torch.testing.assert_close(out[:, :1], latents[:, :1])
torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
torch.testing.assert_close(out[:, 1:], latents[:, 1:] - 0.5 * noise_pred[:, 1:])
def test_shared_step_adapter_uses_official_euler_for_t2v(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
pipeline = SimpleNamespace(scheduler=SimpleNamespace(sigmas=torch.tensor([1.0, 0.5])))
adapter = LTXVideoAudioStepAdapter(
pipeline,
SimpleNamespace(sigmas=torch.tensor([1.0, 0.25])),
latent_num_frames=1,
latent_height=1,
latent_width=1,
image_conditioned=False,
)
video = torch.tensor([[[2.0]]])
audio = torch.tensor([[[4.0]]])
((video, audio),) = adapter.step(
(torch.tensor([[[2.0]]]), torch.tensor([[[4.0]]])),
(torch.tensor(1.0), torch.tensor(1.0)),
(video, audio),
)
torch.testing.assert_close(video, torch.tensor([[[1.0]]]))
torch.testing.assert_close(audio, torch.tensor([[[1.0]]]))
def test_shared_step_adapter_uses_official_euler_for_i2v_audio(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
class FakePipeline:
scheduler = object()
def __init__(self):
self.step_indices = []
def _step_video_latents_i2v(
self,
noise_pred,
latents,
step_index,
latent_num_frames,
latent_height,
latent_width,
):
del noise_pred, latent_num_frames, latent_height, latent_width
self.step_indices.append(step_index)
return latents
pipeline = FakePipeline()
adapter = LTXVideoAudioStepAdapter(
pipeline,
SimpleNamespace(sigmas=torch.tensor([1.0, 0.5, 0.0])),
latent_num_frames=1,
latent_height=1,
latent_width=1,
image_conditioned=True,
)
video = torch.tensor([[[1.0]]])
audio = torch.tensor([[[2.0]]])
video_velocity = torch.zeros_like(video)
audio_velocity = torch.tensor([[[4.0]]])
((video, audio),) = adapter.step(
(video_velocity, audio_velocity),
(torch.tensor(1.0), torch.tensor(1.0)),
(video, audio),
)
torch.testing.assert_close(audio, torch.tensor([[[0.0]]]))
((video, audio),) = adapter.step(
(video_velocity, audio_velocity),
(torch.tensor(0.5), torch.tensor(0.5)),
(video, audio),
)
torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
assert pipeline.step_indices == [0, 1]
def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
denoise_ctx = LTXDenoiseContext(
latents=torch.empty(1, 3, 1),
audio_latents=torch.empty(1, 1, 1),
video_coords=torch.empty(1),
audio_coords=torch.empty(1),
conditioning_mask=torch.tensor([[1.0, 0.0, 0.0]]),
)
actual = pipe._video_guidance_model_sigma(torch.tensor(0.75), denoise_ctx)
torch.testing.assert_close(actual, torch.tensor([[[0.0], [0.75], [0.75]]]))
tests/diffusion/models/ltx2/test_ltx2_latents.py
+51 -7
View File
@@ -4,7 +4,7 @@ import pytest
import torch
from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -23,8 +23,51 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
def test_prepare_video_latents_samples_directly_in_packed_token_space():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
pipeline.transformer_spatial_patch_size = 2
pipeline.transformer_temporal_patch_size = 1
expected_generator = torch.Generator().manual_seed(42)
expected = torch.randn((1, 32, 16), generator=expected_generator)
actual = pipeline.prepare_latents(
batch_size=1,
num_channels_latents=4,
height=64,
width=64,
num_frames=9,
dtype=torch.float32,
device=torch.device("cpu"),
generator=torch.Generator().manual_seed(42),
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
pipeline = _make_pipeline(LTX2Pipeline)
expected_generator = torch.Generator().manual_seed(42)
expected = torch.randn((1, 3, 32), generator=expected_generator)
actual, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=64,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=torch.Generator().manual_seed(42),
)
assert original_num_frames == 3
assert padded_num_frames == 3
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
@@ -66,8 +109,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack():
padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
unpadded = latent_ops.unpad_audio_latents(padded, 10)
unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
unpacked = latent_ops.unpack_audio_latents(unpadded, num_mel_bins=2)
expected = latent_ops.unpack_audio_latents(original, num_mel_bins=2)
assert unpacked.shape == (1, 2, 10, 2)
assert not (unpacked == 999.0).any()
@@ -75,7 +118,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack():
def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
pipeline.preserve_sp_padded_audio_duration = True
latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
@@ -97,7 +141,7 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
latents = torch.zeros(1, 2, 13, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
@@ -106,7 +150,7 @@ def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
def test_prepare_audio_latents_rejects_incompatible_provided_length():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
latents = torch.zeros(1, 11, 4)
with pytest.raises(ValueError, match="incompatible audio frame count"):
tests/diffusion/models/ltx2/test_ltx2_parallel.py
+54 -89
View File
@@ -15,9 +15,9 @@ class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
def test_combine_cfg_noise_matches_x0_space_formula(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23Pipeline)
pipe = object.__new__(LTX2Pipeline)
video_sample = torch.tensor([[[1.0, -2.0]]])
audio_sample = torch.tensor([[[0.5, 3.0]]])
video_pos = torch.tensor([[[0.2, -0.3]]])
@@ -50,93 +50,53 @@ class TestCFGParallelHelpers:
assert torch.allclose(video_combined, expected_video)
assert torch.allclose(audio_combined, expected_audio)
def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
pipe = object.__new__(ltx23.LTX23Pipeline)
video_sample = torch.tensor([[[1.0, -2.0]]])
audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
video_pos = torch.tensor([[[0.2, -0.3]]])
video_neg = torch.tensor([[[-0.4, 0.1]]])
audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
video_sigma = torch.tensor(0.25)
audio_sigma = torch.tensor(0.5)
scale = 4.0
class FakeCfgGroup:
def all_gather(self, tensor, separate_tensors=True):
assert separate_tensors
if tensor.shape == video_pos.shape:
return [video_pos, video_neg]
return [audio_pos, audio_neg]
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
expected_video = ltx2_guidance.combine_velocity_via_x0(
video_sample,
video_pos,
video_neg,
video_sigma,
scale,
)
expected_audio = ltx2_guidance.combine_velocity_via_x0(
audio_sample,
audio_pos,
audio_neg,
audio_sigma,
scale,
def test_cfg_parallel_rejects_non_cfg_guidance(self):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTX_GUIDANCE_EXECUTOR,
LTXGuidancePlan,
)
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
for rank, expected_branch in ((0, "positive"), (1, "negative")):
calls = []
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
def fake_predict_noise(**kwargs):
calls.append(kwargs["branch"])
if kwargs["branch"] == "positive":
return video_pos, audio_pos
return video_neg, audio_neg
with pytest.raises(ValueError, match="CFG-only guidance"):
LTX_GUIDANCE_EXECUTOR.validate_cfg_world_size(plan, 2)
object.__setattr__(pipe, "predict_noise", fake_predict_noise)
video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
true_cfg_scale=scale,
positive_kwargs={"branch": "positive"},
negative_kwargs={"branch": "negative"},
cfg_normalize=False,
video_latents=video_sample,
audio_latents=audio_sample,
video_sigma=video_sigma,
audio_sigma=audio_sigma,
)
def test_cfg_parallel_dummy_warms_supported_cfg_only_plan(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
assert calls == [expected_branch]
torch.testing.assert_close(video_combined, expected_video)
torch.testing.assert_close(audio_combined, expected_audio)
monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
pipe = object.__new__(LTX2Pipeline)
req = SimpleNamespace(is_dummy_run=lambda: True)
request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
assert "_cfg_video_latents" not in pipe.__dict__
assert "_cfg_audio_latents" not in pipe.__dict__
cfg_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
assert cfg_parallel_ready is True
assert pipe._guidance_plan.names == ("cond", "uncond")
assert pipe._guidance_plan.spec.video.cfg_scale == 3.0
assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0
class TestCFGParallelForwardPath:
"""Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
@pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"])
def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name):
from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime
def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipeline_cls = getattr(ltx23, pipeline_cls_name)
pipe = object.__new__(pipeline_cls)
pipe = object.__new__(ltx23.LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 4
monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
pipe.vae_spatial_compression_ratio = 32
pipe.vae_temporal_compression_ratio = 8
monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
class StopAtEncodePromptError(Exception):
pass
@@ -224,7 +184,7 @@ class TestCFGParallelForwardPath:
("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
[(1.0, 1, 1.0), (None, 24, 24.0)],
)
def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
def test_forward_cfg_parallel_uses_shared_video_audio_euler_step(
self,
monkeypatch,
cfg_rank,
@@ -233,13 +193,13 @@ class TestCFGParallelForwardPath:
audio_sampling_rate,
expected_frame_rate,
):
from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime
from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_runtime
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipe = object.__new__(ltx23.LTX23Pipeline)
pipe = object.__new__(ltx23.LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 1
@@ -278,7 +238,7 @@ class TestCFGParallelForwardPath:
return [audio_pos, audio_neg]
raise AssertionError(f"Unexpected gathered tensor: {tensor}")
monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
@@ -314,7 +274,7 @@ class TestCFGParallelForwardPath:
def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
assert padding_side == "left"
assert prompt_embeds.shape[0] == 2
assert prompt_embeds.shape[0] == 1
return prompt_embeds, prompt_embeds, prompt_attention_mask
rope_video_fps: list[float] = []
@@ -380,7 +340,16 @@ class TestCFGParallelForwardPath:
num_frames=1,
frame_rate=frame_rate_input,
num_inference_steps=2,
guidance_scale=4.0,
extra_args={
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 0.0,
"audio_stg_scale": 0.0,
"video_modality_scale": 1.0,
"audio_modality_scale": 1.0,
"video_rescale_scale": 0.0,
"audio_rescale_scale": 0.0,
},
latents=video_latents,
audio_latents=audio_latents,
output_type="latent",
@@ -388,35 +357,31 @@ class TestCFGParallelForwardPath:
request_id="ltx23-cfg-parallel-forward-test",
)
output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
output = pipe.forward(DiffusionRequestBatch(requests=[req]), sigmas=[0.25, 0.25])[0]
expected_video_noise = ltx2_guidance.combine_velocity_via_x0(
video_latents,
video_pos,
video_neg,
pipe.scheduler.sigmas[0],
4.0,
3.0,
)
expected_audio_noise = ltx2_guidance.combine_velocity_via_x0(
audio_latents,
audio_pos,
audio_neg,
pipe.scheduler.sigmas[0],
4.0,
7.0,
)
scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
assert scheduler_call_names == ["video", "audio", "video", "audio"]
assert pipe.scheduler.calls == []
assert len(pipe.transformer.calls) == 2
torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
video_out, audio_out = output.output
torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
sigma_delta = pipe.scheduler.sigmas[-1] - pipe.scheduler.sigmas[0]
expected_video = video_latents + expected_video_noise * sigma_delta
expected_audio = audio_latents + expected_audio_noise * sigma_delta
torch.testing.assert_close(video_out, expected_video.reshape(1, 2, 1, 1, 1))
torch.testing.assert_close(audio_out, expected_audio.reshape(1, 1, 1, 2))
# fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
# to the model's own 24.0 default, not crash on None; a provided rate is passed through.
tests/diffusion/models/ltx2/test_ltx2_pipeline.py
+1001 -203
View File
File diff suppressed because it is too large. Load diff
tests/diffusion/models/ltx2/test_ltx2_transformer.py
+16 -2
View File
@@ -10,11 +10,12 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
def test_ltx_rms_norm_no_affine_matches_official_torch_module():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
assert "weight" not in dict(norm.named_parameters())
assert "weight" in dict(norm.named_buffers())
assert norm.weight is None
assert "weight" not in dict(norm.named_buffers())
assert "weight" not in norm.state_dict()
@@ -50,3 +51,16 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
matched.append(name)
assert matched == ["transformer_blocks.0", "transformer_blocks.1"]
@pytest.mark.parametrize("rope_type", ["interleaved", "split"])
def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
video_timestep = root_plan["timestep"]
audio_timestep = root_plan["audio_timestep"]
assert video_timestep.split_dim == audio_timestep.split_dim == 1
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
tests/diffusion/models/ltx2/test_ltx2_vae.py
+24 -16
View File
@@ -14,7 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
[(False, 0), (True, 1), (RuntimeError("unavailable"), 0)],
[(False, 0), (True, 1), (RuntimeError("unavailable"), None)],
)
def test_non_output_rank_only_enters_collective_vae_decode(
self,
@@ -22,7 +22,8 @@ class TestLTXOutputRank:
distributed_vae_state,
expected_decode_calls,
):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
class FakeVae:
dtype = torch.float32
@@ -40,23 +41,30 @@ class TestLTXOutputRank:
self.decode_calls += 1
return (torch.ones(1, 1),)
monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(ltx.LTX2Pipeline)
monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae = FakeVae()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="np",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
decode_kwargs = {
"latents": torch.ones(1, 1),
"audio_latents": torch.ones(1, 1),
"output_type": "np",
"connector_prompt_embeds": torch.ones(1, 1),
"generator": None,
"device": torch.device("cpu"),
"decode_timestep": 0.0,
"decode_noise_scale": None,
"prompt_batch_size": 1,
}
if isinstance(distributed_vae_state, Exception):
with pytest.raises(RuntimeError, match="unavailable"):
pipe._decode_output(**decode_kwargs)
assert pipe.vae.decode_calls == 0
return
output = pipe._decode_output(**decode_kwargs)
assert pipe.vae.decode_calls == expected_decode_calls
assert output.output[0].numel() == 0