diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 5bfd65649..9b3a87956 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -988,6 +988,42 @@ settings, and both videos with any comparison; the command alone is not quality evidence. A matching seed does not imply matching noise over differently sized full and windowed latents. +## Request-scoped sampler + +H3 supports `euler` (the default) and `res_multistep`. Select the sampler per +request; it applies to both video and audio, using their separate sigma schedules. +For the HTTP examples above, set `num_inference_steps=20` and add +`"sampler":"res_multistep"` to the existing `extra_params` object. For example, +the T2VA fields become: + +```bash +-F 'num_inference_steps=20' \ +-F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0,"sampler":"res_multistep"}' +``` + +For the offline Python API, use `extra_args`: + +```python +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +sampling_params = OmniDiffusionSamplingParams( + num_inference_steps=20, + seed=42, + extra_args={"task": "t2va", "sampler": "res_multistep"}, +) +``` + +Both samplers use one DiT evaluation per step in request mode and step execution. +`res_multistep` retains the previous denoised prediction independently for each +request and each stream; its first and final updates use Euler. Omitting `sampler` +or selecting `euler` preserves the existing Euler update. Reducing the step count +reduces model evaluations, but quality depends on the prompt and configuration; +20 RES steps are a candidate to compare against an Euler-50 reference. + +Requests using a fixed distilled `base_schedule` from a checkpoint or adapter +must keep `sampler="euler"`. RES sampling with those schedules has not been +validated and is rejected with `ValueError`. + ## Request-scoped quality Add one of these fields to any HTTP request above. No Cache-DiT startup option @@ -1408,6 +1444,7 @@ four-step adapter keeps its original sampling and fixed-top-k behavior. | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count | | `fps` | `24` | H3 output FPS is fixed | | `num_inference_steps` | `50` | Matches the reference accuracy workloads | +| `extra_params.sampler` | `euler` | Per-request `euler` or `res_multistep` for both streams; fixed distilled schedules require Euler | | `flow_shift` | `12` | Video sigma shift | | `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` | | `seed` | Task-specific | Use a fixed value for reproducibility | diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 5bba75f93..fbfe0efe3 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -102,7 +102,11 @@ def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_v state.latents = video_rows.clone() state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32) state.step_index = 0 + from vllm_omni.diffusion.models.minimax_h3.sampling import create_h3_sample_solver + state.extra = { + mod._STEP_VIDEO_SOLVER: create_h3_sample_solver("euler", sigmas_video), + mod._STEP_AUDIO_SOLVER: create_h3_sample_solver("euler", sigmas_audio), mod._STEP_BRANCH: branch, # Co-batched requests must share one DiT instance, or denoise_step() # treats the batch as mixed-task and falls back to one forward each. diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index 1b358a070..c4c452b3c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -3,7 +3,7 @@ """MiniMax H3 cfg-distilled full denoise loop. Per step, the positive presentation is forwarded exactly once. Video and audio -target rows chain through the Euler-eta0 update while visual and audio condition +target rows chain through the selected solver while visual and audio condition rows stay pinned to their noised step-0 anchors. """ @@ -25,8 +25,8 @@ from vllm_omni.diffusion.forward_context import ( from vllm_omni.platforms import current_omni_platform from .latent_mask import MiniMaxH3LatentEdit, minimax_h3_prepare_edit_rows +from .sampling import create_h3_sample_solver from .scheduling_minimax_h3_euler_ancestral import ( - minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, ) @@ -333,6 +333,7 @@ def minimax_h3_denoise_loop( audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, step_profiler: Callable[[int], AbstractContextManager] | None = None, + sampler: str = "euler", ) -> tuple[torch.Tensor, torch.Tensor]: """Run the full denoise loop; returns final (video_rows, audio_rows). @@ -365,13 +366,15 @@ def minimax_h3_denoise_loop( if audio_edit is not None: audio_edit = audio_edit.to(device=device, dtype=torch.float32) + video_solver = create_h3_sample_solver(sampler, sigmas_video) + audio_solver = create_h3_sample_solver(sampler, sigmas_audio) num_steps = len(sigmas_video) - 1 for step in range(num_steps): check_request_cancellation() step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: - s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] - s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] + s_v = sigmas_video[step] + s_a = sigmas_audio[step] # Publish where we are so step-gated attention features (the dense # warmup of RAINFUSION_ATTN, the timestep gate of TRTLLM_ATTN) can # see it. Gates use the scheduler-style descending timestep, which @@ -425,7 +428,7 @@ def minimax_h3_denoise_loop( mv_video_t, t_v, ) - new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next) + new_target = video_solver.step(video_rows[update], x0_video, step) video_rows = video_rows.clone() video_rows[update] = new_target if cond_anchor is not None: @@ -443,9 +446,7 @@ def minimax_h3_denoise_loop( mv_audio_t, t_a, ) - new_audio = minimax_h3_euler_eta0_step( - audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next - ) + new_audio = audio_solver.step(audio_rows[audio_update], x0_audio, step) audio_rows = audio_rows.clone() audio_rows[audio_update] = new_audio if positive.locked_audio_rows is None else positive.locked_audio_rows if audio_anchor is not None: diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 563cce959..b1dfc78b9 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -157,8 +157,8 @@ from .packed_tokens import ( minimax_h3_unpatchify_video_tokens, ) from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy +from .sampling import create_h3_sample_solver, normalize_h3_sampler from .scheduling_minimax_h3_euler_ancestral import ( - minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, ) from .time_request import ( @@ -302,6 +302,8 @@ _REFINE_KEYFRAME_CONDITION = "minimax_h3_refine_keyframe_condition" _STEP_BRANCH = "minimax_h3_branch" _STEP_AUDIO_ROWS = "minimax_h3_audio_rows" _STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred" +_STEP_VIDEO_SOLVER = "minimax_h3_video_solver" +_STEP_AUDIO_SOLVER = "minimax_h3_audio_solver" _STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video" _STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio" _STEP_COND_ANCHOR = "minimax_h3_cond_anchor" @@ -1780,6 +1782,7 @@ class MiniMaxH3Pipeline( audio_edit_clean_rows: torch.Tensor | None = None, audio_edit_mask_rows: torch.Tensor | None = None, audio_edit_restore_mask_rows: torch.Tensor | None = None, + sampler: str = "euler", init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, refine: MiniMaxH3LatentRefineSpec | None = None, ) -> dict[str, Any]: @@ -1792,6 +1795,12 @@ class MiniMaxH3Pipeline( one: the rows start from those latents re-noised to the schedule position ``refine`` selects, and the returned schedules begin there. """ + sampler = normalize_h3_sampler(sampler) + if sampler != "euler" and base_schedule is not None: + raise ValueError( + "MiniMax H3 res_multistep sampling with a fixed distilled sigma schedule has not been validated; " + "use sampler='euler' for this schedule." + ) video_sigmas = minimax_h3_time_shift_sigmas( num_steps=num_steps, shift_scale=video_shift, @@ -1977,6 +1986,7 @@ class MiniMaxH3Pipeline( "audio_anchor": ( None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32) ), + "sampler": sampler, "sigmas_video": video_sigmas, "sigmas_audio": audio_sigmas, "video_edit": video_edit, @@ -2048,10 +2058,12 @@ class MiniMaxH3Pipeline( audio_edit_clean_rows: torch.Tensor | None = None, audio_edit_mask_rows: torch.Tensor | None = None, audio_edit_restore_mask_rows: torch.Tensor | None = None, + sampler: str = "euler", init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, refine: MiniMaxH3LatentRefineSpec | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = self._build_denoise_inputs( + sampler=sampler, task=task, text_embeddings=text_embeddings, text_tags=text_tags, @@ -2099,6 +2111,7 @@ class MiniMaxH3Pipeline( with self._resident_dit_layers_on_device(enabled=transformer is self.transformer): with self.progress_bar(total=len(inputs["sigmas_video"]) - 1) as progress: video_rows, audio_rows = minimax_h3_denoise_loop( + sampler=inputs["sampler"], model=transformer, positive=branch, initial_video_rows=inputs["video_rows"], @@ -2862,6 +2875,12 @@ class MiniMaxH3Pipeline( self._prepare_adaln_adapter(sampling) base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) + sampler = normalize_h3_sampler(extra.get("sampler")) + if sampler != "euler" and base_schedule is not None: + raise ValueError( + "MiniMax H3 res_multistep sampling with a fixed distilled sigma schedule has not been validated; " + "use sampler='euler' for this schedule." + ) transformer = getattr( self, "transformers_ref" if task == "ref2va" and hasattr(self, "transformers_ref") else "transformer", None ) @@ -2939,6 +2958,7 @@ class MiniMaxH3Pipeline( "keyframe_frame_indices": list(conditioning.keyframe_frame_indices) or None, "pad_seq_len": _resolve_pad_seq_len(extra.get("pad_seq_len")), "seed": int(sampling.seed if sampling.seed is not None else 42), + "sampler": sampler, "num_steps": num_steps, "video_shift": float(extra.get("flow_shift", self.default_video_shift)), "audio_shift": float(extra.get("audio_flow_shift", self.default_audio_shift)), @@ -2961,7 +2981,10 @@ class MiniMaxH3Pipeline( @staticmethod def _denoise_kwargs(context: dict[str, Any]) -> dict[str, Any]: """Select the denoise-input arguments from a prepared request context.""" - return {key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS} + return { + **{key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS}, + "sampler": context.get("sampler", "euler"), + } @torch.no_grad() def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: @@ -3165,6 +3188,8 @@ class MiniMaxH3Pipeline( _STEP_AUDIO_ROWS: audio_rows, _STEP_COND_ANCHOR: cond_anchor, _STEP_AUDIO_ANCHOR: audio_anchor, + _STEP_VIDEO_SOLVER: create_h3_sample_solver(inputs.get("sampler", "euler"), sigmas_video), + _STEP_AUDIO_SOLVER: create_h3_sample_solver(inputs.get("sampler", "euler"), sigmas_audio), _STEP_SIGMAS_VIDEO: sigmas_video, _STEP_SIGMAS_AUDIO: sigmas_audio, _STEP_VIDEO_EDIT: inputs.get("video_edit"), @@ -3318,7 +3343,7 @@ class MiniMaxH3Pipeline( return video_velocity def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: - """Apply one Euler-eta0 update to this request's video and audio rows.""" + """Apply one request-local solver update to the video and audio rows.""" del kwargs # denoise_step() stages the audio half of this step's velocity; popping # it keeps a second step_scheduler() call from reusing a stale one. @@ -3347,12 +3372,7 @@ class MiniMaxH3Pipeline( noise_pred.float()[update], schedule["t_video"], ) - new_video = minimax_h3_euler_eta0_step( - video_rows[update], - x0_video, - sigma_curr=schedule["sigma_video"], - sigma_next=schedule["sigma_video_next"], - ) + new_video = state.extra[_STEP_VIDEO_SOLVER].step(video_rows[update], x0_video, state.step_index) video_rows = video_rows.clone() video_rows[update] = new_video if cond_anchor is not None: @@ -3371,12 +3391,7 @@ class MiniMaxH3Pipeline( audio_noise_pred.float()[audio_update], schedule["t_audio"], ) - new_audio = minimax_h3_euler_eta0_step( - audio_rows[audio_update], - x0_audio, - sigma_curr=schedule["sigma_audio"], - sigma_next=schedule["sigma_audio_next"], - ) + new_audio = state.extra[_STEP_AUDIO_SOLVER].step(audio_rows[audio_update], x0_audio, state.step_index) audio_rows = audio_rows.clone() audio_rows[audio_update] = new_audio if branch.locked_audio_rows is None else branch.locked_audio_rows if audio_anchor is not None: diff --git a/vllm_omni/diffusion/models/minimax_h3/sampling.py b/vllm_omni/diffusion/models/minimax_h3/sampling.py new file mode 100644 index 000000000..2db9c5de4 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/sampling.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request-local H3 solvers reconstructed from the surviving experiment records. + +RES coefficients follow ComfyUI_RH_MinMaxH3 sampler_core.py at +6abdf97ae8cdce548c751ce850adc5010b583b7c. Euler delegates to the existing +implementation to preserve its exact floating-point operation order. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence + +import torch + +from .scheduling_minimax_h3_euler_ancestral import minimax_h3_euler_eta0_step + + +def normalize_h3_sampler(name: str | None) -> str: + if name is None: + return "euler" + if name not in ("euler", "res_multistep"): + raise ValueError(f"Unsupported MiniMax H3 sampler: {name!r}; expected euler or res_multistep") + return name + + +def res_multistep_coeffs(sigmas: Sequence[float]) -> list[tuple[float, float] | None]: + values = [float(s) for s in sigmas] + if len(values) < 2 or any(not math.isfinite(s) or not 0 <= s <= 1 for s in values): + raise ValueError("H3 sigmas must contain at least two finite values in [0, 1]") + if any(a <= b for a, b in zip(values, values[1:])): + raise ValueError("H3 sigmas must be strictly decreasing") + coeffs = [] + for i, (curr, nxt) in enumerate(zip(values, values[1:])): + if i == 0 or nxt == 0: + coeffs.append(None) + continue + h = math.log(curr / nxt) + c2 = math.log(curr / values[i - 1]) / h + # Stable evaluation near h=0, where subtracting phi1-1 loses precision. + phi1 = math.expm1(-h) / (-h) + phi2 = (phi1 - 1) / (-h) if abs(h) >= 1e-5 else 0.5 - h / 6 + h * h / 24 - h**3 / 120 + coeffs.append((h * (phi1 - phi2 / c2), h * phi2 / c2)) + return coeffs + + +class H3SampleSolver: + """Own the denoised history for one request and one video/audio stream.""" + + def __init__(self, name: str | None, sigmas: Sequence[float]): + self.name = normalize_h3_sampler(name) + self.sigmas = tuple(float(s) for s in sigmas) + self.coeffs = res_multistep_coeffs(self.sigmas) + self.old_denoised: torch.Tensor | None = None + self.step_index = 0 + + def step(self, state: torch.Tensor, denoised: torch.Tensor, step: int) -> torch.Tensor: + if step != self.step_index or not 0 <= step < len(self.coeffs): + raise ValueError("H3 solver steps must be consumed once, in order") + coeff = self.coeffs[step] + curr, nxt = self.sigmas[step : step + 2] + if self.name == "euler" or coeff is None: + out = minimax_h3_euler_eta0_step(state, denoised, sigma_curr=curr, sigma_next=nxt) + else: + old = self.old_denoised + if old is None or state.shape != denoised.shape or old.shape != state.shape: + raise ValueError("H3 solver state and denoised history shapes must match") + if not torch.is_floating_point(state) or not torch.is_floating_point(denoised): + raise ValueError("H3 solver requires floating point tensors") + if not torch.isfinite(state).all() or not torch.isfinite(denoised).all(): + raise ValueError("H3 solver inputs must be finite") + dtype = torch.float32 if state.dtype in (torch.float16, torch.bfloat16) else state.dtype + ratio = state.new_tensor(nxt / curr, dtype=dtype) + hb1, hb2 = coeff + out = (ratio * state.to(dtype) + hb1 * denoised.to(dtype) + hb2 * old.to(dtype)).to(state.dtype) + if not torch.isfinite(out).all(): + raise ValueError("H3 solver output must be finite") + if self.name == "res_multistep": + self.old_denoised = denoised.detach().clone() if step + 1 < len(self.coeffs) else None + self.step_index += 1 + return out + + +def create_h3_sample_solver(name: str | None, sigmas: Sequence[float]) -> H3SampleSolver: + return H3SampleSolver(name, sigmas)