From 5e4daaf1970ef230c50d6fb2c7ef8ed4895545f8 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Thu, 27 Aug 2026 23:54:40 +0800 Subject: [PATCH] [Diffusion] Add LTX-2.5 Diffusion VAE decoder support (#6189) Signed-off-by: mglyn <1203789601@qq.com> Co-authored-by: Hongsheng Liu --- pyproject.toml | 1 - recipes/LTX/LTX-2.5.md | 52 ++- requirements/common.txt | 2 + tests/config/test_omni_config.py | 29 ++ .../models/ltx2/test_ltx25_pipeline.py | 170 ++++++- .../models/ltx2/test_ltx2_conditioning.py | 35 +- .../models/ltx2/test_ltx2_latents.py | 36 ++ .../models/ltx2/test_ltx2_phase_adapter.py | 10 +- .../models/ltx2/test_ltx2_pipeline.py | 129 ++++- tests/diffusion/models/ltx2/test_ltx2_vae.py | 442 +++++++++++++++++- .../ltx/test_ltx25_official_similarity.py | 14 +- tests/e2e/online_serving/test_ltx25.py | 6 +- .../test_async_omni_diffusion_config.py | 58 +++ .../diffusion/models/ltx2/ltx2_components.py | 114 ++++- .../models/ltx2/ltx2_conditioning.py | 15 +- .../models/ltx2/ltx2_diffusion_decoder.py | 245 ++++++++++ .../ltx2_diffusion_decoder_distributed.py | 429 +++++++++++++++++ .../diffusion/models/ltx2/ltx2_latents.py | 16 +- .../models/ltx2/ltx2_phase_adapter.py | 4 +- .../diffusion/models/ltx2/ltx2_request.py | 15 +- .../diffusion/models/ltx2/ltx2_runtime.py | 35 +- vllm_omni/engine/async_omni_engine.py | 36 +- 22 files changed, 1832 insertions(+), 61 deletions(-) create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py diff --git a/pyproject.toml b/pyproject.toml index 9b9493730..2d43edfee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,6 @@ dev = [ "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests "torchdiffeq==0.2.5", # for covo_audio Test "FlagEmbedding==1.4.0", # for hunyuan_image3 tests - "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0) "s3tokenizer==0.3.0", # for tts test # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, # whose top-level stepaudio2 package conflicts with this distribution. diff --git a/recipes/LTX/LTX-2.5.md b/recipes/LTX/LTX-2.5.md index 4ff5f6dab..f0be6948b 100644 --- a/recipes/LTX/LTX-2.5.md +++ b/recipes/LTX/LTX-2.5.md @@ -10,14 +10,22 @@ checkpoint for T2V and first-frame I2V. Every output includes synchronized 48 kHz stereo audio. The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5) -repository is not directly loadable with `--model`. It supplies the official -upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept -both model licenses and authenticate before first use. +repository is not directly loadable with `--model`. It supplies the canonical +Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by +the Full/SFT two-stage pipeline, so accept both model licenses and authenticate +before first use. + +The vLLM-Omni integration code is provided under the project's +[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate +artifacts covered by the +[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md), +not the vLLM-Omni license. Users are responsible for reviewing and complying +with the model license before downloading or using the weights. ## Choose a pipeline | Pipeline | Mode | Default size | Schedule | -|---|---|---:|---:| +| --- | --- | ---: | ---: | | `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps | | `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps | | `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps | @@ -28,6 +36,40 @@ latent upsampler, and run a three-step refinement stage. Select the class with `--model-class-name`; no `--task-type` flag is required. Supplying one initial image selects I2V, while omitting it selects T2V. +## Video decoder + +LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in +to the legacy convolutional VAE, set the startup-only stage-0 model extra +`ltx2_use_conv_vae: true`. The choice applies to all requests and all four +pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub +during startup. `decode_timestep` and `decode_noise_scale` condition only the +legacy ConvVAE and have no effect when DiffVAE is selected. + +For offline Python usage: + +```python +omni = Omni( + model="Lightricks/LTX-2.5-Diffusers", + model_class_name="LTX2Pipeline", + stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}', +) +``` + +For online serving: + +```bash +vllm serve Lightricks/LTX-2.5-Diffusers \ + --omni \ + --model-class-name LTX2Pipeline \ + --stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}' +``` + +Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving +serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial +dimension. DiffVAE also supports distributed VAE decode; see the +[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md). +DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder. + ## Prerequisites ```bash @@ -42,7 +84,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and ## Hardware | GPU | Status | Recommended scope | -|---|---|---| +| --- | --- | --- | | NVIDIA B300 | Verified | All four canonical pipelines | | NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines | | NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines | diff --git a/requirements/common.txt b/requirements/common.txt index 97ccac6ee..99d7412d2 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -13,6 +13,8 @@ transformers >= 5.10.1, < 5.15 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.40.0 +# Hub kernels, including LTX-2.5 NATTEN +kernels==0.14.1 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index b72db7802..c216206f4 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -940,6 +940,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic(): omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false") +def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path): + deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " extras:", + " ltx2_use_conv_vae: true", + ] + ) + ) + + stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + + assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True + + +def test_stage_override_routes_ltx2_conv_vae_extra(): + stage = _from_pipeline_key( + "dreamzero", + cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}}, + ).stage_by_id(0) + + assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True + + def test_structured_diffusion_config_rejects_invalid_compile_granularity(): with pytest.raises(ValidationError, match="diffusion_compile_granularity"): omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py index cbe52e997..4b1fef7d5 100644 --- a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Focused LTX-2.5 pipeline correctness tests.""" import json from dataclasses import replace from types import SimpleNamespace +from typing import Any import pytest import torch @@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test( def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch): - pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None)) + pipe = SimpleNamespace( + component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None), + use_diffusion_decoder=True, + ) od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16) monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) @@ -187,7 +191,19 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc def test_ltx_converted_component_loading_propagates_revision(monkeypatch): revision = "pinned-revision" - calls = {"components": []} + calls: dict[str, Any] = {"components": []} + natten_processor = object() + + class FakeDiffusionDecoder: + def set_attn_processor(self, processor): + calls["diffusion_decoder_processor"] = processor + + def enable_tiling(self): + calls["diffusion_decoder_tiling"] = True + + def set_parallel_size(self, parallel_size, mode): + calls["diffusion_decoder_parallel"] = (parallel_size, mode) + profile = replace( LTX25_DISTILLED_COMPONENT_PROFILE, text_encoder_cls=object, @@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch): vocoder_fallback_cls=None, scheduler_use_dynamic_shifting=False, ) - pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage") + pipeline = SimpleNamespace( + component_profile=profile, + pipeline_kind="distilled_two_stage", + use_diffusion_decoder=True, + ) od_config = SimpleNamespace( model="org/converted-ltx25", revision=revision, dtype=torch.bfloat16, quantization_config=None, + extras={}, + vae_use_tiling=False, + parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"), ) monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) @@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch): ) return object() + def fake_native_diffusion_decoder(model, **kwargs): + calls["native_diffusion_decoder"] = (model, kwargs) + return FakeDiffusionDecoder() + def fake_transformer_config(model, subfolder, local_files_only, *, revision): calls["transformer_config"] = (model, subfolder, local_files_only, revision) return {"component": "transformer"} @@ -234,7 +261,13 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch): monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch) monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer) monkeypatch.setattr(ltx2_components, "_load_component", fake_component) + monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder) monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config) + monkeypatch.setattr( + ltx2_components, + "LTX2VideoVaeNeighborhoodNattenProcessor", + lambda: natten_processor, + ) monkeypatch.setattr( ltx2_components, "create_transformer_from_config", @@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch): "latent_upsampler", } assert all(call[3]["revision"] == revision for call in calls["components"]) + assert "diffusion_decoder" not in calls["prefetch"][1] + assert calls["native_diffusion_decoder"] == ( + od_config.model, + {"local_files_only": False, "dtype": torch.bfloat16, "revision": revision}, + ) + assert calls["diffusion_decoder_processor"] is natten_processor + assert calls["diffusion_decoder_tiling"] is True + assert calls["diffusion_decoder_parallel"] == (2, "tile") assert calls["transformer_config"] == ( od_config.model, profile.transformer_subfolder, @@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch): assert calls["scheduler"][1]["revision"] == revision +def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch): + revision = "diffusers-revision" + config = {"decoder_stage_channels": [8, 8, 8, 8, 8]} + calls: dict[str, Any] = {} + + class FakeDecoder: + def init_distributed(self): + calls["init_distributed"] = True + + expected = FakeDecoder() + + def fake_load_config(model, **kwargs): + calls["config"] = (model, kwargs) + return config + + def fake_resolve(model, repo_id, filename, **kwargs): + calls["artifact"] = (model, repo_id, filename, kwargs) + return "/models/native-diffvae.safetensors" + + def fake_load(cls, path, passed_config, dtype): + calls["load"] = (path, passed_config, dtype) + return expected + + monkeypatch.setattr( + ltx2_components.DistributedLTX2VideoDiffusionDecoderModel, + "load_config", + staticmethod(fake_load_config), + ) + monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve) + monkeypatch.setattr( + ltx2_components.DistributedLTX2VideoDiffusionDecoderModel, + "from_ltx25_native_checkpoint", + classmethod(fake_load), + ) + + actual = ltx2_components._load_ltx25_native_diffusion_decoder( + "Lightricks/LTX-2.5-Diffusers", + local_files_only=False, + dtype=torch.bfloat16, + revision=revision, + ) + + assert actual is expected + assert calls["config"] == ( + "Lightricks/LTX-2.5-Diffusers", + { + "subfolder": "diffusion_decoder", + "local_files_only": False, + "revision": revision, + }, + ) + assert calls["artifact"] == ( + "Lightricks/LTX-2.5-Diffusers", + "Lightricks/LTX-2.5", + "vae/ltx-2.5-video-vae-bf16.safetensors", + { + "model_revision": revision, + "artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION, + }, + ) + assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16) + assert calls["init_distributed"] is True + + +def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch): + class MissingNattenProcessor: + def __init__(self): + raise ImportError("kernels is unavailable") + + monkeypatch.setattr( + ltx2_components, + "LTX2VideoVaeNeighborhoodNattenProcessor", + MissingNattenProcessor, + ) + + with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"): + ltx2_components._create_ltx25_natten_processor() + + def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path): (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"})) (tmp_path / "transformer").mkdir() @@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint(): assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage" +@pytest.mark.parametrize( + "pipeline_cls", + [ + LTX2Pipeline, + LTX2TwoStagePipeline, + LTX2DistilledOneStagePipeline, + LTX2DistilledTwoStagePipeline, + ], +) +@pytest.mark.parametrize( + ("extras", "use_diffusion_decoder"), + [({}, True), ({"ltx2_use_conv_vae": True}, False)], +) +def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in( + tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder +): + (tmp_path / "model_index.json").write_text( + json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]}) + ) + + 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(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = pipeline_cls( + od_config=SimpleNamespace( + model=str(tmp_path), + extras=extras, + enable_diffusion_pipeline_profiler=False, + ) + ) + + assert pipe.use_diffusion_decoder is use_diffusion_decoder + assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder) + + @pytest.mark.parametrize( ("pipeline_kind", "expected"), [ @@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules(): max_sequence_length=16, ) prompt_context = object() - phase_calls = [] + phase_calls: list[Any] = [] def resolve_request_inputs(req, **kwargs): return request_inputs diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py index cbc66e69c..aa5654c74 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for LTX image-to-video input and conditioning behavior.""" @@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning: assert LTX2Pipeline.support_image_input is True + def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(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.use_diffusion_decoder = True + pipe.vae = SimpleNamespace( + encode=lambda image: image, + latents_mean=torch.full((2,), 100.0), + latents_std=torch.full((2,), 10.0), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.diffusion_decoder = SimpleNamespace( + latents_mean=torch.tensor([1.0, 3.0]), + latents_std=torch.tensor([2.0, 4.0]), + config=SimpleNamespace(scaling_factor=2.0), + ) + monkeypatch.setattr( + ltx2_conditioning, + "retrieve_latents", + lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]), + ) + + actual = pipe._encode_i2v_image_latents( + torch.zeros(1, 3, 1, 1), + batch_size=1, + generator=None, + dtype=torch.float32, + ) + + torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0)) + def test_ltx23_i2v_rejects_multi_image_prompt_list(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index a9d19945a..fbac0bfe5 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace import pytest @@ -47,6 +50,39 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout(): assert actual.stride()[1:] == (1, actual.shape[1]) +def test_prepare_video_latents_uses_diffusion_decoder_statistics(): + pipeline = _make_pipeline(LTX2Pipeline) + pipeline.use_diffusion_decoder = True + pipeline.vae_spatial_compression_ratio = 1 + pipeline.vae_temporal_compression_ratio = 1 + pipeline.transformer_spatial_patch_size = 1 + pipeline.transformer_temporal_patch_size = 1 + pipeline.vae = SimpleNamespace( + latents_mean=torch.full((2,), 100.0), + latents_std=torch.full((2,), 10.0), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipeline.diffusion_decoder = SimpleNamespace( + latents_mean=torch.tensor([1.0, 3.0]), + latents_std=torch.tensor([2.0, 4.0]), + config=SimpleNamespace(scaling_factor=2.0), + ) + + actual = pipeline.prepare_latents( + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=1, + noise_scale=0.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]), + ) + + torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0)) + + def test_prepare_audio_latents_samples_directly_in_packed_token_space(): pipeline = _make_pipeline(LTX2Pipeline) diff --git a/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py index 4b1d18cb5..f9b3bedce 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py +++ b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the fixed LTX refinement-phase adapter.""" from types import SimpleNamespace +from typing import Any import pytest import torch @@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization( ): pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config) manifest = object() - captured = {} + captured: dict[str, Any] = {} def resolve_artifact(*_args, **kwargs): - assert not kwargs + assert kwargs == { + "model_revision": None, + "artifact_revision": None, + } return "adapter" def parse(transformer, path): diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index b3da177db..9cafa22b6 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the shared LTX pipeline runtime and public contracts.""" @@ -8,6 +8,7 @@ import os import tempfile from dataclasses import replace from types import SimpleNamespace +from typing import Any import pytest import torch @@ -165,7 +166,19 @@ 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): +@pytest.mark.parametrize( + ("model", "repo_id", "expected_revision"), + [ + ("Lightricks/test", "Lightricks/test", "model-revision"), + ("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"), + ], +) +def test_ltx_artifact_uses_source_revision_and_hub_fallback( + monkeypatch, + model, + repo_id, + expected_revision, +): filename = "ltx-sidecar.safetensors" calls = [] @@ -176,9 +189,22 @@ def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch) monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) assert ( - resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors" + resolve_ltx_artifact( + model, + repo_id, + filename, + model_revision="model-revision", + artifact_revision="artifact-revision", + ) + == "/cache/ltx-sidecar.safetensors" ) - assert calls == [{"repo_id": "Lightricks/test", "filename": filename}] + assert calls == [ + { + "repo_id": repo_id, + "filename": filename, + "revision": expected_revision, + } + ] def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): @@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): 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) + assert resolve_ltx_artifact( + str(tmp_path), + "Lightricks/test", + filename, + model_revision="unused-model-revision", + artifact_revision="unused-artifact-revision", + ) == str(expected) + + +def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch): + 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), + "Lightricks/test", + "ltx-sidecar.safetensors", + model_revision="local-model-revision", + artifact_revision="artifact-revision", + ) + == "/cache/ltx-sidecar.safetensors" + ) + assert calls == [ + { + "repo_id": "Lightricks/test", + "filename": "ltx-sidecar.safetensors", + "revision": "artifact-revision", + } + ] @pytest.mark.parametrize( @@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls): max_sequence_length=16, ) prompt_context = object() - phase_calls = [] + phase_calls: list[Any] = [] def resolve_request_inputs(req, **kwargs): return request_inputs @@ -959,7 +1019,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls): pipeline.vae_spatial_compression_ratio = 32 pipeline.vae_temporal_compression_ratio = 8 pipeline.latent_upsampler = FakeUpsampler() - activated_slots = [] + activated_slots: list[Any] = [] if pipeline_cls is LTX2TwoStagePipeline: pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append) object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) @@ -1133,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas(): ) +def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested(): + request = LTXRequestInputs( + prompt="prompt", + negative_prompt="", + height=128, + width=128, + num_frames=9, + frame_rate=24.0, + num_inference_steps=30, + guidance=LTX2_ONE_STAGE_RECIPE.request_guidance, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=16, + ) + + with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"): + validate_pipeline_request( + request, + pipeline_recipe=LTX2_ONE_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2Pipeline", + min_video_latent_spatial_size=(7, 7), + ) + + # ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid. + validate_pipeline_request( + request, + pipeline_recipe=LTX2_ONE_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2Pipeline", + ) + + validate_pipeline_request( + replace(request, height=256, width=256), + pipeline_recipe=LTX2_ONE_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2Pipeline", + min_video_latent_spatial_size=(7, 7), + ) + + @pytest.mark.parametrize( ("direct_kwargs", "sampling_kwargs", "error"), [ @@ -1284,7 +1397,7 @@ def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, step def test_ltx_two_stage_entries_select_the_official_adapter_slots(): - phase_switches = [] + phase_switches: list[Any] = [] ordinary_pipeline = object.__new__(LTX2TwoStagePipeline) ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append) for phase in LTX2_TWO_STAGE_RECIPE.phases: diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 2e0647ccf..964bb7d06 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior.""" +"""Unit tests for LTX video VAE tiling and distributed decode behavior.""" from types import SimpleNamespace +from typing import Any import pytest import torch @@ -36,6 +37,441 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): assert output.dtype == torch.bfloat16 +class TestLTXDiffusionDecoder: + @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"]) + def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None) + + with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"): + model.set_parallel_size(2, mode=mode) + + def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + calls = [] + model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.distributed_executor = SimpleNamespace( + set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode)) + ) + + model.set_parallel_size(4, mode="tile") + + assert calls == [(4, "tile")] + + def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + LTX2VideoDiffusionDecoder3d, + LTX2VideoDiffusionDecoderModel, + ) + + class IdentityAttention: + def build_block_mask(self, _hidden_states): + return None + + class IdentityBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.attn = IdentityAttention() + + def forward(self, hidden_states, _block_mask): + return hidden_states + + class IdentityUpsample(torch.nn.Module): + def forward(self, hidden_states, *, drop_leading_frame): + assert drop_leading_frame + return hidden_states + + decoder = object.__new__(LTX2VideoDiffusionDecoder3d) + torch.nn.Module.__init__(decoder) + decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])]) + decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()]) + decoder.trailing_pad_latent_frames = 2 + decoder.temporal_compression_ratio = 8 + decoder.stage5_kernel = (11, 3, 3) + + # 17 frames contain one real frame followed by 16 propagated ghost + # frames. The reference keeps 11 for stage-5 NATTEN instead of + # cropping the tensor immediately back to one. + context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4)) + assert context.shape == (1, 11, 3, 3, 4) + + class FakeDecoder(torch.nn.Module): + def forward(self, z, *, generator, num_inference_steps): + assert generator is None + assert num_inference_steps is None + return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2) + + model = object.__new__(LTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.decoder = FakeDecoder() + model.use_tiling = False + model.tile_sample_min_num_frames = 80 + model.tile_sample_min_height = 768 + model.tile_sample_min_width = 768 + model.temporal_compression_ratio = 8 + model.spatial_compression_ratio = 2 + + output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0] + assert output.shape == (1, 3, 1, 4, 6) + + @pytest.mark.parametrize( + ("model_version", "extras", "expected"), + [ + ("2.5", {}, True), + ("2.5", {"ltx2_use_conv_vae": True}, False), + ("2.3", {}, False), + ], + ) + def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected): + from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder + + assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected + + def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self): + from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder + + with pytest.raises(TypeError, match="ltx2_use_conv_vae"): + _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5") + + def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + convert_ltx25_native_diffusion_decoder_state_dict, + ) + + qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4) + qkv_bias = torch.arange(6, dtype=torch.float32) + projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2) + projection_bias = torch.ones(2) + gate = torch.tensor([2.0, 3.0]) + mean = torch.tensor([0.1, 0.2]) + std = torch.tensor([0.3, 0.4]) + native = { + "encoder.ignored": torch.ones(1), + "per_channel_statistics.mean-of-means": mean, + "vae.per_channel_statistics.std-of-means": std, + "decoder.type_emb": torch.ones(2), + "decoder.coarse_head.weight": torch.ones(1), + "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight, + "decoder.det_stages.0.0.attn.qkv.bias": qkv_bias, + "decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]), + "decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]), + "decoder.diff_blocks.0.attn.proj.weight": projection_weight, + "decoder.diff_blocks.0.attn.proj.bias": projection_bias, + "decoder.diff_blocks.0.gate_msa": gate, + "vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2), + } + + converted = convert_ltx25_native_diffusion_decoder_state_dict(native) + + torch.testing.assert_close(converted["latents_mean"], mean) + torch.testing.assert_close(converted["latents_std"], std) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0])) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0])) + torch.testing.assert_close( + converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight + ) + torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias) + assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted + assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted) + + def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + convert_ltx25_native_diffusion_decoder_state_dict, + ) + + with pytest.raises(ValueError, match="not divisible by 3"): + convert_ltx25_native_diffusion_decoder_state_dict( + {"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)} + ) + + def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + seen = {} + + class ConvVae: + dtype = torch.float32 + config = SimpleNamespace(timestep_conditioning=True) + + def decode(self, *_args, **_kwargs): + raise AssertionError("the convolutional VAE decoder must not run") + + class DiffusionDecoder: + dtype = torch.float32 + + def decode(self, latents, *, generator, return_dict): + seen["latents"] = latents + seen["generator"] = generator + seen["return_dict"] = return_dict + return (latents + 1,) + + class AudioVae: + dtype = torch.float32 + + def decode(self, audio_latents, *, return_dict): + assert return_dict is False + return (audio_latents + 2,) + + class VideoProcessor: + def postprocess_video(self, video, *, output_type): + seen["output_type"] = output_type + return video + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.vae = ConvVae() + pipe.diffusion_decoder = DiffusionDecoder() + pipe.audio_vae = AudioVae() + pipe.vocoder = torch.nn.Identity() + pipe.video_processor = VideoProcessor() + + generator = torch.Generator().manual_seed(123) + latents = torch.ones(1, 1) + output = pipe._decode_output( + latents=latents, + audio_latents=torch.ones(1, 1), + output_type="pt", + connector_prompt_embeds=torch.ones(1, 1), + generator=generator, + device=torch.device("cpu"), + decode_timestep=0.7, + decode_noise_scale=0.5, + prompt_batch_size=1, + ) + + torch.testing.assert_close(seen["latents"], latents) + assert seen["generator"] is generator + assert seen["return_dict"] is False + assert seen["output_type"] == "pt" + torch.testing.assert_close(output.output[0], latents + 1) + torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0)) + + def test_diffusion_decode_denormalizes_with_decoder_statistics(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.vae = SimpleNamespace( + latents_mean=torch.full((2,), 100.0), + latents_std=torch.full((2,), 10.0), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.diffusion_decoder = SimpleNamespace( + latents_mean=torch.tensor([10.0, 20.0]), + latents_std=torch.tensor([2.0, 4.0]), + config=SimpleNamespace(scaling_factor=2.0), + ) + pipe.audio_vae = SimpleNamespace( + latents_mean=torch.tensor(0.0), + latents_std=torch.tensor(1.0), + ) + forward_ctx = SimpleNamespace( + latent_num_frames=1, + latent_height=1, + latent_width=1, + original_audio_num_frames=1, + latent_mel_bins=2, + ) + + video, audio = pipe._unpack_and_denormalize_stage( + forward_ctx, + torch.tensor([[[1.0, 2.0]]]), + torch.zeros(1, 1, 2), + ) + + torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0])) + torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2)) + + def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + class ConvVae: + config = SimpleNamespace(timestep_conditioning=False) + + def is_distributed_enabled(self): + raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives") + + class DiffusionDecoder: + dtype = torch.float32 + + def decode(self, *_args, **_kwargs): + raise AssertionError("only the output rank may run DiffVAE decode") + + 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.use_diffusion_decoder = True + pipe.vae = ConvVae() + pipe.diffusion_decoder = DiffusionDecoder() + + output = pipe._decode_output( + latents=torch.ones(1, 1), + audio_latents=torch.ones(1, 1), + output_type="pt", + 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, + ) + + assert output.output[0].numel() == 0 + assert output.output[1].numel() == 0 + + def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + class ConvVae: + config = SimpleNamespace(timestep_conditioning=False) + + def is_distributed_enabled(self): + raise AssertionError("DiffVAE decode must query the active decoder") + + class DiffusionDecoder: + dtype = torch.float32 + + def __init__(self): + self.decode_calls = 0 + + def is_distributed_enabled(self): + return True + + def decode(self, *_args, **_kwargs): + self.decode_calls += 1 + return (torch.empty(0, 3, 0, 0, 0),) + + 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.use_diffusion_decoder = True + pipe.vae = ConvVae() + pipe.diffusion_decoder = DiffusionDecoder() + + output = pipe._decode_output( + latents=torch.ones(1, 1), + audio_latents=torch.ones(1, 1), + output_type="pt", + 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, + ) + + assert pipe.diffusion_decoder.decode_calls == 1 + assert output.output[0].numel() == 0 + assert output.output[1].numel() == 0 + + def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + expected = torch.ones(1, 3, 1, 2, 2) + seen = {} + + def native_tiled_decode(self, z, generator=None, num_inference_steps=None): + seen["z"] = z + seen["generator"] = generator + seen["num_inference_steps"] = num_inference_steps + return expected + + monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode) + model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.is_distributed_enabled = lambda: False + generator = torch.Generator().manual_seed(17) + z = torch.zeros(1, 1, 1, 1, 1) + + result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode( + model, + z, + generator=generator, + num_inference_steps=1, + ) + + assert result is expected + assert seen["z"] is z + assert seen["generator"] is generator + assert seen["num_inference_steps"] == 1 + + def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch): + from diffusers.utils.torch_utils import randn_tensor + + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + model = DistributedLTX2VideoDiffusionDecoderModel( + out_channels=1, + latent_channels=1, + patch_size=1, + decoder_head_dim=8, + decoder_stage_channels=(8, 8, 8, 8, 8), + decoder_stage_depths=(1, 1, 1, 1, 1), + decoder_stage_kernels=((1, 1, 1),) * 4, + decoder_upsample_strides=((1, 1, 1),) * 4, + decoder_upsample_channel_reductions=(1, 1, 1, 1), + decoder_stage5_kernel=(1, 1, 1), + spatial_compression_ratio=1, + temporal_compression_ratio=1, + ) + model.tile_sample_min_num_frames = 2 + model.tile_sample_stride_num_frames = 1 + model.tile_sample_min_height = 2 + model.tile_sample_stride_height = 1 + model.tile_sample_min_width = 2 + model.tile_sample_stride_width = 1 + model.distributed_executor = SimpleNamespace(group=None) + model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1) + monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None) + + z = torch.zeros(1, 1, 2, 3, 3) + distributed_generator = torch.Generator().manual_seed(123) + reference_generator = torch.Generator().manual_seed(123) + tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1) + + assert grid_spec.grid_shape == (1, 3, 3) + assert grid_spec.split_dims == (2, 3, 4) + for task in tasks: + tile_shape = model._tiled_pixel_shape_from_features( + task.tensor, + drop_leading_frame=task.drop_leading_frame, + crop_trailing_ghost=task.crop_trailing_ghost, + ) + expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype) + actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype) + torch.testing.assert_close(actual, expected) + + assert torch.equal(distributed_generator.get_state(), reference_generator.get_state()) + + class TestLTXOutputRank: @pytest.mark.parametrize( ("distributed_vae_state", "expected_decode_calls"), @@ -138,7 +574,7 @@ class TestLTX23VaeDistributedDecode: }, output_dtype=torch.float32, ) - seen = {} + seen: dict[str, Any] = {} def exec_tile(task): return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) diff --git a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py index d7c82ef3a..c8ce29c37 100644 --- a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py +++ b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E accuracy guards against pinned official Lightricks LTX references. @@ -43,7 +43,7 @@ LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5" LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" LTX25_OFFICIAL_COMMON_FILES = { "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", - "video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors", + "video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors", "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors", } LTX25_OFFICIAL_VARIANT_FILES = { @@ -208,6 +208,9 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> model = Path(configured_model).resolve() snapshot_revision = model.name if model.parent.name == "snapshots" else None assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}" + assert (model / "diffusion_decoder" / "config.json").is_file(), ( + f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}" + ) return model, configured_model, configured_revision or snapshot_revision model_id = configured_model or LTX25_OMNI_MODEL_ID revision = configured_revision @@ -223,6 +226,7 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> "connectors/config.json", "connectors/diffusion_pytorch_model.safetensors.index.json", "connectors/diffusion_pytorch_model-*.safetensors", + "diffusion_decoder/config.json", "scheduler/*", "text_encoder/config.json", "text_encoder/generation_config.json", @@ -470,7 +474,7 @@ def _assert_strict_similarity( @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("task", ["t2v", "i2v"]) def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None: - """Compare Omni with the official runtime using the same checkpoint connector weights.""" + """Compare Omni with official canonical DiffVAE and matching connector weights.""" _require_ltx25_model_access() artifact_parent = accuracy_artifact_root / "ltx_official" output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}") @@ -545,6 +549,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path omni_metadata = json.loads((omni_output / "metadata.json").read_text()) assert official_metadata["official_revision"] == official_revision assert official_metadata["pipeline"] == "DistilledPipeline" + assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"]) assert official_metadata["attention_backend"] == ATTENTION_BACKEND assert official_metadata["connector_model"] == str(omni_model) assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline" @@ -594,7 +599,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path @pytest.mark.parametrize("task", ["t2v", "i2v"]) @pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"]) def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None: - """Compare raw official Full/SFT weights with Omni's converted Full pipeline.""" + """Compare raw official Full/SFT and DiffVAE weights with Omni.""" _require_ltx25_model_access() artifact_parent = accuracy_artifact_root / "ltx_official" output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}") @@ -697,6 +702,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi "full_two_stage": "TI2VidTwoStagesPipeline", }[pipeline_mode] ) + assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"]) assert official_metadata["attention_backend"] == ATTENTION_BACKEND assert official_metadata["connector_model"] == str(omni_model) assert omni_metadata["model_class_name"] == model_class_name diff --git a/tests/e2e/online_serving/test_ltx25.py b/tests/e2e/online_serving/test_ltx25.py index da5c53851..4975abcc0 100644 --- a/tests/e2e/online_serving/test_ltx25.py +++ b/tests/e2e/online_serving/test_ltx25.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Online serving smokes for all four canonical LTX-2.5 pipelines.""" @@ -73,8 +73,8 @@ def test_ltx25_pipeline_entries( "form_data": { "model": omni_server.model, "prompt": PROMPT, - "height": 128, - "width": 128, + "height": 256, + "width": 256, "num_frames": 9, "fps": 24, "num_inference_steps": num_inference_steps, diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 452eec286..1127e136d 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -33,6 +33,64 @@ def test_default_stage_config_includes_cache_backend(): assert engine_args["model_stage"] == "diffusion" +def test_default_stage_config_preserves_model_extras(): + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0] + + assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True + + +def test_default_stage_config_preserves_and_overrides_promoted_extras(): + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "extras": { + "auxiliary_text_encoder": "/models/extras-llama", + "default_llama_model_id": "extras/default-llama", + }, + "auxiliary_text_encoder": None, + } + )[0] + + extras = stage_cfg["engine_args"]["extras"] + assert extras["auxiliary_text_encoder"] == "/models/extras-llama" + assert extras["default_llama_model_id"] == "extras/default-llama" + + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "extras": { + "auxiliary_text_encoder": "/models/extras-llama", + "default_llama_model_id": "extras/default-llama", + }, + "auxiliary_text_encoder": "/models/top-level-llama", + "default_llama_model_id": "top-level/default-llama", + } + )[0] + + extras = stage_cfg["engine_args"]["extras"] + assert extras["auxiliary_text_encoder"] == "/models/top-level-llama" + assert extras["default_llama_model_id"] == "top-level/default-llama" + + +def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker): + """Local/unregistered Diffusers checkpoints still honor stage-0 extras.""" + + def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs): + return None, default_stage_cfg_factory(), None + + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + side_effect=resolve_with_default, + ) + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + + _, stage_configs = engine._resolve_stage_configs( + "/models/LTX-2.5-Diffusers", + {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'}, + trust_remote_code=False, + ) + + assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True + + def test_default_cache_config_used_when_missing(): """Ensure default cache_config is synthesized when only backend is given.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index d4d98e7c4..6325df870 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared component construction helpers for the LTX model family.""" @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler +from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder @@ -35,6 +36,12 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from .ltx2_diffusion_decoder import ( + LTX25_NATIVE_ARTIFACT_REVISION, + LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, +) +from .ltx2_diffusion_decoder_distributed import DistributedLTX2VideoDiffusionDecoderModel from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint from .ltx2_transformer import ( LTX2VideoTransformer3DModel, @@ -69,6 +76,18 @@ _LTX_COMPONENT_SUBFOLDERS = ( ) logger = logging.getLogger(__name__) +_LTX2_CONV_VAE_EXTRA = "ltx2_use_conv_vae" +_LTX2_DIFFUSION_DECODER_SUBFOLDER = "diffusion_decoder" + + +def _ltx2_use_diffusion_decoder(od_config: Any, model_version: str) -> bool: + """Select DiffVAE by default for LTX-2.5, with an explicit ConvVAE opt-in.""" + extras = getattr(od_config, "extras", {}) or {} + use_conv_vae = extras.get(_LTX2_CONV_VAE_EXTRA, False) + if not isinstance(use_conv_vae, bool): + raise TypeError(f"{_LTX2_CONV_VAE_EXTRA} must be a bool, got {type(use_conv_vae)!r}") + return model_version == "2.5" and not use_conv_vae + @dataclass(frozen=True) class LTXComponentProfile: @@ -84,6 +103,7 @@ class LTXComponentProfile: text_encoder_cls: type | None = Gemma3ForConditionalGeneration vocoder_fallback_cls: type | None = None artifact_repo_id: str | None = None + artifact_revision: str | None = None latent_upsampler_filename: str | None = None distilled_lora_filename: str | None = None transformer_subfolder: str = "transformer" @@ -202,6 +222,7 @@ LTX25_TWO_STAGE_COMPONENT_PROFILE = replace( name="ltx2_5_two_stage", resident_modules=(*LTX25_FULL_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), artifact_repo_id="Lightricks/LTX-2.5", + artifact_revision=LTX25_NATIVE_ARTIFACT_REVISION, latent_upsampler_filename=("latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"), distilled_lora_filename="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors", ) @@ -239,14 +260,30 @@ def resolve_ltx_artifact( model: str, repo_id: str, filename: str, + *, + model_revision: str | None, + artifact_revision: str | None, ) -> str: - """Resolve an official LTX sidecar from the model root or its Hub repository.""" + """Resolve an official LTX sidecar without crossing repository revisions. + + A local model path selects where the primary components are loaded from; + it does not imply that independently hosted sidecars must be offline. Hub + offline behavior remains controlled by huggingface_hub itself. + """ candidate = Path(model) / filename if candidate.is_file(): return str(candidate) + # Hub revisions are repository-scoped. Reuse the model revision only when + # the model and artifact are in the same repository; otherwise use the + # independently pinned artifact revision. + revision = model_revision if model == repo_id else artifact_revision try: - return hf_hub_download(repo_id=repo_id, filename=filename) + return hf_hub_download( + repo_id=repo_id, + filename=filename, + revision=revision, + ) except Exception as exc: raise FileNotFoundError( f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; " @@ -254,6 +291,18 @@ def resolve_ltx_artifact( ) from exc +def _create_ltx25_natten_processor() -> LTX2VideoVaeNeighborhoodNattenProcessor: + """Create the required production attention backend with an actionable error.""" + try: + return LTX2VideoVaeNeighborhoodNattenProcessor() + except (ImportError, OSError, RuntimeError, ValueError) as exc: + raise RuntimeError( + "LTX-2.5 DiffVAE requires the shi-labs/natten Hub kernel. " + "Install kernels==0.14.1, use a supported GPU, leave " + "DIFFUSERS_DISABLE_REMOTE_CODE unset, and allow Hub access during kernel initialization." + ) from exc + + def _load_ltx_latent_upsampler_single_file(path: str, dtype: torch.dtype) -> LTX2LatentUpsamplerModel: """Load an official single-file upsampler into the Diffusers module.""" with safe_open(path, framework="pt", device="cpu") as handle: @@ -519,18 +568,49 @@ def _load_component( local_files_only: bool, dtype: torch.dtype, revision: str | None, + prefetch_list: tuple[str, ...] = _LTX_COMPONENT_SUBFOLDERS, ) -> Any: return from_pretrained_with_prefetch( component_cls.from_pretrained, model, subfolder=subfolder, - prefetch_list=_LTX_COMPONENT_SUBFOLDERS, + prefetch_list=prefetch_list, local_files_only=local_files_only, revision=revision, torch_dtype=dtype, ) +def _load_ltx25_native_diffusion_decoder( + model: str, + *, + local_files_only: bool, + dtype: torch.dtype, + revision: str | None, +) -> DistributedLTX2VideoDiffusionDecoderModel: + """Load canonical Native weights into the local Diffusers-compatible class.""" + config = DistributedLTX2VideoDiffusionDecoderModel.load_config( + model, + subfolder=_LTX2_DIFFUSION_DECODER_SUBFOLDER, + local_files_only=local_files_only, + revision=revision, + ) + checkpoint_path = resolve_ltx_artifact( + model, + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, + LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, + model_revision=revision, + artifact_revision=LTX25_NATIVE_ARTIFACT_REVISION, + ) + decoder = DistributedLTX2VideoDiffusionDecoderModel.from_ltx25_native_checkpoint( + checkpoint_path, + config, + dtype, + ) + decoder.init_distributed() + return decoder + + def _place_aux_components(pipeline: Any) -> None: parallel_config = getattr(pipeline.od_config, "parallel_config", None) use_managed_placement = bool( @@ -555,6 +635,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: model = od_config.model revision = getattr(od_config, "revision", None) local_files_only = os.path.exists(model) + use_diffusion_decoder = pipeline.use_diffusion_decoder pipeline.weights_sources = [ DiffusersPipelineLoader.ComponentSource( @@ -604,6 +685,27 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: dtype=dtype, revision=revision, ) + if use_diffusion_decoder: + pipeline.diffusion_decoder = _load_ltx25_native_diffusion_decoder( + model, + local_files_only=local_files_only, + dtype=dtype, + revision=revision, + ) + # The canonical decoder is trained and validated with NATTEN. Its + # portable FlexAttention fallback materializes an impractically large + # block mask at production video sizes, so fail early if the matching + # Hub kernel cannot be loaded instead of failing later during decode. + pipeline.diffusion_decoder.set_attn_processor(_create_ltx25_natten_processor()) + vae_patch_parallel_size = int( + getattr(getattr(od_config, "parallel_config", None), "vae_patch_parallel_size", 1) + ) + if vae_patch_parallel_size > 1 or getattr(od_config, "vae_use_tiling", False): + pipeline.diffusion_decoder.enable_tiling() + pipeline.diffusion_decoder.set_parallel_size( + vae_patch_parallel_size, + mode=getattr(getattr(od_config, "parallel_config", None), "vae_parallel_mode", "tile"), + ) pipeline.audio_vae = _load_component( AutoencoderKLLTX2Audio, model, @@ -656,6 +758,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: model, profile.artifact_repo_id, profile.latent_upsampler_filename, + model_revision=revision, + artifact_revision=profile.artifact_revision, ) pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) else: @@ -665,6 +769,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: model, profile.artifact_repo_id, profile.latent_upsampler_filename, + model_revision=revision, + artifact_revision=profile.artifact_revision, ) pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index 71beab643..8b4491e71 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Task-specific conditioning shared by LTX model versions.""" @@ -502,10 +502,12 @@ class LTXI2VConditioningMixin: init_latents = torch.cat(init_latents, dim=0).to(dtype=dtype) if num_videos_per_prompt > 1: init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0) + latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self) return latent_ops.normalize_latents( init_latents, - self.vae.latents_mean, - self.vae.latents_std, + latents_mean, + latents_std, + scaling_factor, ) @staticmethod @@ -570,11 +572,12 @@ class LTXI2VConditioningMixin: ) conditioning_mask = latents.new_zeros(mask_shape) conditioning_mask[:, :, 0] = 1.0 + latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self) latents = latent_ops.normalize_latents( latents, - self.vae.latents_mean, - self.vae.latents_std, - self.vae.config.scaling_factor, + latents_mean, + latents_std, + scaling_factor, ) clean_latents = latents if image is not None: diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py new file mode 100644 index 000000000..73db2ea6b --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved. +# +# Native conversion rules and decoder behavior copied and modified from +# Diffusers commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0). + +"""LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder.""" + +from collections.abc import Mapping +from typing import Any + +import torch +from diffusers.models.autoencoders import ( + LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel, +) +from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( + LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d, +) +from diffusers.models.autoencoders.vae import DecoderOutput +from safetensors import safe_open + +LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5" +LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" +LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors" + +_NATIVE_DECODER_PREFIXES = ("vae.decoder.", "decoder.") +_NATIVE_STATISTICS_KEYS = { + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", +} +_NATIVE_KEY_REPLACEMENTS = ( + ("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1."), + ("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2."), + (".attn.proj.", ".attn.to_out.0."), + (".attn.q_norm.", ".attn.norm_q."), + (".attn.k_norm.", ".attn.norm_k."), +) +_GATE_FOLD_TARGETS = { + ".attn.to_out.0.weight": ".gate_msa", + ".attn.to_out.0.bias": ".gate_msa", + ".mlp.w_down.weight": ".gate_mlp", + ".context_proj.weight": ".gate_ctx", + ".context_proj.bias": ".gate_ctx", +} +_GATE_SUFFIXES = tuple(_GATE_FOLD_TARGETS.values()) + + +def _strip_native_decoder_prefix(key: str) -> str | None: + for prefix in _NATIVE_DECODER_PREFIXES: + if key.startswith(prefix): + return key.removeprefix(prefix) + return None + + +def _native_statistics_target(key: str) -> str | None: + return _NATIVE_STATISTICS_KEYS.get(key.removeprefix("vae.")) + + +def _rename_native_decoder_key(key: str) -> str: + for source, target in _NATIVE_KEY_REPLACEMENTS: + key = key.replace(source, target) + return key + + +def _fold_native_gate(key: str, value: torch.Tensor, gates: Mapping[str, torch.Tensor]) -> torch.Tensor: + target_suffix = next((suffix for suffix in _GATE_FOLD_TARGETS if key.endswith(suffix)), None) + if target_suffix is None: + return value + gate_key = key[: -len(target_suffix)] + _GATE_FOLD_TARGETS[target_suffix] + gate = gates.get(gate_key) + if gate is None: + return value + gate = gate.to(device=value.device, dtype=torch.float32) + value_float = value.to(dtype=torch.float32) + folded = gate.unsqueeze(1) * value_float if value.ndim == 2 else gate * value_float + return folded.to(dtype=value.dtype) + + +def convert_ltx25_native_diffusion_decoder_state_dict( + native_state_dict: Mapping[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Convert the canonical Native LTX-2.5 DiffVAE state dict to Diffusers.""" + decoder_state_dict: dict[str, torch.Tensor] = {} + converted: dict[str, torch.Tensor] = {} + for key, value in native_state_dict.items(): + statistics_target = _native_statistics_target(key) + if statistics_target is not None: + converted[statistics_target] = value + continue + decoder_key = _strip_native_decoder_prefix(key) + if decoder_key is not None: + decoder_state_dict[decoder_key] = value + + gates = {key: value for key, value in decoder_state_dict.items() if key.endswith(_GATE_SUFFIXES)} + for key, value in decoder_state_dict.items(): + # ``type_emb`` is an unused shipping-checkpoint residue. Preview heads + # and legacy static gates are not parameters of the runtime decoder. + if key == "type_emb" or key.startswith("coarse_") or ".coarse_" in key or key in gates: + continue + + converted_key = _rename_native_decoder_key(key) + value = _fold_native_gate(converted_key, value, gates) + if converted_key.endswith((".qkv.weight", ".qkv.bias")): + leaf = "weight" if converted_key.endswith(".weight") else "bias" + prefix = converted_key[: -len(f"qkv.{leaf}")] + if value.shape[0] % 3 != 0: + raise ValueError( + f"Fused LTX-2.5 DiffVAE parameter {key!r} has leading dimension " + f"{value.shape[0]}, which is not divisible by 3." + ) + chunk = value.shape[0] // 3 + converted[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone() + converted[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone() + converted[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone() + continue + + converted[f"decoder.{converted_key}"] = value + + return converted + + +def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch.Tensor]: + """Read and convert only DiffVAE tensors from the canonical Native VAE file.""" + with safe_open(path, framework="pt", device="cpu") as handle: + native_state_dict = { + key: handle.get_tensor(key) + for key in handle.keys() + if _strip_native_decoder_prefix(key) is not None or _native_statistics_target(key) is not None + } + return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict) + + +class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d): + """Diffusers decoder core with the short-clip NATTEN context fix.""" + + def forward_stage_4( + self, + hidden_states: torch.Tensor, + drop_leading_frame: bool = True, + crop_trailing_ghost: bool = True, + ) -> torch.Tensor: + blocks = self.det_stages[-1] + block_mask = blocks[0].attn.build_block_mask(hidden_states) + for block in blocks: + hidden_states = block(hidden_states, block_mask) + hidden_states = self.upsamples[-1](hidden_states, drop_leading_frame=drop_leading_frame) + + num_pad = self.trailing_pad_latent_frames + if crop_trailing_ghost and num_pad > 0: + content_frames = max(hidden_states.shape[1] - num_pad * self.temporal_compression_ratio, 1) + keep_frames = min(hidden_states.shape[1], max(content_frames, self.stage5_kernel[0])) + hidden_states = hidden_states[:, :keep_frames] + return hidden_states + + +class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel): + """Diffusers 0.40 decoder plus canonical Native loading and short-clip support.""" + + def __init__( + self, + out_channels: int = 3, + latent_channels: int = 128, + patch_size: int = 4, + scaling_factor: float = 1.0, + decoder_head_dim: int = 64, + decoder_stage_channels: tuple[int, ...] = (2048, 1024, 512, 512, 256), + decoder_stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8), + decoder_stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)), + decoder_upsample_strides: tuple[tuple[int, int, int], ...] = ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)), + decoder_upsample_channel_reductions: tuple[int, ...] = (2, 2, 1, 2), + decoder_stage5_kernel: tuple[int, int, int] = (11, 11, 11), + decoder_t_emb_dim: int = 384, + decoder_timestep_scale_multiplier: float = 1000.0, + decoder_model_output_type: str = "x0", + decoder_num_inference_steps: int = 1, + spatial_compression_ratio: int = 32, + temporal_compression_ratio: int = 8, + ) -> None: + super().__init__( + out_channels=out_channels, + latent_channels=latent_channels, + patch_size=patch_size, + scaling_factor=scaling_factor, + decoder_head_dim=decoder_head_dim, + decoder_stage_channels=decoder_stage_channels, + decoder_stage_depths=decoder_stage_depths, + decoder_stage_kernels=decoder_stage_kernels, + decoder_upsample_strides=decoder_upsample_strides, + decoder_upsample_channel_reductions=decoder_upsample_channel_reductions, + decoder_stage5_kernel=decoder_stage5_kernel, + decoder_t_emb_dim=decoder_t_emb_dim, + decoder_timestep_scale_multiplier=decoder_timestep_scale_multiplier, + decoder_model_output_type=decoder_model_output_type, + decoder_num_inference_steps=decoder_num_inference_steps, + spatial_compression_ratio=spatial_compression_ratio, + temporal_compression_ratio=temporal_compression_ratio, + ) + # Diffusers constructs its private decoder core internally. Switching + # to this behavior-only subclass preserves parameters and state-dict keys. + self.decoder.__class__ = LTX2VideoDiffusionDecoder3d + self.decoder.stage5_kernel = tuple(decoder_stage5_kernel) + + @classmethod + def from_ltx25_native_checkpoint( + cls, + checkpoint_path: str, + config: Mapping[str, Any], + dtype: torch.dtype, + ) -> "LTX2VideoDiffusionDecoderModel": + """Construct this decoder class directly from the canonical Native checkpoint.""" + with torch.device("meta"): + model = cls.from_config(dict(config)) + state_dict = load_ltx25_native_diffusion_decoder_state_dict(checkpoint_path) + try: + model.load_state_dict(state_dict, strict=True, assign=True) + except RuntimeError as exc: + raise ValueError(f"Invalid LTX-2.5 Native DiffVAE checkpoint {checkpoint_path!r}.") from exc + model.to(device="cpu", dtype=dtype) + return model + + def decode( + self, + z: torch.Tensor, + generator: torch.Generator | None = None, + num_inference_steps: int | None = None, + return_dict: bool = True, + ) -> DecoderOutput | tuple[torch.Tensor]: + decoded = super().decode( + z, + generator=generator, + num_inference_steps=num_inference_steps, + return_dict=False, + )[0] + + # Short clips retain replicated trailing context through stage 5. Do + # not expose those context pixels to the caller. + target_num_frames = (z.shape[2] - 1) * self.temporal_compression_ratio + 1 + target_height = z.shape[3] * self.spatial_compression_ratio + target_width = z.shape[4] * self.spatial_compression_ratio + decoded = decoded[:, :, :target_num_frames, :target_height, :target_width] + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py new file mode 100644 index 000000000..7f486dbdf --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved. +# +# Distributed tiling is copied and modified from Diffusers' serial LTX-2.5 +# tiling at commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0). + +"""LTX-2.5-specific distributed execution for the diffusion VAE decoder.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from diffusers.utils import logging +from diffusers.utils.torch_utils import randn_tensor + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> list[tuple[int, int]]: + """Build overlapping intervals, merging a too-small trailing remnant.""" + if length <= tile_size: + return [(0, length)] + starts = list(range(0, length, stride)) + while len(starts) > 1 and length - starts[-1] < min_size: + starts.pop() + return [(start, min(start + tile_size, length)) for start in starts[:-1]] + [(starts[-1], length)] + + +@dataclass(frozen=True) +class LTX2VideoDiffusionTilePlan: + """Geometry shared by all ranks for one distributed DiffVAE decode.""" + + temporal_tiles: tuple[tuple[int, int], ...] + height_tiles: tuple[tuple[int, int], ...] + width_tiles: tuple[tuple[int, int], ...] + num_frames: int + height: int + width: int + scale_t: int + scale_h: int + scale_w: int + stride_t: int + stride_h: int + stride_w: int + blend_frames: int + blend_height: int + blend_width: int + single_step_x0: bool + + +@dataclass +class LTX2VideoDiffusionTileTask(TileTask): + """One stage-4 + stage-5 tile and its rank-independent noise.""" + + drop_leading_frame: bool = False + crop_trailing_ghost: bool = False + noise_generator: torch.Generator | list[torch.Generator] | None = None + x_t: torch.Tensor | None = None + + +class DistributedLTX2VideoDiffusionDecoderModel(LTX2VideoDiffusionDecoderModel, DistributedVaeMixin): + """LTX-2.5 diffusion decoder with model-specific tile parallelism. + + Stages 1-3 run over the complete low-resolution feature volume on every + participating rank. Stage 4 and the diffusion stage run as independent + overlapping tile tasks, and rank 0 performs the reference blend/merge. + """ + + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None: + if mode != "tile": + raise ValueError(f"LTX-2.5 DiffVAE only supports vae_parallel_mode='tile', got {mode!r}.") + super().set_parallel_size(parallel_size, mode=mode) + + def _build_tiled_decode_plan( + self, + features: torch.Tensor, + num_inference_steps: int, + ) -> LTX2VideoDiffusionTilePlan: + """Build the rank-independent tile schedule around Diffusers' decoder.""" + decoder = self.decoder + patch_size = decoder.patch_size + upsample_stride = decoder.upsamples[-1].stride + scale_t, scale_h, scale_w = ( + upsample_stride[0], + upsample_stride[1] * patch_size, + upsample_stride[2] * patch_size, + ) + tile_t = self.tile_sample_min_num_frames // scale_t + tile_h = self.tile_sample_min_height // scale_h + tile_w = self.tile_sample_min_width // scale_w + stride_t = self.tile_sample_stride_num_frames // scale_t + stride_h = self.tile_sample_stride_height // scale_h + stride_w = self.tile_sample_stride_width // scale_w + min_sizes = [ + max(kernel_4, -(-kernel_5 // stride)) + for kernel_4, kernel_5, stride in zip( + self.config.decoder_stage_kernels[-1], + self.config.decoder_stage5_kernel, + upsample_stride, + ) + ] + + ghost_frames = decoder.trailing_pad_latent_frames * math.prod( + upsample.stride[0] for upsample in decoder.upsamples[:-1] + ) + num_frames = features.shape[1] - ghost_frames + height, width = features.shape[2:4] + return LTX2VideoDiffusionTilePlan( + temporal_tiles=tuple(_tile_intervals(num_frames, tile_t, stride_t, min_sizes[0])), + height_tiles=tuple(_tile_intervals(height, tile_h, stride_h, min_sizes[1])), + width_tiles=tuple(_tile_intervals(width, tile_w, stride_w, min_sizes[2])), + num_frames=num_frames, + height=height, + width=width, + scale_t=scale_t, + scale_h=scale_h, + scale_w=scale_w, + stride_t=stride_t, + stride_h=stride_h, + stride_w=stride_w, + blend_frames=(tile_t - stride_t) * scale_t, + blend_height=(tile_h - stride_h) * scale_h, + blend_width=(tile_w - stride_w) * scale_w, + single_step_x0=num_inference_steps == 1 and decoder.model_output_type == "x0", + ) + + @staticmethod + def _iter_tiled_decode_coords(plan: LTX2VideoDiffusionTilePlan): + for t_idx in range(len(plan.temporal_tiles)): + for h_idx in range(len(plan.height_tiles)): + for w_idx in range(len(plan.width_tiles)): + yield t_idx, h_idx, w_idx + + @staticmethod + def _get_tiled_feature_slice( + features: torch.Tensor, + coord: tuple[int, int, int], + plan: LTX2VideoDiffusionTilePlan, + ) -> tuple[torch.Tensor, bool, bool]: + t_idx, h_idx, w_idx = coord + t0, t1 = plan.temporal_tiles[t_idx] + h0, h1 = plan.height_tiles[h_idx] + w0, w1 = plan.width_tiles[w_idx] + is_origin = t0 == 0 + is_trailing = t1 == plan.num_frames + feature_t1 = features.shape[1] if is_trailing else t1 + return features[:, t0:feature_t1, h0:h1, w0:w1], is_origin, is_trailing + + def _tiled_pixel_shape_from_features( + self, + feature_tile: torch.Tensor, + *, + drop_leading_frame: bool, + crop_trailing_ghost: bool, + ) -> tuple[int, int, int, int, int]: + """Predict one output tile's shape without materializing stage 4.""" + decoder = self.decoder + stride_t, stride_h, stride_w = decoder.upsamples[-1].stride + num_frames = feature_tile.shape[1] * stride_t + if stride_t == 2 and drop_leading_frame: + num_frames -= 1 + if crop_trailing_ghost and decoder.trailing_pad_latent_frames > 0: + content_frames = max( + num_frames - decoder.trailing_pad_latent_frames * decoder.temporal_compression_ratio, + 1, + ) + num_frames = min(num_frames, max(content_frames, decoder.stage5_kernel[0])) + return ( + feature_tile.shape[0], + decoder.out_channels, + num_frames, + feature_tile.shape[2] * stride_h * decoder.patch_size, + feature_tile.shape[3] * stride_w * decoder.patch_size, + ) + + @staticmethod + def _slice_tiled_noise( + x_t_full: torch.Tensor, + coord: tuple[int, int, int], + tile_pixel_shape: tuple[int, int, int, int, int], + plan: LTX2VideoDiffusionTilePlan, + ) -> torch.Tensor: + t_idx, h_idx, w_idx = coord + t0, _ = plan.temporal_tiles[t_idx] + h0, _ = plan.height_tiles[h_idx] + w0, _ = plan.width_tiles[w_idx] + pixel_t0 = t0 * plan.scale_t - (1 if t0 != 0 and plan.scale_t == 2 else 0) + return x_t_full[ + :, + :, + pixel_t0 : pixel_t0 + tile_pixel_shape[2], + h0 * plan.scale_h : h0 * plan.scale_h + tile_pixel_shape[3], + w0 * plan.scale_w : w0 * plan.scale_w + tile_pixel_shape[4], + ] + + def _merge_tiled_decode( + self, + tiles: dict[tuple[int, int, int], torch.Tensor], + plan: LTX2VideoDiffusionTilePlan, + ) -> torch.Tensor: + """Blend distributed RGB tiles in Diffusers' serial order.""" + frame_groups = [] + for t_idx in range(len(plan.temporal_tiles)): + rows = [ + [tiles[(t_idx, h_idx, w_idx)] for w_idx in range(len(plan.width_tiles))] + for h_idx in range(len(plan.height_tiles)) + ] + result_rows = [] + for h_idx, row in enumerate(rows): + result_row = [] + for w_idx, tile in enumerate(row): + if h_idx > 0: + tile = self.blend_v(rows[h_idx - 1][w_idx], tile, plan.blend_height) + if w_idx > 0: + tile = self.blend_h(row[w_idx - 1], tile, plan.blend_width) + keep_height = plan.stride_h * plan.scale_h if h_idx < len(rows) - 1 else tile.shape[3] + keep_width = plan.stride_w * plan.scale_w if w_idx < len(row) - 1 else tile.shape[4] + result_row.append(tile[:, :, :, :keep_height, :keep_width]) + result_rows.append(torch.cat(result_row, dim=4)) + frame_groups.append(torch.cat(result_rows, dim=3)) + + result = [] + for t_idx, group in enumerate(frame_groups): + if t_idx > 0: + group = self.blend_t(frame_groups[t_idx - 1], group, plan.blend_frames) + if t_idx < len(frame_groups) - 1: + keep_frames = plan.stride_t * plan.scale_t - (1 if t_idx == 0 and plan.scale_t == 2 else 0) + group = group[:, :, :keep_frames] + result.append(group) + return torch.cat(result, dim=2) + + def _default_generator(self, device: torch.device) -> torch.Generator: + if device.type == "cpu": + return torch.default_generator + device_module = getattr(torch, device.type, None) + default_generators = getattr(device_module, "default_generators", None) + if default_generators is None: + raise ValueError( + f"Distributed LTX-2.5 diffusion decode on {device.type!r} requires an explicit torch.Generator." + ) + device_index = device.index + if device_index is None: + current_device = getattr(device_module, "current_device", None) + device_index = current_device() if current_device is not None else 0 + return default_generators[device_index] + + def _sync_generators( + self, + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, + ) -> torch.Generator | list[torch.Generator]: + """Make every rank start tile-noise generation from rank 0's state.""" + generators = self._default_generator(device) if generator is None else generator + generator_list = generators if isinstance(generators, list) else [generators] + for item in generator_list: + state_on_device = item.get_state().to(device=device) + dist.broadcast(state_on_device, src=0, group=self.distributed_executor.group) + item.set_state(state_on_device.cpu()) + return generators + + @staticmethod + def _clone_generator(generator: torch.Generator) -> torch.Generator: + cloned = torch.Generator(device=generator.device) + cloned.set_state(generator.get_state()) + return cloned + + def _clone_generators( + self, + generators: torch.Generator | list[torch.Generator], + ) -> torch.Generator | list[torch.Generator]: + if isinstance(generators, list): + return [self._clone_generator(item) for item in generators] + return self._clone_generator(generators) + + def _distributed_tile_split( + self, + z: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None, + num_inference_steps: int, + ) -> tuple[list[LTX2VideoDiffusionTileTask], GridSpec]: + decoder = self.decoder + features = decoder.forward_stages_1_to_3(z) + plan = self._build_tiled_decode_plan(features, num_inference_steps) + generators = self._sync_generators(generator, z.device) + + x_t_full = None + if not plan.single_step_x0: + pixel_frames = plan.num_frames * plan.scale_t - (1 if plan.scale_t == 2 else 0) + x_t_full = randn_tensor( + ( + z.shape[0], + decoder.out_channels, + pixel_frames, + plan.height * plan.scale_h, + plan.width * plan.scale_w, + ), + generator=generators, + device=z.device, + dtype=z.dtype, + ) + + tasks = [] + for tile_id, coord in enumerate(self._iter_tiled_decode_coords(plan)): + feature_tile, is_origin, is_trailing = self._get_tiled_feature_slice(features, coord, plan) + tile_pixel_shape = self._tiled_pixel_shape_from_features( + feature_tile, + drop_leading_frame=is_origin, + crop_trailing_ghost=is_trailing, + ) + tile_generator = None + x_t = None + if plan.single_step_x0: + # Save the canonical serial state for this tile, then advance + # the request generator exactly as the single-rank tiled path + # does. Task placement therefore cannot change tile noise. + tile_generator = self._clone_generators(generators) + discarded_noise = randn_tensor( + tile_pixel_shape, + generator=generators, + device=z.device, + dtype=z.dtype, + ) + del discarded_noise + else: + x_t = self._slice_tiled_noise(x_t_full, coord, tile_pixel_shape, plan) + + tasks.append( + LTX2VideoDiffusionTileTask( + tile_id=tile_id, + grid_coord=coord, + tensor=feature_tile, + workload=math.prod(tile_pixel_shape), + drop_leading_frame=is_origin, + crop_trailing_ghost=is_trailing, + noise_generator=tile_generator, + x_t=x_t, + ) + ) + + return tasks, GridSpec( + split_dims=(2, 3, 4), + grid_shape=(len(plan.temporal_tiles), len(plan.height_tiles), len(plan.width_tiles)), + tile_spec={"plan": plan}, + output_dtype=z.dtype, + ) + + def _distributed_tile_exec( + self, + task: LTX2VideoDiffusionTileTask, + num_inference_steps: int, + ) -> torch.Tensor: + context = self.decoder.forward_stage_4( + task.tensor, + drop_leading_frame=task.drop_leading_frame, + crop_trailing_ghost=task.crop_trailing_ghost, + ) + tile_pixel_shape = ( + context.shape[0], + self.decoder.out_channels, + context.shape[1], + context.shape[2] * self.decoder.patch_size, + context.shape[3] * self.decoder.patch_size, + ) + if task.x_t is None: + x_t = randn_tensor( + tile_pixel_shape, + generator=task.noise_generator, + device=context.device, + dtype=context.dtype, + ) + else: + x_t = task.x_t + return self.decoder.denoise(context, x_t, num_inference_steps) + + def _distributed_tile_merge( + self, + tiles: dict[tuple[int, int, int], torch.Tensor], + grid_spec: GridSpec, + ) -> torch.Tensor: + plan = grid_spec.tile_spec["plan"] + if not isinstance(plan, LTX2VideoDiffusionTilePlan): + raise TypeError(f"Expected an LTX2VideoDiffusionTilePlan, got {type(plan)!r}.") + return self._merge_tiled_decode(tiles, plan) + + def tiled_decode( + self, + z: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, + num_inference_steps: int | None = None, + ) -> torch.Tensor: + if not self.is_distributed_enabled(): + return super().tiled_decode(z, generator=generator, num_inference_steps=num_inference_steps) + + logger.debug("LTX-2.5 diffusion decoder running with distributed stage-4/stage-5 tiles") + num_inference_steps = num_inference_steps or self.decoder.default_num_inference_steps + result = self.distributed_executor.execute( + z, + DistributedOperator( + split=lambda tensor: self._distributed_tile_split(tensor, generator, num_inference_steps), + exec=lambda task: self._distributed_tile_exec(task, num_inference_steps), + merge=self._distributed_tile_merge, + ), + broadcast_result=False, + ) + if result.numel() == 0: + # The base decode method crops five dimensions before the LTX + # runtime discards non-output-rank results. + return torch.empty((0, self.decoder.out_channels, 0, 0, 0), device=z.device, dtype=z.dtype) + return result diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py index 1b25a8e01..82de85504 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared latent layout and normalization primitives for LTX pipelines.""" @@ -69,6 +69,13 @@ def unpack_latents( return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) +def resolve_video_latent_statistics(pipeline: Any) -> tuple[torch.Tensor, torch.Tensor, float]: + """Return statistics owned by the active video decoder.""" + decoder = pipeline.diffusion_decoder if getattr(pipeline, "use_diffusion_decoder", False) else pipeline.vae + scaling_factor = getattr(getattr(decoder, "config", None), "scaling_factor", 1.0) + return decoder.latents_mean, decoder.latents_std, scaling_factor + + def normalize_latents( latents: torch.Tensor, latents_mean: torch.Tensor, @@ -202,11 +209,12 @@ def prepare_video_latents( ) -> torch.Tensor: if latents is not None: if latents.ndim == 5: + latents_mean, latents_std, scaling_factor = resolve_video_latent_statistics(pipeline) latents = normalize_latents( latents, - pipeline.vae.latents_mean, - pipeline.vae.latents_std, - pipeline.vae.config.scaling_factor, + latents_mean, + latents_std, + scaling_factor, ) latents = pack_latents( latents, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py index 50d413ee7..6b0433bbe 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Construction and execution of a fixed LTX refinement-phase adapter. @@ -420,6 +420,8 @@ def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None: pipeline.od_config.model, profile.artifact_repo_id, profile.distilled_lora_filename, + model_revision=getattr(pipeline.od_config, "revision", None), + artifact_revision=profile.artifact_revision, ) if layer_fused and pipeline.od_config.dtype is not torch.bfloat16: diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index 35e0cef0b..22230e486 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Request normalization and validation shared by LTX pipeline variants.""" @@ -86,6 +86,7 @@ def validate_pipeline_request( vae_spatial_compression_ratio: int, vae_temporal_compression_ratio: int, pipeline_name: str, + min_video_latent_spatial_size: tuple[int, int] | None = None, request_sigmas: list[float] | None = None, request_phase_sigmas: tuple[list[float] | None, ...] | None = None, ) -> None: @@ -97,6 +98,18 @@ def validate_pipeline_request( f"{request_inputs.width}x{request_inputs.height}." ) + if min_video_latent_spatial_size is not None: + min_latent_height, min_latent_width = min_video_latent_spatial_size + latent_height = request_inputs.height // vae_spatial_compression_ratio + latent_width = request_inputs.width // vae_spatial_compression_ratio + if latent_height < min_latent_height or latent_width < min_latent_width: + raise ValueError( + f"{pipeline_name} DiffVAE requires video latent spatial dimensions of at least " + f"{min_latent_width}x{min_latent_height} for neighborhood attention, got " + f"{latent_width}x{latent_height} from request resolution " + f"{request_inputs.width}x{request_inputs.height}." + ) + if request_inputs.num_frames < 1 or (request_inputs.num_frames - 1) % vae_temporal_compression_ratio != 0: raise ValueError( f"{pipeline_name} `num_frames` must be {vae_temporal_compression_ratio} * k + 1, " diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index d3addd432..fc589fcb8 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared recipe-driven runtime for LTX pipeline variants.""" @@ -29,6 +29,7 @@ from vllm_omni.platforms import current_omni_platform from . import ltx2_latents as latent_ops from .ltx2_components import ( LTXComponentProfile, + _ltx2_use_diffusion_decoder, detect_ltx_model_version, initialize_pipeline_components, resolve_ltx_component_profile, @@ -168,6 +169,7 @@ class LTXRuntime( "Use the default ulysses_mode='strict' for LTX sequence parallelism." ) self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None)) + self.use_diffusion_decoder = _ltx2_use_diffusion_decoder(od_config, self.model_version) self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version) self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version) if getattr(od_config, "cache_backend", "none") == "cache_dit" and not self.pipeline_recipe.supports_cache_dit: @@ -178,6 +180,8 @@ class LTXRuntime( self._dit_modules = list(self.component_profile.dit_modules) self._encoder_modules = list(self.component_profile.encoder_modules) self._vae_modules = list(self.component_profile.vae_modules) + if self.use_diffusion_decoder: + self._vae_modules.append("diffusion_decoder") self._resident_modules = list(self.component_profile.resident_modules) if self.model_version in ("2.3", "2.5"): self.preserve_sp_padded_audio_duration = True @@ -249,12 +253,17 @@ class LTXRuntime( raise ValueError(f"{self.__class__.__name__} does not support `image` input.") request_sigmas = self._resolve_request_sigmas(req, sigmas) request_phase_sigmas = self._resolve_request_phase_sigmas(req, stage_1_sigmas, stage_2_sigmas) + min_video_latent_spatial_size = None + if getattr(self, "use_diffusion_decoder", False): + first_stage_kernel = self.diffusion_decoder.config.decoder_stage_kernels[0] + min_video_latent_spatial_size = (int(first_stage_kernel[-2]), int(first_stage_kernel[-1])) validate_pipeline_request( request_inputs, pipeline_recipe=self.pipeline_recipe, vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, pipeline_name=self.__class__.__name__, + min_video_latent_spatial_size=min_video_latent_spatial_size, request_sigmas=request_sigmas, request_phase_sigmas=request_phase_sigmas, ) @@ -542,7 +551,10 @@ class LTXRuntime( return self._make_output((latents, audio_latents)) latents = latents.to(connector_prompt_embeds.dtype) - if not self.vae.config.timestep_conditioning: + use_diffusion_decoder = getattr(self, "use_diffusion_decoder", False) + if use_diffusion_decoder: + timestep_decode = None + elif not self.vae.config.timestep_conditioning: timestep_decode = None else: noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) @@ -559,14 +571,22 @@ class LTXRuntime( dist_initialized = torch.distributed.is_initialized() is_output_rank = not dist_initialized or torch.distributed.get_rank() == 0 vae_decode_needs_all_ranks = False - is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None) + video_decoder = self.diffusion_decoder if use_diffusion_decoder else self.vae + is_distributed_vae_enabled = getattr(video_decoder, "is_distributed_enabled", None) if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled): # Distributed tiled decode is collective, so every rank must enter it. vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled()) should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks if should_decode_video: - video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0] + if use_diffusion_decoder: + video = self.diffusion_decoder.decode( + latents.to(self.diffusion_decoder.dtype), + generator=generator, + return_dict=False, + )[0] + else: + video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0] else: video = torch.empty(0, device=latents.device, dtype=latents.dtype) @@ -800,11 +820,12 @@ class LTXRuntime( self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) + latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self) latents = latent_ops.denormalize_latents( latents, - self.vae.latents_mean, - self.vae.latents_std, - self.vae.config.scaling_factor, + latents_mean, + latents_std, + scaling_factor, ) audio_latents = latent_ops.unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index a2aa4e7ae..3c66290a9 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1043,6 +1043,17 @@ class AsyncOmniEngine: fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), ) + extras = dict(kwargs.get("extras") or {}) + for key, default in ( + ("auxiliary_text_encoder", None), + ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), + ): + top_level_value = kwargs.get(key) + if top_level_value is not None: + extras[key] = top_level_value + else: + extras.setdefault(key, default) + stage_engine_args = { "max_num_seqs": kwargs.get("max_num_seqs") or 1, "parallel_config": parallel_config, @@ -1103,10 +1114,7 @@ class AsyncOmniEngine: "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), "streaming_output": kwargs.get("diffusion_streaming_output", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), - "extras": { - "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None), - "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), - }, + "extras": extras, **( { "profiler_config": asdict(kwargs["profiler_config"]) @@ -1196,11 +1204,29 @@ class AsyncOmniEngine: # Parse --stage-overrides JSON string if provided stage_overrides = parse_stage_overrides(stage_overrides_json) + # Unregistered diffusion checkpoints use the single-stage fallback + # below instead of StageConfigFactory, so fold stage-0 model extras + # into the fallback's input as well. Registered pipelines still get + # the complete override mapping through load_and_resolve_stage_configs. + default_stage_kwargs = kwargs + stage_zero_overrides = (stage_overrides or {}).get("0", {}) + if "extras" in stage_zero_overrides: + override_extras = stage_zero_overrides["extras"] + if not isinstance(override_extras, Mapping): + raise TypeError("stage 0 extras must be a mapping") + default_stage_kwargs = { + **kwargs, + "extras": { + **(kwargs.get("extras") or {}), + **override_extras, + }, + } + config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs( model, kwargs, trust_remote_code=trust_remote_code, - default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), + default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(default_stage_kwargs), deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, strategy_config_path=strategy_config_path,