diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index 0f1f33a25..e30b627ee 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -296,14 +296,14 @@ done ```bash vllm serve Lightricks/LTX-2 --omni --port 8098 \ - --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 + --enforce-eager --boundary-ratio 1.0 ``` For multi-GPU memory reduction, you can enable HSDP: ```bash vllm serve Lightricks/LTX-2 --omni --port 8098 \ - --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 \ + --enforce-eager --boundary-ratio 1.0 \ --use-hsdp --hsdp-shard-size 2 ``` diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index 93de44a88..71afd9e7c 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/examples/online_serving/text_to_video/README.md @@ -254,7 +254,7 @@ done ```bash vllm serve Lightricks/LTX-2 --omni --port 8098 \ - --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 + --enforce-eager --boundary-ratio 1.0 ``` #### Start with Optimization Presets diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 2f9155fe0..41b1d5a84 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -17,7 +17,11 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import ( LTX23_COMPONENT_PROFILE, ) from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin -from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseExecutor, LTXPhaseResult +from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( + LTXDenoiseExecutor, + LTXPhaseResult, + prepare_scheduler_stage, +) from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( LTX_LEGACY_VELOCITY_GUIDANCE, LTX_OFFICIAL_X0_GUIDANCE, @@ -324,6 +328,38 @@ def test_ltx2_two_stage_reuses_prompt_context_between_phases(): torch.testing.assert_close(output.output[1], torch.tensor([4.0])) +def test_ltx_custom_sigmas_bypass_scheduler_shifting(): + from diffusers import FlowMatchEulerDiscreteScheduler + + scheduler = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=True, + base_shift=0.95, + max_shift=2.05, + shift_terminal=0.1, + ) + pipeline = SimpleNamespace( + scheduler=scheduler, + _make_video_audio_scheduler=lambda *args, **kwargs: object(), + ) + sigmas = [1.0, 0.75, 0.25, 0.0] + + audio_scheduler, _, timesteps = prepare_scheduler_stage( + pipeline, + SimpleNamespace(num_inference_steps=3), + device=torch.device("cpu"), + sigmas=sigmas, + timesteps=None, + latent_num_frames=1, + latent_height=1, + latent_width=1, + ) + + expected = torch.tensor(sigmas, dtype=torch.float32) + torch.testing.assert_close(pipeline.scheduler.sigmas, expected) + torch.testing.assert_close(audio_scheduler.sigmas, expected) + torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps) + + class TestLTXRequestParsing: @pytest.mark.parametrize( ("t2v_cls", "i2v_cls"), @@ -513,6 +549,7 @@ class TestLTX23ForwardStages: from vllm_omni.inputs.data import OmniDiffusionSamplingParams pipe = _make_ltx_request_pipe(LTX23Pipeline) + request_sigmas = [1.0, 0.4] req = DiffusionRequestBatch( [ OmniDiffusionRequest( @@ -525,12 +562,13 @@ class TestLTX23ForwardStages: num_inference_steps=3, guidance_scale=4.5, output_type="latent", + sigmas=request_sigmas, ), request_id="ltx23-t2v-forward-stage-delegation", ) ] ) - sigmas = [1.0, 0.5] + fallback_sigmas = [1.0, 0.5] timesteps = [1000, 500] attention_kwargs = {"scale": 1.0} seen = {} @@ -545,7 +583,7 @@ class TestLTX23ForwardStages: output = pipe.forward( req, - sigmas=sigmas, + sigmas=fallback_sigmas, timesteps=timesteps, noise_scale=0.25, attention_kwargs=attention_kwargs, @@ -561,7 +599,7 @@ class TestLTX23ForwardStages: assert seen["request_inputs"].output_type == "latent" assert seen["kwargs"] == { "noise_scale": 0.25, - "sigmas": sigmas, + "sigmas": request_sigmas, "timesteps": timesteps, "attention_kwargs": attention_kwargs, } diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 9b53d8941..c2807797d 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -82,6 +82,7 @@ class QualityTestConfig: gpu: str = "H100" # minimum GPU requirement negative_prompt: str = "" guidance_scale: float | None = None + sigmas: list[float] | None = None enable_cpu_offload: bool = False def baseline_ref(self) -> str: @@ -215,6 +216,19 @@ QUALITY_CONFIGS = [ width=256, num_frames=25, num_inference_steps=8, + # Preserve the final scheduler trajectory used when this FP8 quality + # gate was established. Explicit LTX sigmas bypass dynamic shifting. + sigmas=[ + 1.0, + 0.92185378074646, + 0.8327768445014954, + 0.7303033471107483, + 0.6111654043197632, + 0.4709382653236389, + 0.30347853899002075, + 0.10000002384185791, + 0.0, + ], ), ] @@ -322,6 +336,7 @@ def _generate_video(omni, config: QualityTestConfig): num_inference_steps=config.num_inference_steps, num_frames=config.num_frames, guidance_scale=config.guidance_scale, + sigmas=config.sigmas, ), ) @@ -446,6 +461,33 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch assert peak_mem == 0.0 +def test_generate_video_forwards_sigmas(monkeypatch): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False) + monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False) + monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False) + captured = SimpleNamespace(sampling=None) + + class DummyOmni: + def generate(self, _prompt, sampling): + captured.sampling = sampling + return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])] + + config = QualityTestConfig( + id="ltx-sigmas", + model="unused", + quantization="fp8", + task="t2v", + prompt="test", + max_lpips=0.1, + sigmas=[1.0, 0.5], + ) + _generate_video(DummyOmni(), config) + + assert captured.sampling.sigmas == [1.0, 0.5] + + _marks = hardware_marks(res={"cuda": "H100"}) _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index af8df4b41..d0dfa3e54 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -177,6 +177,17 @@ class I2VVideoAudioScheduler: return ((video_out, audio_out),) +def _set_scheduler_sigmas(scheduler: Any, sigmas: torch.Tensor) -> torch.Tensor: + sigmas = sigmas.to(torch.float32) + timesteps = sigmas[:-1] * scheduler.config.get("num_train_timesteps", 1000) + scheduler.sigmas = sigmas + scheduler.timesteps = timesteps + scheduler.num_inference_steps = len(timesteps) + scheduler._step_index = None + scheduler._begin_index = None + return timesteps + + def prepare_scheduler_stage( pipeline: Any, request_inputs: LTXRequestInputs, @@ -188,18 +199,9 @@ def prepare_scheduler_stage( latent_height: int, latent_width: int, ) -> tuple[Any, Any, torch.Tensor]: - sigmas = ( - np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps) - if sigmas is None - else sigmas - ) - mu = calculate_shift( - pipeline.scheduler.config.get("max_image_seq_len", 4096), - pipeline.scheduler.config.get("base_image_seq_len", 1024), - pipeline.scheduler.config.get("max_image_seq_len", 4096), - pipeline.scheduler.config.get("base_shift", 0.95), - pipeline.scheduler.config.get("max_shift", 2.05), - ) + if sigmas is not None and timesteps is not None: + raise ValueError("Only one of `sigmas` or `timesteps` may be provided.") + audio_scheduler = copy.deepcopy(pipeline.scheduler) video_audio_scheduler = pipeline._make_video_audio_scheduler( audio_scheduler, @@ -207,6 +209,27 @@ def prepare_scheduler_stage( latent_height, latent_width, ) + if sigmas is not None: + scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device) + if scheduler_sigmas.ndim != 1 or scheduler_sigmas.numel() < 2: + raise ValueError("An LTX custom sigma schedule must contain at least two boundary values.") + if scheduler_sigmas[-1] != 0: + scheduler_sigmas = torch.cat([scheduler_sigmas, scheduler_sigmas.new_zeros(1)]) + timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas) + _set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone()) + return audio_scheduler, video_audio_scheduler, timesteps_tensor + + sigmas = np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps) + # Official LTX2 one-stage omits the latent, so mu stays at max_shift across resolutions. + # Mu varies with latent token count only when the scheduler receives a latent, as in official HQ two-stage; + # explicit sigma schedules bypass mu entirely. + mu = calculate_shift( + pipeline.scheduler.config.get("max_image_seq_len", 4096), + pipeline.scheduler.config.get("base_image_seq_len", 1024), + pipeline.scheduler.config.get("max_image_seq_len", 4096), + pipeline.scheduler.config.get("base_shift", 0.95), + pipeline.scheduler.config.get("max_shift", 2.05), + ) retrieve_timesteps( audio_scheduler, request_inputs.num_inference_steps, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index 02cab34d2..d8a4356b3 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -71,6 +71,17 @@ class LTXRequestMixin: supports_request_batch = False supports_guidance_rescale = False + @staticmethod + def _resolve_request_sigmas( + req: DiffusionRequestBatch, + fallback: list[float] | None, + ) -> list[float] | None: + request_sigmas = [sampling.sigmas for sampling in req.sampling_params_list] + first = request_sigmas[0] if request_sigmas else None + if any(item != first for item in request_sigmas[1:]): + raise ValueError("Batched LTX requests must use identical custom sigmas.") + return first if first is not None else fallback + def check_inputs( self, prompt: str | list[str] | None, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index c32a6918a..5de1f4862 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -96,6 +96,7 @@ class LTXOneStagePipeline(LTXPipelineRuntime): image: Any | None = None, ) -> DiffusionOutput | list[DiffusionOutput]: del return_dict + sigmas = self._resolve_request_sigmas(req, sigmas) request_inputs = self._resolve_request_inputs( req, prompt=prompt,