diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 02d51197b..9594e62e5 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -422,6 +422,7 @@ source_file_dependencies: diffusion_minimax_h3_function: - *minimax_h3 + - tests/helpers/assertions.py - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py - tests/e2e/online_serving/minimax_h3/ - tests/diffusion/models/minimax_h3/ diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index a14a57c6c..104d6df26 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -208,3 +208,14 @@ steps: timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' mirror_hardwares: h100_2 + + - label: "Diffusion · MiniMax-H3 Latent Upscale and Refine" + source_file_dependencies: diffusion_minimax_h3_function + timeout_in_minutes: 40 + commands: + - | + set -eu + VLLM_TEST_MINIMAX_H3_LATENT_UPSCALER=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="LBH-123-AI/Minimax_h3_latent_Upscaler", filename="minimax_h3_latent_upscaler_3d_conv_v1/minimax_h3_latent_upscaler_3d_conv_v1_bf16.safetensors"))') + export VLLM_TEST_MINIMAX_H3_LATENT_UPSCALER + timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' + mirror_hardwares: h100_2 diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 29fbf9050..5bfd65649 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -1515,9 +1515,184 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ --cache-config '{"rel_l1_thresh":0.17}' ``` +## Latent super-resolution + +H3's video VAE is a ~5B-parameter model, so the usual way to raise output +resolution -- decode, upscale in pixel space, re-encode -- pays for that VAE +twice. The community +[Minimax H3 latent upscaler](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler) +is a learned 3D-convolutional resizer trained on H3's 24-channel latents: it +runs between the denoise loop and the VAE, so the clip is decoded once, at the +larger size. Unlike bilinear or bicubic interpolation of a latent, it does not +introduce ghosting or double-image artifacts. + +The network runs one normalization below the pipeline latent: an H3 latent in +vLLM-Omni is already normalized, and the upscaler was trained on that tensor +normalized again by the VAE's `latents_mean` / `latents_std`. The stage applies +and undoes that itself, so nothing about a request changes -- but a port that +skips it inflates the latent about 5x and decodes as a magenta grid at one tile +per latent cell. + +The checkpoint is optional and is not part of the H3 release. Download one file +from that repository (`minimax_h3_latent_upscaler_3d_bf16.safetensors`, 691 MB) +and name it at startup: + +```bash +vllm serve "${MODEL_ROOT}/FL2VA" \ + --omni \ + --trust-remote-code \ + --additional-config '{"latent_upscaler_path":"/models/minimax_h3_latent_upscaler_3d_bf16.safetensors"}' +``` + +A request then asks for a size. The three modes are mutually exclusive and all +resolve to the same 16-pixel latent grid, so the three requests below are the +same 960x544 -> 1920x1088 upscale: + +```bash +curl -sS -X POST "${API_URL}" \ + -F 'prompt=...' \ + -F 'width=960' -F 'height=544' -F 'fps=24' \ + -F 'extra_params={"task":"t2va","duration":8.7,"latent_upscale":2.0}' \ + -o t2va_1080p.mp4 + +# or name the output size in pixels +# "latent_upscale":{"width":1920,"height":1088} +# or name a pixel budget and keep the aspect ratio +# "latent_upscale":{"megapixels":2.0} +``` + +Over HTTP the field is `extra_params`: the videos request model declares that +name and nothing else, so an `extra_args` form field is dropped without an +error and the request silently returns the un-upscaled size. Offline, the same +value goes in `extra_args`: + +```python +omni = Omni( + model=os.path.join(os.environ["MODEL_ROOT"], "FL2VA"), + trust_remote_code=True, + additional_config={"latent_upscaler_path": "/models/minimax_h3_latent_upscaler_3d_bf16.safetensors"}, +) +outputs = omni.generate( + "A quiet cinematic night scene with matching ambient sound.", + OmniDiffusionSamplingParams( + height=544, + width=960, + fps=24, + num_inference_steps=50, + seed=42, + extra_args={"task": "t2va", "duration": 8.7, "latent_upscale": 2.0}, + ), +) +``` + +### The hi-res route: generate small, refine large + +Upscaling alone decodes the enlarged latent directly, which is fast but can +only interpolate detail the first pass never generated. Adding `latent_refine` +resumes the denoise schedule at the new size, so the DiT puts real detail back: + +1. generate at a cheap size -- far fewer DiT tokens, +2. upscale the latent, +3. re-noise it and run the tail of the schedule at the target size. + +```bash +curl -sS -X POST "${API_URL}" \ + -F 'prompt=...' \ + -F 'width=960' -F 'height=544' -F 'fps=24' \ + -F 'num_inference_steps=50' \ + -F 'extra_params={"task":"t2va","duration":8.7,"latent_upscale":2.0,"latent_refine":0.4}' \ + -o t2va_1080p.mp4 +``` + +`latent_refine` is the fraction of the request's steps the second pass re-runs, +the img2img convention, so it sets both how much the pass may change and what +it costs: `0.4` of `num_inference_steps=50` is 50 cheap steps followed by 20 +expensive ones. Low values stay close to the first pass and mostly sharpen; +high values re-generate and may drift from it. `1.0` re-runs the whole schedule +from noise, which discards the first pass entirely. + +Detail does not keep climbing with strength. Measured on a 15s 2688x1536 clip, +six frames apart, against the same seed and prompt: 0.5 returned 0.98x the +high-frequency energy of 0.35 and 0.65 returned 0.92x, while the mean per-pixel +drift from the 0.35 result rose from 6.3 to 12.6 (0-255). What caps detail is +what the DiT can resolve at the target size, not how noisy a start it is given; +past that, a higher strength only erases more of the first pass and re-imagines +it from less. On this checkpoint 0.35 was the best of the three -- treat higher +values as "re-generate", not as "sharpen". + +Why this is faster than generating at the target size: the DiT cost follows the +video token count, which is quadratic in the latent area for attention. The +example above is 31,620 video tokens at 960x544 against 126,480 at 1920x1088 -- +4x the tokens, so roughly 4x the linear cost and up to 16x the attention cost +per step. Paying that on 20 steps instead of 50 is where the saving comes from. +The gain therefore depends on the resolutions, the strength, and how +attention-bound the shape is; measure it on your own workload. + +Audio is re-noised and refined alongside the video at the same schedule +position. The two modalities are shifted apart (`flow_shift` 12.0 against +`audio_flow_shift` 3.0), so each resumes at its own sigma for that position -- +the state the first pass actually passed through, which is what keeps the pass +in distribution. + +`latent_refine` works without `latent_upscale` too, as a plain detail pass at +the generated size. For FL2VA, the keyframe condition is pinned to the output +latent size, so a refine after an upscale re-encodes the keyframes from the +originals at the new size; REF2VA references carry their own shapes and are +reused as they are. + +### Options + +| `--additional-config` key | Default | Meaning | +| --- | --- | --- | +| `latent_upscaler_path` | unset | Checkpoint file, or a directory holding exactly one. Unset disables the stage and loads no extra weights. | +| `latent_upscaler_dtype` | the engine dtype | `bf16`, `fp16` or `fp32`. | +| `latent_upscaler_chunk_frames` | `0` | `0` runs the whole clip in one pass; a positive value bounds activation memory with approximate temporal chunking. | +| `latent_upscaler_resident` | `false` | Keep the ~700 MB of weights on the GPU between requests instead of in host memory. | +| `latent_refine_max_tokens_per_rank` | `65536` | Reject larger refine layouts before denoising. Set `0` only after validating a larger deployment. | +| `latent_upscale` | unset | A default target applied to every request, in the request's own format. A request overrides it, and `false` opts out. | +| `latent_refine` | unset | A default refine strength applied to every request. A request overrides it, and `false` opts out. | + +The `latent_upscale` request value is a multiplier (`2.0`), `false`, or an +object naming one mode: `{"scale": 2.0}`, `{"width": …, "height": …}` or +`{"megapixels": …}`, each accepting an `align` (default `32` pixels). The +checkpoint was trained between 1.0x and 4.0x and only upscales; a smaller +target is rejected, and `scale` 1.0 is a no-op that skips the stage. The +`latent_refine` value is a strength in `(0, 1]`, `false`, or +`{"strength": 0.4}`; it needs no checkpoint of its own. + +When `latent_refine` is enabled, both resolved output dimensions must be +divisible by 32 pixels for the DiT's spatial patches. A custom `align` that +produces an incompatible target is rejected before denoising; upscale-only +requests can still use the VAE's 16-pixel grid. Numeric upscale fields must be +positive and finite, and `width`, `height`, and `align` must be integers. + +Temporal chunking can bound activation memory on long clips -- a 15-second clip is +102 latent frames, more than three 32-frame chunks. It is an approximation rather than a +partition: the network's GroupNorms pool statistics over whatever clip they are +handed, so a chunked pass can differ from a single pass across the whole output +and not only at the chunk seams. How much that costs in practice was not +measured here; the default single pass buys the exact result for the +memory a full-length activation volume takes (10.1 GiB against 5.6 GiB at +2688x1536 for 15 seconds, measured on one B300). Set +`latent_upscaler_chunk_frames: 32` explicitly if that memory cost is too high. + ## Known limitations - TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached. +- The latent upscaler is a community checkpoint, not part of the H3 release, + and is unavailable unless `latent_upscaler_path` names one. It runs + replicated on every DiT rank rather than sharded. +- `latent_refine` is request-mode only. It is a second denoise loop with its + own resolution and its own truncated schedule, and the step contract carries + one schedule per request, so step execution rejects it. `latent_upscale` + without `latent_refine` works in both modes. +- Large refine layouts can abort workers with `CUDA error: an illegal memory + access was encountered`. On B300, observed video token counts of 54,144 per + rank complete while 108,288 and 109,360 abort. The preflight limit of 65,536 + counts the full padded layout, including text, audio, and visual references, + per Ulysses rank. It rejects the observed failing sizes before the base pass. + This is an empirical B300 guard, not a proven kernel limit on every device; + deployments with different hardware can adjust the limit after validation. - Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads both task-specific DiTs, and loads shared components once from `FL2VA`. - Request mode executes one generation request per diffusion batch. Use diff --git a/tests/diffusion/cache/test_cache_dit_request_runtime.py b/tests/diffusion/cache/test_cache_dit_request_runtime.py index 8bc6ee846..7c638dfd4 100644 --- a/tests/diffusion/cache/test_cache_dit_request_runtime.py +++ b/tests/diffusion/cache/test_cache_dit_request_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 import pytest @@ -45,7 +45,7 @@ def _recording_backend(events): def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch): - events = [] + events: list[tuple] = [] FakeBackend = _recording_backend(events) monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend) pipeline = object() @@ -83,7 +83,7 @@ def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch): - events = [] + events: list[tuple] = [] monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events)) pipeline = object() runtime = RequestScopedCacheDiTRuntime(pipeline) @@ -99,3 +99,21 @@ def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch ] assert events[2][-1] == 40 assert not runtime.is_enabled + + +def test_denoise_pass_refresh_preserves_profile_and_uses_actual_steps(monkeypatch): + events: list[tuple] = [] + monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events)) + pipeline = object() + runtime = RequestScopedCacheDiTRuntime(pipeline) + runtime.refresh(10) + assert events == [] + runtime.prepare(_spec("high", 50)) + events.clear() + runtime.refresh(20) + runtime.refresh(50) + assert events == [("refresh", pipeline, 20), ("refresh", pipeline, 50)] + runtime.disable() + events.clear() + runtime.refresh(20) + assert events == [] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 0c2af16fb..4405997f7 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -3551,7 +3551,7 @@ def test_request_cancellation_at_prepare_and_decode_boundaries(preencode, cancel def prepare(*args): if cancel_phase == "prepare": registry.cancel(["request"]) - return {"num_outputs": 1, "seed": 1101, "preencode_mp4": preencode} + return {"num_outputs": 1, "seed": 1101, "preencode_mp4": preencode, "height": 2, "width": 2} def diffuse(**kwargs): registry.cancel(["request"]) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3.py new file mode 100644 index 000000000..a782fc181 --- /dev/null +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Real-weight L3 coverage for optional latent upscaling and second-pass refinement. + +Set VLLM_TEST_MINIMAX_H3_LATENT_UPSCALER to the LBH BF16 upscaler checkpoint. +VLLM_TEST_MINIMAX_H3_MODEL may point to a local FL2VA model directory. +""" + +import json +import os +from pathlib import Path + +import pytest + +from tests.helpers.assertions import assert_video_first_frames_differ, assert_video_valid +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient + +MODEL = os.environ.get("VLLM_TEST_MINIMAX_H3_MODEL", "MiniMaxAI/MiniMax-H3") +UPSCALER = os.environ.get("VLLM_TEST_MINIMAX_H3_LATENT_UPSCALER") +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] + + +@pytest.mark.skipif(not UPSCALER, reason="set VLLM_TEST_MINIMAX_H3_LATENT_UPSCALER to the real upscaler checkpoint") +@pytest.mark.parametrize( + "omni_server", + [ + pytest.param( + OmniServerParams( + model=MODEL, + server_args=[ + "--trust-remote-code", + "--task-type", + "fl2va", + "--num-gpus", + "2", + "--tensor-parallel-size", + "2", + "--usp", + "1", + "--ring", + "1", + "--text-encoder-tp-size", + "2", + "--vae-patch-parallel-size", + "2", + "--vae-parallel-mode", + "tile", + "--vae-use-tiling", + "--enable-cpu-offload", + "--enforce-eager", + "--diffusion-attention-backend", + "FLASH_ATTN", + "--additional-config", + json.dumps({"latent_upscaler_path": UPSCALER, "latent_upscaler_dtype": "bf16"}), + ], + env_dict={"VLLM_WORKER_MULTIPROC_METHOD": "spawn"}, + ), + id="minimax_h3_latent_upscale_refine_tp2", + marks=hardware_marks(res={"cuda": "H100"}, num_cards=2), + ), + ], + indirect=True, +) +@pytest.mark.parametrize("task", ["t2va", "fl2va"]) +def test_latent_upscale_and_refine( + omni_server: OmniServer, + online_client: OnlineOmniClient, + tmp_path: Path, + task: str, +) -> None: + """Exercise upscale, refinement, and request-local opt-out on one live server.""" + videos = [] + frame_counts = [] + image_reference = None + if task == "fl2va": + image = generate_synthetic_image(512, 288, seed=42) + image_reference = f"data:image/jpeg;base64,{image['base64']}" + # The 1024x576 refine keyframe spans multiple 256px VAE encoder tiles: + # both VAE ranks must encode it before broadcasting the condition latents. + # Four base steps and the minimum four-second clip keep real-weight CI affordable. + # Use size rather than width/height: upscale intentionally changes output dimensions. + for name, upscale, refine, width, height in [ + ("original", False, False, 512, 288), + ("upscaled", 2.0, False, 1024, 576), + ("refined", 2.0, 0.5, 1024, 576), + ("disabled_again", False, False, 512, 288), + ]: + responses = online_client.send_video_diffusion_request( + { + "model": omni_server.model, + "image_reference": image_reference, + "form_data": { + "model": omni_server.model, + "prompt": "A woman in a sunlit room turns toward the camera, cinematic lighting.", + "size": "512x288", + "fps": 24, + "num_inference_steps": 4, + "flow_shift": 12, + "seed": 1101, + "extra_params": json.dumps( + { + "task": task, + "duration": 4.0, + "aspect_ratio": "16:9", + "audio_flow_shift": 3.0, + "latent_upscale": upscale, + "latent_refine": refine, + } + ), + }, + "expected_audio": {"sample_rate": 32000, "channels": 2}, + } + ) + video = responses[0].videos[0] + (tmp_path / f"{name}.mp4").write_bytes(video) + metadata = assert_video_valid(video, width=width, height=height, fps=24) + frame_counts.append(metadata["num_frames"]) + videos.append(video) + + assert frame_counts[0] > 0 + assert len(set(frame_counts)) == 1, f"Upscale/refinement changed the frame count: {frame_counts}" + # Equal output resolution alone would miss a silently ignored second denoise pass. + assert_video_first_frames_differ(videos[2], videos[1], min_mean_absolute_error=1 / 255) diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index d877c9068..f6559e0c3 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -367,6 +367,25 @@ def assert_video_first_frame_matches( ) +def assert_video_first_frames_differ( + video: bytes, + reference: bytes, + *, + min_mean_absolute_error: float, +) -> None: + """Compare decoded pixels, ignoring differences in MP4 encoding metadata.""" + with av.open(BytesIO(video)) as container: + actual = next(container.decode(video=0)).to_ndarray(format="rgb24") + with av.open(BytesIO(reference)) as container: + expected = next(container.decode(video=0)).to_ndarray(format="rgb24") + + assert actual.shape == expected.shape + mean_absolute_error = float(np.abs(actual.astype(np.float32) - expected).mean() / 255.0) + assert mean_absolute_error > min_mean_absolute_error, ( + f"Expected first-frame MAE > {min_mean_absolute_error}, got {mean_absolute_error:.6f}." + ) + + def assert_audio_diffusion_response( response: "DiffusionResponse", request_config: dict[str, Any], diff --git a/vllm_omni/diffusion/cache/cachedit/runtime.py b/vllm_omni/diffusion/cache/cachedit/runtime.py index e3403434a..4e76590cc 100644 --- a/vllm_omni/diffusion/cache/cachedit/runtime.py +++ b/vllm_omni/diffusion/cache/cachedit/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 """Declarative request-scoped Cache-DiT lifecycle.""" @@ -87,6 +87,12 @@ class RequestScopedCacheDiTRuntime: self._backend = None self._installation_key = None + def refresh(self, num_inference_steps: int) -> None: + """Start an independent denoise pass with the installed profile.""" + if self.is_enabled: + assert self._backend is not None + self._backend.refresh(self._pipeline, num_inference_steps=num_inference_steps) + __all__ = [ "CacheDiTRequestSpec", diff --git a/vllm_omni/diffusion/models/minimax_h3/latent_upscaler.py b/vllm_omni/diffusion/models/minimax_h3/latent_upscaler.py new file mode 100644 index 000000000..419dd243a --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/latent_upscaler.py @@ -0,0 +1,672 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Learned latent super-resolution for MiniMax-H3 video latents. + +The released upscaler is a pure-3D convolutional resizer trained on H3's +24-channel VAE latents: it lifts a latent to a target spatial size without a +decode/encode round trip through the ~5B-parameter H3 VAE. Weights come from +``LBH-123-AI/Minimax_h3_latent_Upscaler`` and are a plain ``state_dict``, so the +module layout below reproduces the checkpoint's parameter names exactly. + +The network works in a space one normalization *below* the pipeline latent. +A vLLM-Omni H3 latent is already normalized -- :class:`MiniMaxH3VideoVAE` +denormalizes at decode time -- and the reference ComfyUI node applies the VAE's +``latents_mean`` / ``latents_std`` on top of a latent in that same convention, +so the upscaler was trained on the doubly-normalized tensor. Feeding it the +pipeline latent directly inflates the output by roughly 5x (measured: output +std 5.4 against an input std of 1.0), which decodes as a magenta grid at one +tile per latent cell. :meth:`MiniMaxH3LatentUpscaler.upscale` therefore +normalizes in and denormalizes out, exactly as the reference node does. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path + +import regex as re +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# H3's video VAE downsamples 16x spatially, so one latent cell is a 16px tile. +MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE = 16 +# The published checkpoints were trained between 1x and 4x. +MINIMAX_H3_LATENT_UPSCALE_MAX_SCALE = 4.0 +_EMBED_DIM = 64 +_GROUP_NORM_GROUPS = 32 + + +class MiniMaxH3LatentUpscalerError(ValueError): + """A latent upscaler artifact that cannot be served.""" + + +@dataclass(frozen=True) +class MiniMaxH3LatentUpscalerArch: + """Shape of a released latent upscaler, recovered from its ``state_dict``.""" + + in_channels: int = 24 + channels: int = 512 + in_blocks: int = 12 + out_blocks: int = 12 + temporal_every: int = 2 + temporal_kernel: int = 5 + + +def _normalization(channels: int) -> nn.GroupNorm: + return nn.GroupNorm(_GROUP_NORM_GROUPS, channels) + + +class _ResBlockEmb3D(nn.Module): + """3D residual block whose output norm is modulated by the scale embedding.""" + + def __init__(self, channels: int, emb_channels: int) -> None: + super().__init__() + self.in_layers = nn.Sequential( + _normalization(channels), + nn.SiLU(), + nn.Conv3d(channels, channels, 3, padding=1), + ) + self.emb_layers = nn.Sequential( + nn.SiLU(), + nn.Linear(emb_channels, 2 * channels), + ) + self.out_norm = _normalization(channels) + # Index 1 is the training-time dropout. It is a no-op at inference but + # has to keep its slot so ``out_layers.2`` still names the convolution. + self.out_layers = nn.Sequential( + nn.SiLU(), + nn.Identity(), + nn.Conv3d(channels, channels, 3, padding=1), + ) + self.skip = nn.Identity() + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + h = self.in_layers(x) + scale, shift = torch.chunk(self.emb_layers(emb).to(h.dtype)[..., None, None, None], 2, dim=1) + h = self.out_norm(h) * (1 + scale) + shift + return self.skip(x) + self.out_layers(h) + + +class _TemporalConv3D(nn.Module): + """Zero-initialized depthwise temporal residual, for cross-frame coherence.""" + + def __init__(self, channels: int, kernel_size: int) -> None: + super().__init__() + self.norm = _normalization(channels) + self.dwconv = nn.Conv3d( + channels, + channels, + kernel_size=(kernel_size, 1, 1), + padding=(kernel_size // 2, 0, 0), + groups=channels, + ) + self.pwconv = nn.Conv3d(channels, channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.pwconv(self.dwconv(F.silu(self.norm(x)))) + + +class MiniMaxH3LatentResizer3D(nn.Module): + """The released pure-3D latent resizer. + + ``in_blocks`` run at the source resolution, a trilinear interpolation moves + the feature volume to the target size, and ``out_blocks`` refine it there. + The requested scale enters every residual block through a two-layer + embedding of ``scale - 1``. + """ + + def __init__(self, arch: MiniMaxH3LatentUpscalerArch) -> None: + super().__init__() + self.arch = arch + self.conv_in = nn.Conv3d(arch.in_channels, arch.channels, 3, padding=1) + self.embed = nn.Sequential( + nn.Linear(1, _EMBED_DIM), + nn.SiLU(), + nn.Linear(_EMBED_DIM, _EMBED_DIM), + ) + self.in_blocks = self._build_stage(arch, arch.in_blocks) + self.out_blocks = self._build_stage(arch, arch.out_blocks) + self.norm_out = _normalization(arch.channels) + self.conv_out = nn.Conv3d(arch.channels, arch.in_channels, 3, padding=1) + + @staticmethod + def _build_stage(arch: MiniMaxH3LatentUpscalerArch, num_blocks: int) -> nn.ModuleList: + # Residual and temporal blocks share one flat ModuleList in the + # checkpoint, so their interleaving fixes the parameter names. + stage = nn.ModuleList() + for index in range(num_blocks): + stage.append(_ResBlockEmb3D(arch.channels, _EMBED_DIM)) + if arch.temporal_every > 0 and index % arch.temporal_every == 0: + stage.append(_TemporalConv3D(arch.channels, arch.temporal_kernel)) + return stage + + @property + def temporal_kernel(self) -> int: + return self.arch.temporal_kernel if self.arch.temporal_every > 0 else 0 + + def forward(self, latent: torch.Tensor, *, scale: float, target_size: tuple[int, int, int]) -> torch.Tensor: + """Resize ``[B, C, T, H, W]`` to ``target_size`` under a scale hint.""" + emb = self.embed( + torch.tensor([[scale - 1.0]], dtype=latent.dtype, device=latent.device), + ).expand(latent.shape[0], -1) + + hidden = self.conv_in(latent) + for block in self.in_blocks: + hidden = block(hidden, emb) if isinstance(block, _ResBlockEmb3D) else block(hidden) + hidden = F.interpolate(hidden, size=target_size, mode="trilinear", align_corners=False) + for block in self.out_blocks: + hidden = block(hidden, emb) if isinstance(block, _ResBlockEmb3D) else block(hidden) + return self.conv_out(F.silu(self.norm_out(hidden))) + + +def detect_minimax_h3_upscaler_arch(state_dict: Mapping[str, torch.Tensor]) -> MiniMaxH3LatentUpscalerArch: + """Recover the block layout of a released upscaler from its tensors.""" + conv_in = state_dict.get("conv_in.weight") + if conv_in is None or conv_in.ndim != 5: + raise MiniMaxH3LatentUpscalerError("not a MiniMax-H3 3D latent upscaler: expected a rank-5 'conv_in.weight'") + if any("attn" in key for key in state_dict): + raise MiniMaxH3LatentUpscalerError( + "this checkpoint carries attention blocks, which no released MiniMax-H3 upscaler uses" + ) + + residual_blocks: dict[str, set[int]] = {"in_blocks": set(), "out_blocks": set()} + temporal_blocks: dict[str, set[int]] = {"in_blocks": set(), "out_blocks": set()} + temporal_kernel = 5 + for key, tensor in state_dict.items(): + match = re.match(r"(in_blocks|out_blocks)\.(\d+)\.(in_layers|dwconv)\b", key) + if match is None: + continue + stage, index, kind = match.group(1), int(match.group(2)), match.group(3) + if kind == "in_layers": + residual_blocks[stage].add(index) + else: + temporal_blocks[stage].add(index) + if key.endswith("dwconv.weight"): + temporal_kernel = int(tensor.shape[2]) + + has_temporal = bool(temporal_blocks["in_blocks"] or temporal_blocks["out_blocks"]) + return MiniMaxH3LatentUpscalerArch( + in_channels=int(conv_in.shape[1]), + channels=int(conv_in.shape[0]), + in_blocks=len(residual_blocks["in_blocks"]), + out_blocks=len(residual_blocks["out_blocks"]), + temporal_every=2 if has_temporal else 0, + temporal_kernel=temporal_kernel, + ) + + +@dataclass(frozen=True) +class MiniMaxH3LatentUpscaleTarget: + """A resolved upscale target in latent cells, plus the scale hint to feed the network.""" + + latent_height: int + latent_width: int + scale: float + + @property + def height(self) -> int: + return self.latent_height * MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE + + @property + def width(self) -> int: + return self.latent_width * MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE + + +def resolve_minimax_h3_latent_upscale_target( + *, + latent_height: int, + latent_width: int, + scale: float | None = None, + height: int | None = None, + width: int | None = None, + megapixels: float | None = None, + align: int = 32, +) -> MiniMaxH3LatentUpscaleTarget: + """Resolve one of the three sizing modes into latent dimensions. + + ``scale`` multiplies the source size; ``height``/``width`` name the target + in pixels; ``megapixels`` names a total pixel budget and keeps the aspect + ratio. For the two size-based modes the network still needs a single scale + hint, which is the mean of the per-axis ratios -- what the reference node + feeds the same weights. Targets snap to ``align`` pixels and then to the + VAE's 16px cell. + """ + source_height = latent_height * MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE + source_width = latent_width * MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE + modes = [ + name + for name, value in (("scale", scale), ("dimensions", height or width), ("megapixels", megapixels)) + if value is not None + ] + if len(modes) != 1: + raise MiniMaxH3LatentUpscalerError( + f"latent upscaling takes exactly one of scale, height/width, or megapixels, got {modes or ['nothing']}" + ) + if align < 1: + raise MiniMaxH3LatentUpscalerError(f"latent upscale align must be positive, got {align}") + + if scale is not None: + target_height = source_height * float(scale) + target_width = source_width * float(scale) + elif megapixels is not None: + if megapixels <= 0: + raise MiniMaxH3LatentUpscalerError(f"latent upscale megapixels must be positive, got {megapixels}") + aspect = source_width / source_height + target_height = math.sqrt(float(megapixels) * 1024 * 1024 / aspect) + target_width = target_height * aspect + else: + if height is None or width is None: + raise MiniMaxH3LatentUpscalerError("latent upscaling by dimensions needs both height and width") + target_height = float(height) + target_width = float(width) + + cell = MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE + if not math.isfinite(target_height) or not math.isfinite(target_width): + raise MiniMaxH3LatentUpscalerError("latent upscale target dimensions must be finite") + out_height = max(1, int(round(round(target_height / align) * align / cell))) + out_width = max(1, int(round(round(target_width / align) * align / cell))) + effective = (out_height / latent_height + out_width / latent_width) / 2.0 + if out_height < latent_height or out_width < latent_width: + raise MiniMaxH3LatentUpscalerError( + f"the MiniMax-H3 latent upscaler only upscales: {latent_width}x{latent_height} -> " + f"{out_width}x{out_height} latent cells" + ) + if effective > MINIMAX_H3_LATENT_UPSCALE_MAX_SCALE: + logger.warning( + "MiniMax-H3 latent upscale factor %.3f exceeds the %.1fx the checkpoint was trained for", + effective, + MINIMAX_H3_LATENT_UPSCALE_MAX_SCALE, + ) + return MiniMaxH3LatentUpscaleTarget(latent_height=out_height, latent_width=out_width, scale=effective) + + +class MiniMaxH3LatentUpscaler(nn.Module): + """Serving wrapper around :class:`MiniMaxH3LatentResizer3D`. + + Owns the temporal chunking that keeps long clips inside VRAM and the + residency policy that keeps the extra ~700MB off the device between + requests. + """ + + def __init__( + self, + resizer: MiniMaxH3LatentResizer3D, + *, + device: torch.device, + dtype: torch.dtype, + latents_mean: Sequence[float], + latents_std: Sequence[float], + chunk_frames: int = 0, + chunk_overlap: int | None = None, + resident: bool = False, + ) -> None: + super().__init__() + if chunk_frames < 0: + raise MiniMaxH3LatentUpscalerError(f"latent upscale chunk_frames must be >= 0, got {chunk_frames}") + # Plain tuples rather than buffers: this module manages its own device + # placement, and a buffer would also show up in the pipeline state dict. + self.latents_mean = tuple(float(value) for value in latents_mean) + self.latents_std = tuple(float(value) for value in latents_std) + if len(self.latents_mean) != len(self.latents_std): + raise MiniMaxH3LatentUpscalerError( + f"latents_mean/latents_std length mismatch: {len(self.latents_mean)} vs {len(self.latents_std)}" + ) + if any(value == 0.0 for value in self.latents_std): + raise MiniMaxH3LatentUpscalerError("latents_std must not contain zeros") + self.resizer = resizer.to(dtype=dtype).eval().requires_grad_(False) + self.device = device + self.dtype = dtype + self.chunk_frames = chunk_frames + self.chunk_overlap = resizer.temporal_kernel if chunk_overlap is None else chunk_overlap + self.resident = resident + self.resizer.to(device if resident else torch.device("cpu")) + + def load_to_device(self) -> None: + self.resizer.to(self.device) + + def offload_to_cpu(self) -> None: + self.resizer.to(torch.device("cpu")) + + @torch.inference_mode() + def upscale(self, latent: torch.Tensor, target: MiniMaxH3LatentUpscaleTarget) -> torch.Tensor: + """Upscale a normalized ``[B, 24, T, H, W]`` H3 latent to ``target``.""" + if latent.ndim != 5: + raise MiniMaxH3LatentUpscalerError(f"expected a rank-5 H3 latent, got shape {tuple(latent.shape)}") + size = (int(latent.shape[2]), target.latent_height, target.latent_width) + if size[1:] == tuple(latent.shape[-2:]): + return latent + + if int(latent.shape[1]) != len(self.latents_mean): + raise MiniMaxH3LatentUpscalerError( + f"expected {len(self.latents_mean)} latent channels, got {int(latent.shape[1])}" + ) + source_dtype = latent.dtype + work = latent.to(device=self.device, dtype=self.dtype) + mean, std = self._norm_tensors(work.device, work.dtype) + try: + if not self.resident: + self.resizer.to(self.device) + upscaled = self._run((work - mean) / std, target.scale, size) * std + mean + finally: + if not self.resident: + self.resizer.to(torch.device("cpu")) + return upscaled.to(dtype=source_dtype) + + def _norm_tensors(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + shape = (1, -1, 1, 1, 1) + return ( + torch.tensor(self.latents_mean, device=device, dtype=dtype).view(shape), + torch.tensor(self.latents_std, device=device, dtype=dtype).view(shape), + ) + + def _run(self, latent: torch.Tensor, scale: float, size: tuple[int, int, int]) -> torch.Tensor: + """Resize the clip, in temporal chunks when it does not fit in one pass. + + Chunking is not free: every GroupNorm in the network pools its + statistics over the whole clip it is given, so a chunked pass is an + approximation of the single pass everywhere, not only at the seams. + ``chunk_frames=0`` buys the exact result for the memory a full-length + activation volume costs. + """ + frames = int(latent.shape[2]) + overlap = self.chunk_overlap + if self.chunk_frames == 0 or frames <= self.chunk_frames or overlap == 0: + return self.resizer(latent, scale=scale, target_size=size) + + logger.debug( + "MiniMax-H3 latent upscale: %d frames in chunks of %d with %d-frame overlap", + frames, + self.chunk_frames, + overlap, + ) + # Replicate padding gives the first and last chunk the same amount of + # temporal context as the interior ones, so the clip does not flicker at + # its ends. Unlike the reference node the context is symmetric: every + # chunk sees `overlap` frames on both sides rather than only before. + padded = F.pad(latent, (0, 0, 0, 0, overlap, overlap), mode="replicate") + batch, channels = latent.shape[0], latent.shape[1] + accumulated = torch.zeros(batch, channels, frames, size[1], size[2], device=latent.device, dtype=latent.dtype) + weights = torch.zeros(1, 1, frames, 1, 1, device=latent.device, dtype=latent.dtype) + + for start in range(0, frames, self.chunk_frames): + core_start, core_end = start, min(frames, start + self.chunk_frames) + blend_start, blend_end = max(0, core_start - overlap), min(frames, core_end + overlap) + # Padded coordinates: original frame `i` sits at `i + overlap`. + segment = padded[:, :, blend_start : blend_end + 2 * overlap].contiguous() + decoded = self.resizer( + segment, + scale=scale, + target_size=(int(segment.shape[2]), size[1], size[2]), + )[:, :, overlap : overlap + (blend_end - blend_start)] + + weight = self._blend_weights( + blend_start, + blend_end, + core_start, + core_end, + device=latent.device, + dtype=latent.dtype, + ) + accumulated[:, :, blend_start:blend_end] += decoded * weight + weights[:, :, blend_start:blend_end] += weight + del segment, decoded + + return accumulated / weights.clamp(min=1e-8) + + @staticmethod + def _blend_weights( + blend_start: int, + blend_end: int, + core_start: int, + core_end: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Linear ramps over the frames a chunk shares with its neighbours.""" + weight = torch.ones(blend_end - blend_start, device=device, dtype=dtype) + lead = core_start - blend_start + if lead > 0: + weight[:lead] = torch.arange(1, lead + 1, device=device, dtype=dtype) / (lead + 1) + trail = blend_end - core_end + if trail > 0: + weight[-trail:] = torch.arange(trail, 0, -1, device=device, dtype=dtype) / (trail + 1) + return weight.view(1, 1, -1, 1, 1) + + +_UPSCALER_DTYPES = { + "fp32": torch.float32, + "float32": torch.float32, + "fp16": torch.float16, + "float16": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, +} + + +def _resolve_checkpoint_file(path: str | Path) -> Path: + """Find the single upscaler checkpoint at ``path``.""" + candidate = Path(path).expanduser() + if candidate.is_file(): + return candidate + if not candidate.is_dir(): + raise MiniMaxH3LatentUpscalerError(f"MiniMax-H3 latent upscaler not found at {candidate}") + files = sorted(file for suffix in (".safetensors", ".pth") for file in candidate.glob(f"*{suffix}")) + if len(files) != 1: + raise MiniMaxH3LatentUpscalerError( + f"{candidate} holds {len(files)} latent upscaler checkpoints; point the path at one file" + ) + return files[0] + + +def _read_state_dict(file: Path) -> dict[str, torch.Tensor]: + if file.suffix == ".safetensors": + with safe_open(file, framework="pt", device="cpu") as checkpoint: + state_dict = {key: checkpoint.get_tensor(key) for key in checkpoint.keys()} + elif file.suffix == ".pth": + state_dict = torch.load(file, map_location="cpu", weights_only=True) + if isinstance(state_dict, Mapping) and "model" in state_dict: + state_dict = state_dict["model"] + else: + raise MiniMaxH3LatentUpscalerError(f"unsupported latent upscaler checkpoint {file}") + # Some releases nest the resizer under a training-time wrapper. + if any(key.startswith("upscaler.") for key in state_dict): + return { + key.removeprefix("upscaler."): value for key, value in state_dict.items() if key.startswith("upscaler.") + } + return dict(state_dict) + + +def load_minimax_h3_latent_upscaler( + path: str | Path, + *, + device: torch.device, + dtype: torch.dtype, + latents_mean: Sequence[float], + latents_std: Sequence[float], + chunk_frames: int = 0, + chunk_overlap: int | None = None, + resident: bool = False, +) -> MiniMaxH3LatentUpscaler: + """Build a serving-ready upscaler from a released checkpoint.""" + file = _resolve_checkpoint_file(path) + state_dict = _read_state_dict(file) + arch = detect_minimax_h3_upscaler_arch(state_dict) + # Built on meta and assigned from the checkpoint: the pipeline is + # constructed under a CUDA default-device context, and a real + # initialization would spend ~2.7GB of device memory on values the next + # line overwrites anyway. + with torch.device("meta"): + resizer = MiniMaxH3LatentResizer3D(arch) + resizer.load_state_dict(state_dict, strict=True, assign=True) + logger.info( + "Loaded MiniMax-H3 latent upscaler %s: %d channels, %d+%d blocks, temporal kernel %d, %s", + file.name, + arch.channels, + arch.in_blocks, + arch.out_blocks, + arch.temporal_kernel, + dtype, + ) + return MiniMaxH3LatentUpscaler( + resizer, + device=device, + dtype=dtype, + latents_mean=latents_mean, + latents_std=latents_std, + chunk_frames=chunk_frames, + chunk_overlap=chunk_overlap, + resident=resident, + ) + + +def resolve_minimax_h3_latent_upscaler( + od_config, + *, + device: torch.device, + latent_stats: Callable[[], tuple[Sequence[float], Sequence[float]]], +) -> MiniMaxH3LatentUpscaler | None: + """Build the upscaler declared by ``--additional-config``, if any. + + ``latent_stats`` yields the VAE's per-channel ``latents_mean`` / + ``latents_std`` and is called only once a checkpoint is configured, so a + deployment without this stage never reaches for them. + + Recognized keys: ``latent_upscaler_path`` (required to enable the stage), + ``latent_upscaler_dtype``, ``latent_upscaler_chunk_frames`` (0 disables + temporal chunking) and ``latent_upscaler_resident`` (keep the weights on + the device between requests instead of parking them in host memory). + """ + additional = getattr(od_config, "additional_config", None) or {} + path = additional.get("latent_upscaler_path") + if not path: + return None + raw_dtype = additional.get("latent_upscaler_dtype") + if raw_dtype is None: + dtype = getattr(od_config, "dtype", torch.bfloat16) + elif str(raw_dtype) not in _UPSCALER_DTYPES: + raise MiniMaxH3LatentUpscalerError( + f"latent_upscaler_dtype must be one of {sorted(_UPSCALER_DTYPES)}, got {raw_dtype!r}" + ) + else: + dtype = _UPSCALER_DTYPES[str(raw_dtype)] + latents_mean, latents_std = latent_stats() + return load_minimax_h3_latent_upscaler( + path, + device=device, + dtype=dtype, + latents_mean=latents_mean, + latents_std=latents_std, + chunk_frames=int(additional.get("latent_upscaler_chunk_frames", 0)), + resident=bool(additional.get("latent_upscaler_resident", False)), + ) + + +def _parse_upscale_number(value, *, key: str) -> float | int: + integer = key in {"width", "height", "align"} + try: + number = float(value) + if isinstance(value, bool) or not math.isfinite(number) or number <= 0: + raise ValueError + if integer and not number.is_integer(): + raise ValueError + return int(value) if integer else number + except (TypeError, ValueError, OverflowError) as exc: + kind = "positive integer" if integer else "positive finite number" + raise MiniMaxH3LatentUpscalerError(f"latent_upscale {key} must be a {kind}, got {value!r}") from exc + + +def parse_minimax_h3_latent_upscale_request(value) -> dict[str, float | int] | None: + """Normalize ``extra_args['latent_upscale']`` into resolver keywords. + + Accepts a bare multiplier (``2.0``), ``false``/``null`` to opt out of a + server-side default, or an object naming one sizing mode: + ``{"scale": 2}``, ``{"width": 2688, "height": 1536}`` or + ``{"megapixels": 4}``, each optionally with ``align``. + """ + if value is None or value is False: + return None + if isinstance(value, bool): + raise MiniMaxH3LatentUpscalerError("latent_upscale must be a number or an object, not true") + if isinstance(value, (int, float)): + return {"scale": _parse_upscale_number(value, key="scale")} + if not isinstance(value, Mapping): + raise MiniMaxH3LatentUpscalerError(f"latent_upscale must be a number or an object, got {type(value).__name__}") + known = {"scale", "width", "height", "megapixels", "align"} + unknown = set(value) - known + if unknown: + raise MiniMaxH3LatentUpscalerError(f"unknown latent_upscale keys {sorted(unknown)}; expected {sorted(known)}") + resolved: dict[str, float | int] = {} + for key in ("scale", "megapixels", "width", "height", "align"): + if value.get(key) is not None: + resolved[key] = _parse_upscale_number(value[key], key=key) + return resolved or None + + +@dataclass(frozen=True) +class MiniMaxH3LatentRefineSpec: + """How much of the sigma schedule a second denoise pass re-runs.""" + + strength: float + + def start_index(self, num_sigma_points: int) -> int: + """The schedule position the refine pass starts from. + + ``strength`` is the fraction of the request's steps to re-run, the + img2img convention, so the cost of the pass is predictable from the + request. At least one step always runs. + """ + total_steps = num_sigma_points - 1 + if total_steps < 1: + raise MiniMaxH3LatentUpscalerError("a refine pass needs a schedule with at least two sigma points") + return total_steps - max(1, round(self.strength * total_steps)) + + +def parse_minimax_h3_latent_refine_request(value) -> MiniMaxH3LatentRefineSpec | None: + """Normalize ``extra_args['latent_refine']`` into a refine spec. + + Accepts a bare strength (``0.4``), ``false``/``null`` to opt out of a + server-side default, or ``{"strength": 0.4}``. + """ + if value is None or value is False: + return None + if isinstance(value, bool): + raise MiniMaxH3LatentUpscalerError("latent_refine must be a number or an object, not true") + if isinstance(value, Mapping): + unknown = set(value) - {"strength"} + if unknown: + raise MiniMaxH3LatentUpscalerError(f"unknown latent_refine keys {sorted(unknown)}; expected ['strength']") + if value.get("strength") is None: + return None + value = value["strength"] + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise MiniMaxH3LatentUpscalerError(f"latent_refine strength must be a number, got {type(value).__name__}") + strength = float(value) + if not 0.0 < strength <= 1.0: + raise MiniMaxH3LatentUpscalerError(f"latent_refine strength must be in (0, 1], got {strength}") + return MiniMaxH3LatentRefineSpec(strength=strength) + + +__all__ = [ + "MINIMAX_H3_LATENT_UPSCALE_MAX_SCALE", + "MINIMAX_H3_VAE_SPATIAL_DOWNSAMPLE", + "MiniMaxH3LatentRefineSpec", + "MiniMaxH3LatentResizer3D", + "MiniMaxH3LatentUpscaleTarget", + "MiniMaxH3LatentUpscaler", + "MiniMaxH3LatentUpscalerArch", + "MiniMaxH3LatentUpscalerError", + "detect_minimax_h3_upscaler_arch", + "load_minimax_h3_latent_upscaler", + "parse_minimax_h3_latent_refine_request", + "parse_minimax_h3_latent_upscale_request", + "resolve_minimax_h3_latent_upscale_target", + "resolve_minimax_h3_latent_upscaler", +] 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 d08a5d255..563cce959 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -18,6 +18,7 @@ import numpy as np import torch import torch.distributed as dist import torch.nn as nn +from PIL import Image from transformers import Qwen2TokenizerFast, Qwen3VLProcessor from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -27,6 +28,7 @@ from vllm_omni.diffusion.cache.cachedit import ( CacheDiTBackend, RequestScopedCacheDiTRuntime, ) +from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook from vllm_omni.diffusion.cancellation import check_request_cancellation from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group @@ -86,7 +88,10 @@ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( prepare_encoder_inputs, ) from vllm_omni.model_executor.models.minimax_h3.long_video import validate_encoded_frame_limit -from vllm_omni.model_executor.models.minimax_h3.preprocessing import build_minimax_h3_presentation +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + build_minimax_h3_presentation, + load_minimax_h3_images, +) from vllm_omni.model_executor.models.minimax_h3.reference_video import ( MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, deserialize_prepared_reference_videos, @@ -120,6 +125,15 @@ from .latent_mask import ( minimax_h3_prepare_edit_rows, minimax_h3_video_edit_masks, ) +from .latent_upscaler import ( + MiniMaxH3LatentRefineSpec, + MiniMaxH3LatentUpscalerError, + MiniMaxH3LatentUpscaleTarget, + parse_minimax_h3_latent_refine_request, + parse_minimax_h3_latent_upscale_request, + resolve_minimax_h3_latent_upscale_target, + resolve_minimax_h3_latent_upscaler, +) from .lora import TurboSpec, load_minimax_h3_turbo_lora from .minimax_h3_transformer import ( MiniMaxH3Attention, @@ -137,6 +151,7 @@ from .packed_sequence import ( minimax_h3_packed_sequence_ref2va_blocks, ) from .packed_tokens import ( + minimax_h3_pack_audio_latent, minimax_h3_patchify_video_latent, minimax_h3_unpack_audio_tokens, minimax_h3_unpatchify_video_tokens, @@ -153,8 +168,6 @@ from .time_request import ( from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE, _VideoVAEPartProxy if TYPE_CHECKING: - from PIL import Image - from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -282,6 +295,9 @@ _MINIMAX_H3_DENOISE_INPUT_KEYS = ( "audio_edit_restore_mask_rows", ) +# Request-context key for the FL2VA keyframes re-encoded at the refine size. +_REFINE_KEYFRAME_CONDITION = "minimax_h3_refine_keyframe_condition" + # ``StepRequestState.extra`` keys owned by the step-execution path. _STEP_BRANCH = "minimax_h3_branch" _STEP_AUDIO_ROWS = "minimax_h3_audio_rows" @@ -351,6 +367,16 @@ def resolve_minimax_h3_diffusion_model_path( return str(model_root / subdir) +def _minimax_h3_output_canvas( + shape: Mapping[str, Any], + target: MiniMaxH3LatentUpscaleTarget | None, +) -> tuple[int, int]: + """The decoded frame size, which the upscaler moves when it runs.""" + if target is None: + return int(shape["height"]), int(shape["width"]) + return target.height, target.width + + def _minimax_h3_post_process(output, output_type: str = "np"): """Convert the joint video/audio output without capturing worker state. @@ -1076,6 +1102,19 @@ class MiniMaxH3Pipeline( ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae + # Optional learned latent super-resolution, run between the denoise + # loop and the VAE. Absent unless --additional-config names a + # checkpoint, so a plain H3 deployment carries none of its weights. + # The upscaler works one normalization below the pipeline latent, so it + # needs the same per-channel statistics the VAE denormalizes with. + self.latent_upscaler = resolve_minimax_h3_latent_upscaler( + od_config, + device=self.device, + latent_stats=lambda: ( + self.video_vae.config_dict["latents_mean"], + self.video_vae.config_dict["latents_std"], + ), + ) self._dlo_component_cache = None offloads_text_encoder = should_offload_component(od_config, TEXT_ENCODER_COMPONENT) @@ -1138,8 +1177,8 @@ class MiniMaxH3Pipeline( # ``weights_sources``. The text encoder uses the shared component # loader so online quantization and offload processing follow the same # path as the DiT. - for component_name in ("video_vae", "audio_vae"): - component = getattr(self, component_name) + for component_name in ("video_vae", "audio_vae", "latent_upscaler"): + component = getattr(self, component_name, None) if component is None: continue loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) @@ -1458,20 +1497,33 @@ class MiniMaxH3Pipeline( components = ModuleDiscovery.discover(self) dits = components.dits + # This optional stage owns its own weight placement. Register it only + # with model-level sequential offload; generic VAE discovery would + # otherwise move it onto the GPU during unrelated layerwise setup. + upscaler = getattr(self, "latent_upscaler", None) stages = [*components.encoders, *components.vaes] + if upscaler is not None: + stages.append(upscaler) modules = [*dits, *stages] + # The upscaler normally parks its own weights on the host. Keep it as + # an execution stage so activating it evicts the DiT, but do not scan + # its parameters on every DiT step unless residency was requested. + selected_stages = [stage for stage in stages if stage is not upscaler or upscaler.resident] selection_options: dict[str, Any] = {} if offload_components is not None: if DIT_COMPONENT in offload_components and not dits: raise ValueError("MiniMax-H3 has no loaded DiT for selected module offload") if TEXT_ENCODER_COMPONENT in offload_components and not components.encoders: raise ValueError("MiniMax-H3 has no loaded text encoder for selected module offload") + selected_explicit_stages = [*components.encoders] if TEXT_ENCODER_COMPONENT in offload_components else [] + if upscaler is not None and upscaler.resident: + selected_explicit_stages.append(upscaler) selection_options = { "offload_dit_modules": dits if DIT_COMPONENT in offload_components else (), - "offload_encoder_modules": ( - components.encoders if TEXT_ENCODER_COMPONENT in offload_components else () - ), + "offload_encoder_modules": selected_explicit_stages, } + elif upscaler is not None: + selection_options["offload_encoder_modules"] = selected_stages apply_sequential_offload( dit_modules=dits, encoder_modules=stages, @@ -1556,22 +1608,26 @@ class MiniMaxH3Pipeline( rows: list[torch.Tensor] = [] shapes: list[tuple[int, int, int]] = [] _, rank, _ = _dit_rank_world() + # encode_image retains parallel tiling when there are enough tiles. + # Every VAE rank must finish its codec collectives before the latent + # broadcast, including when refinement enlarges the keyframe images. + encode_images = bool(images) and (rank == 0 or self.video_vae.is_distributed_enabled()) # Keep image and video references in one residency window when both # appear in a request; otherwise the video branch would reload the VAE. # Encoding touches only the CNN encoder half, so the 9GB ViT decoder # stays off the device for the whole window. - needs_video_vae = video_count > 0 or (rank == 0 and bool(images)) + needs_video_vae = video_count > 0 or encode_images video_vae_context = ( self._component_on_device(self.video_vae.encoder_component) if needs_video_vae else nullcontext() ) with video_vae_context: if images: image_rows = None - if rank == 0: + if encode_images: image_rows = torch.cat([self.video_vae.encode_image(image) for image in images]) rows.append( _broadcast_tensor( - image_rows, + image_rows if rank == 0 else None, dtype=torch.float32, device=self.device, ) @@ -1642,6 +1698,43 @@ class MiniMaxH3Pipeline( ) return video_rows, audio_rows + @staticmethod + def _renoise_rows( + init_latents: tuple[torch.Tensor, torch.Tensor], + *, + noise_video: torch.Tensor, + noise_audio: torch.Tensor, + sigma_video: float, + sigma_audio: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Put finished latents back on the flow at the given sigmas. + + The rectified-flow forward process is ``x_s = (1 - s) * x0 + s * noise`` + -- the interpolation :func:`minimax_h3_euler_eta0_step` walks back down + -- so this lands the latents exactly where a full pass would have been + at that sigma, which is what makes resuming mid-schedule in-distribution. + """ + video_latent, audio_latent = init_latents + video_prior = minimax_h3_patchify_video_latent( + video_latent.detach().to(device="cpu", dtype=torch.float32), + patch_size=(1, 2, 2), + ) + audio_prior = minimax_h3_pack_audio_latent(audio_latent.detach().to(device="cpu", dtype=torch.float32)) + if video_prior.shape != noise_video.shape: + raise ValueError( + f"refine video rows {tuple(video_prior.shape)} do not match the " + f"target layout {tuple(noise_video.shape)}" + ) + if audio_prior.shape != noise_audio.shape: + raise ValueError( + f"refine audio rows {tuple(audio_prior.shape)} do not match the " + f"target layout {tuple(noise_audio.shape)}" + ) + return ( + (1.0 - sigma_video) * video_prior + sigma_video * noise_video, + (1.0 - sigma_audio) * audio_prior + sigma_audio * noise_audio, + ) + @contextmanager def _resident_dit_layers_on_device(self, *, enabled: bool = True): controller = getattr(self, "_dlo_residency_controller", None) @@ -1687,12 +1780,28 @@ 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, + init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, + refine: MiniMaxH3LatentRefineSpec | None = None, ) -> dict[str, Any]: """Build the packed layout, initial rows, anchors, and sigma schedules. Shared by request-mode :meth:`diffuse` and step-mode :meth:`prepare_encode` so both paths start from identical state. + + ``init_latents`` and ``refine`` turn the pass into a second, partial + one: the rows start from those latents re-noised to the schedule + position ``refine`` selects, and the returned schedules begin there. """ + video_sigmas = minimax_h3_time_shift_sigmas( + num_steps=num_steps, + shift_scale=video_shift, + base_schedule=base_schedule, + ) + audio_sigmas = minimax_h3_time_shift_sigmas( + num_steps=num_steps, + shift_scale=audio_shift, + base_schedule=base_schedule, + ) initial_video, initial_audio = self._initial_noise( seed=seed, latent_t=latent_t, @@ -1700,6 +1809,23 @@ class MiniMaxH3Pipeline( latent_w=latent_w, audio_t=audio_t, ) + if init_latents is not None: + if refine is None: + raise ValueError("init_latents needs a refine spec to place them on the schedule") + # Video and audio are shifted apart (12.0 against 3.0 by default), + # so the two schedules hold different sigmas at the same position. + # The loop steps them by index, so the pass has to resume at one + # index and re-noise each modality to its own sigma there. + start = refine.start_index(len(video_sigmas)) + video_sigmas = video_sigmas[start:] + audio_sigmas = audio_sigmas[start:] + initial_video, initial_audio = self._renoise_rows( + init_latents, + noise_video=initial_video, + noise_audio=initial_audio, + sigma_video=video_sigmas[0], + sigma_audio=audio_sigmas[0], + ) if task == "ref2va": if ref_blocks is None: if visual_condition_shape is None or ref_audio_t is None: @@ -1839,16 +1965,6 @@ class MiniMaxH3Pipeline( audio_edit_restore_mask_rows, ) - video_sigmas = minimax_h3_time_shift_sigmas( - num_steps=num_steps, - shift_scale=video_shift, - base_schedule=base_schedule, - ) - audio_sigmas = minimax_h3_time_shift_sigmas( - num_steps=num_steps, - shift_scale=audio_shift, - base_schedule=base_schedule, - ) return { "branch": branch, # The request-mode loop moves these onto the device itself; step mode @@ -1932,6 +2048,8 @@ 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, + init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, + refine: MiniMaxH3LatentRefineSpec | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = self._build_denoise_inputs( task=task, @@ -1965,9 +2083,19 @@ class MiniMaxH3Pipeline( audio_edit_clean_rows=audio_edit_clean_rows, audio_edit_mask_rows=audio_edit_mask_rows, audio_edit_restore_mask_rows=audio_edit_restore_mask_rows, + init_latents=init_latents, + refine=refine, ) branch = inputs["branch"] transformer = self._transformer_for_task(task) + # Each pass (including each output and hi-res refine) owns its cache. + # Refine can change both the packed shape and the schedule length. + registry = getattr(transformer, "_hook_registry", None) + if registry is not None: + registry.reset_hook(TeaCacheHook._HOOK_NAME) + cache_runtime = getattr(self, "_cache_dit_runtime", None) + if cache_runtime is not None: + cache_runtime.refresh(len(inputs["sigmas_video"]) - 1) 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( @@ -2097,6 +2225,188 @@ class MiniMaxH3Pipeline( audio = self._offload_model_cpu_stage_output(audio) return video, audio + def _resolve_latent_upscale( + self, + extra: Mapping[str, Any], + *, + latent_h: int, + latent_w: int, + ) -> MiniMaxH3LatentUpscaleTarget | None: + """Resolve the requested latent super-resolution target, if any. + + A request opts in with ``extra_args['latent_upscale']``; a deployment + can set the same value under ``--additional-config`` to upscale every + request, which a request then overrides (``false`` opts back out). + Resolving here means an unserviceable size is rejected before the + denoise loop rather than after it. + """ + additional = getattr(getattr(self, "od_config", None), "additional_config", None) or {} + raw = extra["latent_upscale"] if "latent_upscale" in extra else additional.get("latent_upscale") + try: + spec = parse_minimax_h3_latent_upscale_request(raw) + if spec is None: + return None + if self.latent_upscaler is None: + raise MiniMaxH3LatentUpscalerError( + "latent_upscale needs a checkpoint: serve with " + '--additional-config \'{"latent_upscaler_path": ""}\'' + ) + target = resolve_minimax_h3_latent_upscale_target( + latent_height=latent_h, + latent_width=latent_w, + **spec, + ) + if self._resolve_latent_refine(extra) is not None and (target.latent_height % 2 or target.latent_width % 2): + raise MiniMaxH3LatentUpscalerError( + "latent_refine requires upscale latent height and width divisible by 2 " + "(32 pixels); use align=32 or a compatible target size" + ) + except MiniMaxH3LatentUpscalerError as exc: + raise OmniClientError(str(exc)) from exc + if (target.latent_height, target.latent_width) == (latent_h, latent_w): + return None + logger.info( + "MiniMax H3 latent upscale %dx%d -> %dx%d (scale %.3f)", + latent_w * 16, + latent_h * 16, + target.width, + target.height, + target.scale, + ) + return target + + def _resolve_latent_refine(self, extra: Mapping[str, Any]) -> MiniMaxH3LatentRefineSpec | None: + """Resolve the requested second denoise pass, if any. + + Read like ``latent_upscale``: ``extra_args['latent_refine']`` wins over + an ``--additional-config`` default, and ``false`` opts back out. + """ + additional = getattr(getattr(self, "od_config", None), "additional_config", None) or {} + raw = extra["latent_refine"] if "latent_refine" in extra else additional.get("latent_refine") + try: + return parse_minimax_h3_latent_refine_request(raw) + except MiniMaxH3LatentUpscalerError as exc: + raise OmniClientError(str(exc)) from exc + + def _validate_refine_token_budget( + self, + *, + task: str, + target: MiniMaxH3LatentUpscaleTarget | None, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + text_len: int, + keyframe_count: int, + ref_blocks: list[dict[str, Any]] | None, + ) -> None: + """Reject refine layouts known to exceed the tested per-rank limit. + + Count the same rows as the packed layout without materializing its + large position tensors before the first denoise pass. + """ + additional = getattr(getattr(self, "od_config", None), "additional_config", None) or {} + limit = additional.get("latent_refine_max_tokens_per_rank", 65_536) + if type(limit) is not int or limit < 0: + raise OmniClientError("latent_refine_max_tokens_per_rank must be a non-negative integer") + if limit == 0: + return + height = target.latent_height if target is not None else latent_h + width = target.latent_width if target is not None else latent_w + frame_rows = (height // 2) * (width // 2) + used = text_len + 2 * audio_t + latent_t * frame_rows + if task == "fl2va": + used += keyframe_count * frame_rows + elif task == "ref2va": + for block in ref_blocks or (): + kind = block["kind"] + if kind == "image": + used += (block["latent_h"] // 2) * (block["latent_w"] // 2) + elif kind == "audio": + used += 2 * block["ref_audio_t"] + else: + used += 2 * block["ref_audio_t"] + used += block["latent_t"] * (block["latent_h"] // 2) * (block["latent_w"] // 2) + padded = ((used + MINIMAX_H3_SEQ_ALIGN - 1) // MINIMAX_H3_SEQ_ALIGN) * MINIMAX_H3_SEQ_ALIGN + parallel = getattr(self, "parallel_config", None) + degree = int(getattr(parallel, "ulysses_degree", 1)) + per_rank = (padded + degree - 1) // degree + if per_rank > limit: + raise OmniClientError( + f"MiniMax H3 latent_refine needs about {per_rank:,} packed tokens per Ulysses rank " + f"(limit {limit:,}); reduce the target size or duration, increase --usp, " + "or adjust latent_refine_max_tokens_per_rank for a validated deployment" + ) + + def _refine_keyframe_condition( + self, + context: dict[str, Any], + ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]: + """Encode the FL2VA keyframes at the refine size, once per request. + + The encode broadcasts across the DiT group, so every rank has to reach + it the same number of times; caching on the request context keeps that + true while sparing the repeat for each additional output. + """ + cached = context.get(_REFINE_KEYFRAME_CONDITION) + if cached is None: + target = context["latent_upscale"] + cached = self._encode_visual_conditions( + [ + image.resize((target.width, target.height), Image.Resampling.LANCZOS) + for image in context["keyframe_images"] + ], + None, + video_count=0, + ) + context[_REFINE_KEYFRAME_CONDITION] = cached + return cached + + def _refined_latents( + self, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + *, + context: dict[str, Any], + seed: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Re-denoise finished latents at their current size. + + This is the second half of the hi-res route: the first pass generates + at a cheap size, the upscaler moves the latent, and this pass resumes + the schedule at the target size to put back the detail the upscaler can + only approximate. Without an upscale it is an ordinary detail pass. + """ + refine = context["latent_refine"] + target = context["latent_upscale"] + kwargs = dict(self._denoise_kwargs(context)) + kwargs["seed"] = seed + if target is not None: + kwargs["latent_h"] = target.latent_height + kwargs["latent_w"] = target.latent_width + # A pinned pad_seq_len sizes the first pass's layout. This pass packs + # several times as many video rows, so carrying the pin over would + # fail the seq_len >= used check with an error that names neither + # the refine pass nor the size that outgrew it. + kwargs["pad_seq_len"] = None + if context["task"] == "fl2va": + condition, shapes = self._refine_keyframe_condition(context) + kwargs["visual_condition"] = condition + kwargs["visual_condition_shapes"] = shapes + kwargs["visual_condition_shape"] = shapes[0] if len(shapes) == 1 else None + return self.diffuse(**kwargs, init_latents=(video_latent, audio_latent), refine=refine) + + def _upscaled_latent( + self, + video_latent: torch.Tensor, + target: MiniMaxH3LatentUpscaleTarget | None, + ) -> torch.Tensor: + if target is None: + return video_latent + with self._component_on_device(self.latent_upscaler): + return self.latent_upscaler.upscale(video_latent, target) + @staticmethod def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]: """Split a request prompt into its text and multimodal parts.""" @@ -2398,6 +2708,19 @@ class MiniMaxH3Pipeline( else: conditioning = self._extract_encoder_conditioning(raw_prompt) context = self._prepare_encoder_conditioning_inputs(conditioning, sampling) + if ( + context.get("latent_refine") is not None + and context.get("latent_upscale") is not None + and context["task"] == "fl2va" + ): + if not self.load_vae_encoder: + raise OmniClientError("MiniMax H3 FL2VA latent_refine with upscale requires a local VAE encoder") + _, media = self._extract_prompt(raw_prompt) + images = media.get("image") + images = list(images) if isinstance(images, (list, tuple)) else [images] if images is not None else [] + context["keyframe_images"] = load_minimax_h3_images(images) + if len(context["keyframe_images"]) != len(context["keyframe_frame_indices"] or ()): + raise OmniClientError("MiniMax H3 latent_refine requires the original FL2VA keyframe images") if window_text is not None: context["continuation_text_conditioning"] = window_text return context @@ -2566,9 +2889,32 @@ class MiniMaxH3Pipeline( if continuation is not None and quality_plan.cache_dit is not None: raise OmniClientError("MiniMax H3 continuation requires uncached denoising; set quality=lossless") self._cache_dit_runtime.prepare(quality_plan.cache_dit) + upscale_target = self._resolve_latent_upscale( + extra, latent_h=conditioning.height // 16, latent_w=conditioning.width // 16 + ) + latent_refine = self._resolve_latent_refine(extra) + if latent_refine is not None: + if continuation is not None: + raise OmniClientError("MiniMax H3 latent_refine does not support latent-tail continuation") + if conditioning.video_edit_clean_rows is not None or conditioning.audio_edit_clean_rows is not None: + raise OmniClientError("MiniMax H3 latent_refine does not support latent-mask editing") + self._validate_refine_token_budget( + task=task, + target=upscale_target, + latent_t=conditioning.latent_t, + latent_h=conditioning.height // 16, + latent_w=conditioning.width // 16, + audio_t=conditioning.audio_t, + text_len=int(conditioning.hidden_states.shape[0]), + keyframe_count=len(conditioning.keyframe_frame_indices), + ref_blocks=list(conditioning.ref_blocks) or None, + ) return { "continuation": continuation, "task": task, + "latent_upscale": upscale_target, + "latent_refine": latent_refine, + "keyframe_images": [], "height": conditioning.height, "width": conditioning.width, "num_frames": conditioning.num_frames, @@ -2629,6 +2975,9 @@ class MiniMaxH3Pipeline( check_request_cancellation() denoise_kwargs = self._denoise_kwargs(context) num_outputs = context["num_outputs"] + upscale_target = context.get("latent_upscale") + latent_refine = context.get("latent_refine") + height, width = _minimax_h3_output_canvas(context, upscale_target) videos = [] audios = [] for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): @@ -2646,13 +2995,22 @@ class MiniMaxH3Pipeline( text_conditioning=context.get("continuation_text_conditioning"), ) check_request_cancellation() + if upscale_target is not None: + video_latent = self._upscaled_latent(video_latent, upscale_target) + if latent_refine is not None: + video_latent, audio_latent = self._refined_latents( + video_latent, + audio_latent, + context=context, + seed=output_seed, + ) if context["preencode_mp4"]: videos.append( self.decode_to_mp4( video_latent, audio_latent, - height=context["height"], - width=context["width"], + height=height, + width=width, video_codec_options=context["video_codec_options"], batch_frames=context["preencode_batch_frames"], ) @@ -2662,8 +3020,8 @@ class MiniMaxH3Pipeline( video, audio = self.decode( video_latent, audio_latent, - height=context["height"], - width=context["width"], + height=height, + width=width, ) # Rebind rather than append the expression: the local would # otherwise keep the decoded frames' device storage alive for the @@ -2751,6 +3109,17 @@ class MiniMaxH3Pipeline( "co-batched requests would reuse incompatible cache state. Drop --step-execution " "or omit quality=high." ) + # A refine pass is a second denoise loop, at its own resolution and over + # its own truncated schedule. The step contract gives the scheduler one + # latent and one schedule per request, so there is nowhere to put it. + # Latent upscaling alone has no such problem and stays supported: it + # runs once in ``post_decode``. + if self._resolve_latent_refine(getattr(state.sampling, "extra_args", None) or {}) is not None: + raise OmniClientError( + "MiniMax H3 step execution does not support latent_refine; it is a second denoise " + "loop and the step contract carries one schedule per request. Drop --step-execution, " + "or keep latent_upscale without latent_refine." + ) context = self._prepare_request_inputs( state.prompt, state.sampling, @@ -2810,6 +3179,7 @@ class MiniMaxH3Pipeline( "preencode_mp4": context.get("preencode_mp4", False), "preencode_batch_frames": context.get("preencode_batch_frames", 17), "video_codec_options": context.get("video_codec_options"), + "latent_upscale": context.get("latent_upscale"), }, } ) @@ -3029,12 +3399,16 @@ class MiniMaxH3Pipeline( latent_w=shape["latent_w"], audio_t=shape["audio_t"], ) + upscale_target = shape.get("latent_upscale") + if upscale_target is not None: + video_latent = self._upscaled_latent(video_latent, upscale_target) + height, width = _minimax_h3_output_canvas(shape, upscale_target) if shape.get("preencode_mp4", False): video = self.decode_to_mp4( video_latent, audio_latent, - height=shape["height"], - width=shape["width"], + height=height, + width=width, video_codec_options=shape.get("video_codec_options"), batch_frames=shape.get("preencode_batch_frames", 17), ) @@ -3043,8 +3417,8 @@ class MiniMaxH3Pipeline( video, audio = self.decode( video_latent, audio_latent, - height=shape["height"], - width=shape["width"], + height=height, + width=width, ) video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) self._release_stage_cache()