diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_scheduler_finite_gate.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_scheduler_finite_gate.py new file mode 100644 index 000000000..8a90a0500 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_scheduler_finite_gate.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_scheduler_default_skips_large_finite_scans_and_preserves_results(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import ( + scheduling_minimax_h3_euler_ancestral as scheduler, + ) + + monkeypatch.delenv("VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS", raising=False) + original_isfinite = torch.isfinite + scanned_numels = [] + + def record_isfinite(tensor): + scanned_numels.append(tensor.numel()) + return original_isfinite(tensor) + + monkeypatch.setattr(torch, "isfinite", record_isfinite) + state = torch.linspace(-1.0, 1.0, 8) + velocity = torch.linspace(0.25, -0.25, 8) + timestep = torch.tensor(0.5) + default_x0 = scheduler.minimax_h3_rf_v_to_x0(state, velocity, timestep) + default = scheduler.minimax_h3_euler_eta0_step( + state, + default_x0, + sigma_curr=0.5, + sigma_next=0.25, + ) + + assert scanned_numels == [1] + + checked_x0 = scheduler.minimax_h3_rf_v_to_x0( + state, + velocity, + timestep, + validate_finite_tensors=True, + ) + checked = scheduler.minimax_h3_euler_eta0_step( + state, + checked_x0, + sigma_curr=0.5, + sigma_next=0.25, + validate_finite_tensors=True, + ) + torch.testing.assert_close(default_x0, checked_x0, rtol=0, atol=0) + torch.testing.assert_close(default, checked, rtol=0, atol=0) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")]) +def test_scheduler_finite_gate_rejects_nonfinite_inputs(monkeypatch, invalid_value): + from vllm_omni.diffusion.models.minimax_h3 import ( + scheduling_minimax_h3_euler_ancestral as scheduler, + ) + + monkeypatch.setenv("VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS", "1") + invalid = torch.ones(8) + invalid[0] = invalid_value + + with pytest.raises(ValueError, match="must be finite"): + scheduler.minimax_h3_rf_v_to_x0( + invalid, + torch.ones_like(invalid), + torch.tensor(0.5), + ) + with pytest.raises(ValueError, match="must be finite"): + scheduler.minimax_h3_euler_eta0_step( + invalid, + torch.ones_like(invalid), + sigma_curr=0.5, + sigma_next=0.25, + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index b051951a7..45f1e653e 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -17,6 +17,7 @@ import torch from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, + minimax_h3_validate_finite_tensors_enabled, ) MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 @@ -187,6 +188,7 @@ def minimax_h3_denoise_loop( if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor + validate_finite_tensors = minimax_h3_validate_finite_tensors_enabled() num_steps = len(sigmas_video) - 1 for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() @@ -214,8 +216,15 @@ def minimax_h3_denoise_loop( video_rows[update], mv_video_t, torch.tensor(t_v, dtype=torch.float32, device=device), + validate_finite_tensors=validate_finite_tensors, + ) + new_target = minimax_h3_euler_eta0_step( + video_rows[update], + x0_video, + sigma_curr=s_v, + sigma_next=s_v_next, + validate_finite_tensors=validate_finite_tensors, ) - new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next) video_rows = video_rows.clone() video_rows[update] = new_target if cond_anchor is not None: @@ -225,9 +234,14 @@ def minimax_h3_denoise_loop( audio_rows[audio_update], mv_audio_t, torch.tensor(t_a, dtype=torch.float32, device=device), + validate_finite_tensors=validate_finite_tensors, ) new_audio = minimax_h3_euler_eta0_step( - audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next + audio_rows[audio_update], + x0_audio, + sigma_curr=s_a, + sigma_next=s_a_next, + validate_finite_tensors=validate_finite_tensors, ) audio_rows = audio_rows.clone() audio_rows[audio_update] = new_audio diff --git a/vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py b/vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py index 9a3a1c9bc..4cd0586b6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py +++ b/vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py @@ -2,10 +2,24 @@ from __future__ import annotations import math +import os from typing import Any import torch +MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV = "VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS" + + +def minimax_h3_validate_finite_tensors_enabled() -> bool: + value = os.getenv(MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV, "") + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _resolve_validate_finite_tensors(value: bool | None) -> bool: + if value is None: + return minimax_h3_validate_finite_tensors_enabled() + return bool(value) + def _require_finite_tensor(tensor: torch.Tensor, name: str) -> None: if not bool(torch.isfinite(tensor).all().item()): @@ -50,6 +64,8 @@ def minimax_h3_rf_v_to_x0( xt: torch.Tensor, v: torch.Tensor, timestep: torch.Tensor, + *, + validate_finite_tensors: bool | None = None, ) -> torch.Tensor: if xt.shape != v.shape: raise ValueError(f"xt and v shapes must match, got {xt.shape} vs {v.shape}") @@ -57,15 +73,18 @@ def minimax_h3_rf_v_to_x0( raise ValueError("xt must be a floating point tensor") if not torch.is_floating_point(v): raise ValueError("v must be a floating point tensor") - _require_finite_tensor(xt, "xt") - _require_finite_tensor(v, "v") + validate_finite_tensors = _resolve_validate_finite_tensors(validate_finite_tensors) + if validate_finite_tensors: + _require_finite_tensor(xt, "xt") + _require_finite_tensor(v, "v") _validate_unit_timestep(timestep, "timestep") cond_t = timestep.to(device=xt.device, dtype=xt.dtype) while cond_t.ndim < xt.ndim: cond_t = cond_t.unsqueeze(-1) sigma_t = 1 - cond_t x0 = xt + sigma_t * v - _require_finite_tensor(x0, "x0") + if validate_finite_tensors: + _require_finite_tensor(x0, "x0") return x0 @@ -75,6 +94,7 @@ def minimax_h3_euler_eta0_step( *, sigma_curr: float, sigma_next: float, + validate_finite_tensors: bool | None = None, ) -> torch.Tensor: if state.shape != denoised.shape: raise ValueError(f"state and denoised shapes must match, got {state.shape} vs {denoised.shape}") @@ -82,8 +102,10 @@ def minimax_h3_euler_eta0_step( raise ValueError("state must be a floating point tensor") if not torch.is_floating_point(denoised): raise ValueError("denoised must be a floating point tensor") - _require_finite_tensor(state, "state") - _require_finite_tensor(denoised, "denoised") + validate_finite_tensors = _resolve_validate_finite_tensors(validate_finite_tensors) + if validate_finite_tensors: + _require_finite_tensor(state, "state") + _require_finite_tensor(denoised, "denoised") sigma_curr = _validate_sigma(sigma_curr, "sigma_curr") sigma_next = _validate_sigma(sigma_next, "sigma_next") if sigma_curr == 0.0: @@ -98,7 +120,8 @@ def minimax_h3_euler_eta0_step( sigma_ratio = sigma_next_t / sigma_curr_t out = sigma_ratio * state.to(dtype=compute_dtype) + (1.0 - sigma_ratio) * denoised.to(dtype=compute_dtype) out = out.to(dtype=state.dtype) - _require_finite_tensor(out, "euler_eta0_step output") + if validate_finite_tensors: + _require_finite_tensor(out, "euler_eta0_step output") return out @@ -127,6 +150,7 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter: audio_sigma_curr: float | None = None, audio_sigma_next: float | None = None, ) -> dict[str, torch.Tensor]: + validate_finite_tensors = minimax_h3_validate_finite_tensors_enabled() visual_timestep = timestep if video_timestep is None else video_timestep audio_timestep = timestep if audio_timestep is None else audio_timestep visual_sigma_curr = sigma_curr if video_sigma_curr is None else video_sigma_curr @@ -148,11 +172,13 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter: input_visual_latent, noise_pred_visual, visual_timestep, + validate_finite_tensors=validate_finite_tensors, ) denoised_audio = minimax_h3_rf_v_to_x0( input_audio_latent, noise_pred_audio, audio_timestep, + validate_finite_tensors=validate_finite_tensors, ) return { "output_visual_latent": minimax_h3_euler_eta0_step( @@ -160,12 +186,14 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter: denoised_visual, sigma_curr=visual_sigma_curr, sigma_next=visual_sigma_next, + validate_finite_tensors=validate_finite_tensors, ), "output_audio_latent": minimax_h3_euler_eta0_step( input_audio_latent, denoised_audio, sigma_curr=audio_sigma_curr, sigma_next=audio_sigma_next, + validate_finite_tensors=validate_finite_tensors, ), } @@ -173,7 +201,9 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter: EntryClass = MiniMaxH3EulerAncestralEta0SchedulerAdapter __all__ = [ + "MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV", "MiniMaxH3EulerAncestralEta0SchedulerAdapter", "minimax_h3_euler_eta0_step", "minimax_h3_rf_v_to_x0", + "minimax_h3_validate_finite_tensors_enabled", ]