From 85ecb827ab9cf76ac4ea3aa793d7db047ed19a57 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 2 Aug 2026 17:17:31 +0000 Subject: [PATCH 01/16] Let the MiniMax-H3 video VAE encode a single frame `encode` padded a lone frame up to `clip_length` by repeating it, ran the temporal path over 17 copies and dropped `token_drop` latent frames, returning two latent frames rather than one. That is why the blocks reached past it into `_encode_clip` / `_encode`, which in turn needed `@apply_forward_hook` of their own to onload the VAE under offloading. Encode a still through the spatial encoder alone, so `encode` is correct for it and only `encode` / `decode` carry the hook, as in every other autoencoder. Co-Authored-By: Claude Opus 5 (1M context) --- .../autoencoders/autoencoder_kl_minimax_h3.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index 586138fc884e..56d4707092b3 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -707,14 +707,8 @@ def _stitch_tiles( result_rows.append(torch.cat(result_row, dim=-1)) return torch.cat(result_rows, dim=-2) - @apply_forward_hook def _encode_clip(self, x: torch.Tensor) -> torch.Tensor: - r""" - Encode one temporal clip, spatially tiled when tiling is enabled. - - MiniMax-H3 encodes a keyframe or an image reference through this method rather than through [`~encode`], - because a single frame must not go through the temporal chunking, so it carries the offload hook too. - """ + r"""Encode one temporal clip, spatially tiled when tiling is enabled.""" if not self.use_tiling: return self.quant_conv(self.encoder(x)) @@ -768,17 +762,19 @@ def _decode_clip(self, z: torch.Tensor) -> torch.Tensor: return self._stitch_tiles(rows, y_overlaps, x_overlaps) - @apply_forward_hook def _encode(self, x: torch.Tensor) -> torch.Tensor: r""" Encode a video in `clip_length`-frame chunks and drop the `token_drop` trailing latent frames. - MiniMax-H3 encodes a video reference through this method rather than through [`~encode`], because the - posterior is sampled under a fixed generator rather than through the distribution object, so it carries the - offload hook too. + A single frame has no temporal extent to chunk, so it goes through the spatial encoder alone. Padding it up to + `clip_length` by repetition instead would run the temporal path over `clip_length` copies of the same image and + return `clip_length // temporal_compression_ratio - token_drop` latent frames rather than one — which is not + the conditioning MiniMax-H3 was trained with. """ clip_length = self.config.clip_length num_frames = x.shape[2] + if num_frames == 1: + return self._encode_clip(x) if num_frames % clip_length != 0: pad_frames = x[:, :, -1:].repeat(1, 1, (-num_frames) % clip_length, 1, 1) x = torch.cat([x, pad_frames], dim=2) From ed785ded37fc5c510a983ec5b7cc045cae6a1457 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 2 Aug 2026 17:17:40 +0000 Subject: [PATCH 02/16] Leave the MiniMax-H3 encoders encoding only The keyframe and reference encoders drew the conditioning noise, mixed it in at the noise-augmentation level and packed the result into rows, so they owned a slice of latent preparation and had to be handed the target latent geometry and the request generator to do it. They now return the encoded latents and nothing else, one tensor per condition. Both also go through the public `vae.encode` rather than `_encode_clip` / `_encode`, and `encode_keyframes` takes the VAE directly instead of the whole component bag. The draw order is unchanged, which is what keeps a seeded request reproducible: one draw per condition, in packed order, ahead of the video and audio noise. Co-Authored-By: Claude Opus 5 (1M context) --- .../modular_pipelines/minimax_h3/encoders.py | 143 +++++------------- 1 file changed, 40 insertions(+), 103 deletions(-) diff --git a/src/diffusers/modular_pipelines/minimax_h3/encoders.py b/src/diffusers/modular_pipelines/minimax_h3/encoders.py index da2d611eaba0..26e3e2b55e5c 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/encoders.py +++ b/src/diffusers/modular_pipelines/minimax_h3/encoders.py @@ -17,22 +17,17 @@ from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor from ...models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio -from ...models.autoencoders.vae import DiagonalGaussianDistribution -from ...schedulers import MiniMaxH3Scheduler from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline from .packing import ( MINIMAX_H3_KEYFRAME_ENCODE_SEED, - MINIMAX_H3_KEYFRAME_NOISE_AUG, MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD, MINIMAX_H3_TEXT_ENCODER_LAYER, MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG, - keyframe_condition_noise, - patchify_video_latents, ) from .packing_ref2va import ( MiniMaxH3PreparedReference, @@ -219,17 +214,14 @@ class MiniMaxH3KeyframeVaeEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Encodes the `fl2va` keyframes into packed conditioning rows and noises them to MiniMax-H3's " - "conditioning level. The rows are the anchors of the whole denoising loop: the loop only ever writes the " - "generated rows, so they are never updated again." + "Encodes the `fl2va` keyframes into conditioning latents. They become the anchors of the whole denoising " + "loop, which only ever writes the generated rows, so they are never updated again — the prepare-latents " + "step noises them to MiniMax-H3's conditioning level and packs them." ) @property def expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("vae", AutoencoderKLMiniMaxH3), - ComponentSpec("scheduler", MiniMaxH3Scheduler), - ] + return [ComponentSpec("vae", AutoencoderKLMiniMaxH3)] @property def inputs(self) -> list[InputParam]: @@ -240,15 +232,6 @@ def inputs(self) -> list[InputParam]: required=True, description="The keyframes put onto the target canvas, in packed order.", ), - InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), - InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), - InputParam.template( - "generator", - description=( - "The generator of the request. The conditioning noise is drawn from it before the target noise " - "of the prepare-latents step." - ), - ), ] @property @@ -256,67 +239,58 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "condition_latents", - type_hint=torch.Tensor, - description="The noise-augmented video conditioning rows, in packed order.", + type_hint=list[torch.Tensor], + description=( + "The normalized video conditioning latents, one `(1, latent_channels, 1, latent_height, " + "latent_width)` tensor per keyframe, in packed order." + ), ) ] @staticmethod - def encode_keyframes(components, images: list, device: torch.device | None = None) -> torch.Tensor: + def encode_keyframes(vae, images: list, device: torch.device) -> list[torch.Tensor]: r""" - Encode the `fl2va` keyframes into packed conditioning rows. + Encode the `fl2va` keyframes into normalized conditioning latents. - The keyframes go through the video VAE's spatial encoder only — they are single frames, so none of its - 17-frame temporal chunking applies — and the posterior is *sampled*, under a generator seeded with 42 - independently of the request seed. The sampled latent is rounded to float16 before being normalized, as in the - reference implementation; both are part of reproducing the released model's conditioning. + A keyframe is a single frame, so `vae.encode` runs its spatial encoder alone with none of the 17-frame + temporal chunking. The posterior is *sampled*, under a generator seeded with 42 independently of the request + seed, and the sampled latent is rounded to float16 before being normalized, as in the reference + implementation; both are part of reproducing the released model's conditioning. Args: + vae (`AutoencoderKLMiniMaxH3`): The video VAE. images (`list[PIL.Image.Image]`): The keyframes, already prepared onto the target canvas, in packed order. - device (`torch.device`, *optional*): The device to run the VAE on. + device (`torch.device`): The device to run the VAE on. Returns: - `torch.Tensor` of shape `(num_condition_rows, latent_channels * prod(patch_size))`: the float32 - conditioning rows. + `list[torch.Tensor]`: one `(1, latent_channels, 1, latent_height, latent_width)` float32 CPU tensor per + keyframe, in packed order. One entry per condition is what the prepare-latents step draws its noise + against, so the list is the unit the request's generator is consumed in. """ - device = device or components._execution_device - latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1) - latents_std = torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1) + latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1) pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) - rows = [] + keyframe_latents = [] for image in images: pixels = torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None] pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std - # `vae.encode` chunks along time for videos; a keyframe is one frame and is encoded by the (tiled) - # spatial encoder alone, which is what the released model conditions on. - moments = components.vae._encode_clip(pixels) - posterior = DiagonalGaussianDistribution(moments) + posterior = vae.encode(pixels, return_dict=False)[0] latents = posterior.sample(generator=torch.Generator().manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) # The sampled latent is rounded to float16 before it is normalized: ~11 bits of every conditioning # latent, so the released model's conditioning cannot be reproduced without it. latents = latents.to(torch.float16).float().cpu() - rows.append(patchify_video_latents((latents - latents_mean) / latents_std, components.patch_size)) - return torch.cat(rows) + keyframe_latents.append((latents - latents_mean) / latents_std) + return keyframe_latents @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - condition_latents = self.encode_keyframes(components, block_state.keyframes, device=device) - noise = keyframe_condition_noise( - ((1, block_state.latent_height, block_state.latent_width),) * len(block_state.keyframes), - components.patch_size, - components.vae_latent_channels, - generator=block_state.generator, - device=device, - ) - block_state.condition_latents = components.scheduler.scale_noise( - condition_latents.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise - ) + block_state.condition_latents = self.encode_keyframes(components.vae, block_state.keyframes, device) self.set_block_state(state, block_state) return components, state @@ -480,11 +454,11 @@ class MiniMaxH3Ref2VAReferenceEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Encodes the `ref2va` references into packed conditioning rows — image and video references through the " - "video VAE, soundtracks through the audio VAE — and noises the visual ones to MiniMax-H3's conditioning " - "level. Audio references ride along clean, at `t = 1.0`. Both are anchors of the whole denoising loop, " - "which only ever writes the generated rows. The latent geometry of every reference is resolved here, so " - "this runs before the packed layout is built." + "Encodes the `ref2va` references — image and video references through the video VAE, soundtracks through " + "the audio VAE. They are the anchors of the whole denoising loop, which only ever writes the generated " + "rows; the prepare-latents step noises the visual ones to MiniMax-H3's conditioning level and packs them, " + "while soundtracks ride along clean at `t = 1.0`. The latent geometry of every reference is resolved " + "here, so this runs before the packed layout is built." ) @property @@ -492,7 +466,6 @@ def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKLMiniMaxH3), ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio), - ComponentSpec("scheduler", MiniMaxH3Scheduler), ] @property @@ -504,13 +477,6 @@ def inputs(self) -> list[InputParam]: required=True, description="The prepared references, in packed order. Their latent geometry is filled in here.", ), - InputParam.template( - "generator", - description=( - "The generator of the request. The conditioning noise is drawn from it before the target noise " - "of the prepare-latents step." - ), - ), ] @property @@ -518,9 +484,10 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "condition_latents", - type_hint=torch.Tensor, + type_hint=list[torch.Tensor], description=( - "The noise-augmented video conditioning rows of the image and video references, in packed order, " + "The encoded video conditioning latents of the image and video references, one `(1, " + "latent_channels, num_latent_frames, latent_height, latent_width)` tensor each in packed order, " "or None when the references carry none." ), ), @@ -565,7 +532,7 @@ def encode_references( audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean).view(1, 1, -1) audio_latents_std = torch.tensor(components.audio_vae.config.latents_std).view(1, 1, -1) - video_rows, audio_rows = [], [] + video_latents, audio_rows = [], [] for reference in references: if reference.kind != "audio": if reference.kind == "image": @@ -576,21 +543,14 @@ def encode_references( pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std # A single frame is encoded by the (tiled) spatial encoder alone; a video goes through the temporal # chunking, which is what turns `17 * n + 5` frames into `5 * n + 2` latent frames. - moments = ( - components.vae._encode_clip(pixels) - if reference.kind == "image" - else components.vae._encode(pixels) - ) - posterior = DiagonalGaussianDistribution(moments) + posterior = components.vae.encode(pixels, return_dict=False)[0] latents = posterior.sample(generator=torch.Generator().manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) # The sampled latent is rounded to float16 before it is normalized: ~11 bits of every conditioning # latent, so the released model's conditioning cannot be reproduced without it. latents = latents.to(torch.float16).float().cpu() reference.num_latent_frames = latents.shape[2] reference.latent_height, reference.latent_width = latents.shape[3], latents.shape[4] - video_rows.append( - patchify_video_latents((latents - latents_mean) / latents_std, components.patch_size) - ) + video_latents.append((latents - latents_mean) / latents_std) if reference.has_audio: posterior = components.audio_vae.encode(reference.waveform.to(device)[:, None], return_dict=False)[0] @@ -600,39 +560,16 @@ def encode_references( normalized = (latents - audio_latents_mean) / audio_latents_std audio_rows.append(normalized.reshape(-1, components.audio_latent_channels)) - return ( - torch.cat(video_rows) if video_rows else None, - torch.cat(audio_rows) if audio_rows else None, - ) + return video_latents or None, torch.cat(audio_rows) if audio_rows else None @torch.no_grad() def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - condition_latents, audio_condition_latents = self.encode_references( + block_state.condition_latents, block_state.audio_condition_latents = self.encode_references( components, block_state.prepared_references, device=device ) - if condition_latents is not None: - noise = keyframe_condition_noise( - tuple( - (reference.num_latent_frames, reference.latent_height, reference.latent_width) - for reference in block_state.prepared_references - if reference.kind != "audio" - ), - components.patch_size, - components.vae_latent_channels, - generator=block_state.generator, - device=device, - ) - condition_latents = components.scheduler.scale_noise( - condition_latents.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise - ) - if audio_condition_latents is not None: - audio_condition_latents = audio_condition_latents.to(device) - - block_state.condition_latents = condition_latents - block_state.audio_condition_latents = audio_condition_latents self.set_block_state(state, block_state) return components, state From b97a4999bfc72fa5ec5121be883ce03acadeae8e Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 2 Aug 2026 17:17:52 +0000 Subject: [PATCH 03/16] Turn the MiniMax-H3 setup step into a resize step `MiniMaxH3SetupStep` ran ahead of both encoders and handed them derived values, so neither could be popped out and run on raw inputs. It also resolved the frame count and the latent geometry, which no encoder reads. What is left is a resize step, in the shape `QwenImageEditResizeStep` uses: raw `image` / `last_image` in, a `VaeImageProcessor` of its own, canvas-sized keyframes out. It is wrapped in a conditional block so a text-only request skips it rather than running it empty. `exif_transpose` and `convert("RGB")` are gone because `load_image` already applies both. The follower keyframe keeps MiniMax-H3's own cover-crop rather than moving to `resize_mode="crop"`: the processor sizes with floor division and centres with `w // 2 - src_w // 2` where MiniMax-H3 rounds and centres with `(src_w - w) // 2`, which differs by a pixel on 106 of 218 sampled aspect ratios and would move the conditioning latents off the reference implementation. The stretched anchor is pixel-identical either way and does go through the processor. Co-Authored-By: Claude Opus 5 (1M context) --- .../minimax_h3/before_encoder.py | 146 ++++++------------ .../minimax_h3/modular_blocks_minimax_h3.py | 72 +++++++-- .../test_modular_pipeline_minimax_h3.py | 2 +- 3 files changed, 110 insertions(+), 110 deletions(-) diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py b/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py index 88978bb6ff77..f32e521002cd 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py @@ -16,9 +16,11 @@ import torch from PIL import Image, ImageOps +from ...configuration_utils import FrozenDict +from ...image_processor import VaeImageProcessor from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState -from ..modular_pipeline_utils import InputParam, OutputParam +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline from .packing import ( MINIMAX_H3_CANVAS_MULTIPLE, @@ -26,10 +28,7 @@ MINIMAX_H3_MAX_DURATION, MINIMAX_H3_MIN_DURATION, align_num_frames, - audio_latent_num_frames, - prepare_keyframe_image, resolve_canvas_size, - video_latent_num_frames, ) from .packing_ref2va import ( MINIMAX_H3_MAX_REFERENCE_AUDIOS, @@ -51,57 +50,27 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def _latent_geometry(components, height: int, width: int, num_frames: int) -> tuple[int, int, int, int]: - r"""The latent geometry the packed layout, the noise draws and the decoders all key off.""" - ratio = components.vae_spatial_compression_ratio - return video_latent_num_frames(num_frames), height // ratio, width // ratio, audio_latent_num_frames(num_frames) - - -def _latent_geometry_outputs() -> list[OutputParam]: - r"""The declaration of what [`_latent_geometry`] resolves, shared by the two setup blocks.""" - return [ - OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."), - OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."), - OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."), - OutputParam("num_audio_latents", type_hint=int, description="Number of generated audio latents per channel."), - ] - - -class MiniMaxH3SetupStep(ModularPipelineBlocks): +class MiniMaxH3ResizeStep(ModularPipelineBlocks): model_name = "minimax-h3" @property def description(self) -> str: return ( - "Resolves the plan shared by the `t2va` and `fl2va` tasks: the canvas (MiniMax-H3's own 768-short-edge " - "geometry for the aspect ratio of the first keyframe, or 16:9 without keyframes), the `17 * n + 5` frame " - "count the video VAE can decode, the latent geometry every later block keys off, and the keyframes put " - "onto that canvas." + "Puts the `fl2va` keyframes onto the target canvas — MiniMax-H3's own 768-short-edge geometry for the " + "aspect ratio of the first keyframe unless `height` and `width` say otherwise. The canvas resolved here " + "is the one the whole request generates at." ) - @staticmethod - def _check_inputs(block_state) -> None: - if (block_state.height is None) != (block_state.width is None): - raise ValueError("`height` and `width` have to be passed together, or neither of them.") - if block_state.height is not None and ( - block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE - ): - raise ValueError( - f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got " - f"{block_state.height}x{block_state.width}." - ) - # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has - # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds. - aligned_num_frames = align_num_frames(block_state.num_frames) - duration = aligned_num_frames / MINIMAX_H3_FPS - if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: - raise ValueError( - f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at " - f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " - f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and " - f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames} (rounded up to " - f"{aligned_num_frames})." - ) + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] @property def inputs(self) -> list[InputParam]: @@ -124,15 +93,6 @@ def inputs(self) -> list[InputParam]: ), InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."), InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."), - InputParam( - name="num_frames", - type_hint=int, - default=124, - description=( - "Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video " - "VAE can decode; the resulting duration must stay between 5 and 15 seconds." - ), - ), ] @property @@ -140,57 +100,61 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."), OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."), - OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), - *_latent_geometry_outputs(), OutputParam( "keyframes", type_hint=list, - description="The keyframes put onto the target canvas, in packed order (empty for `t2va`).", + description="The keyframes put onto the target canvas, in packed order.", ), OutputParam( "keyframe_anchors", type_hint=tuple, - description="Which end of the video every keyframe is anchored to, in packed order.", + description=( + "Which end of the video every keyframe is anchored to, in packed order. Positional with " + "`keyframes`, so both are resolved here." + ), ), ] @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - self._check_inputs(block_state) - keyframes = [ - ImageOps.exif_transpose(keyframe).convert("RGB") - for keyframe in (block_state.image, block_state.last_image) - if keyframe is not None - ] + keyframes = [keyframe for keyframe in (block_state.image, block_state.last_image) if keyframe is not None] block_state.keyframe_anchors = tuple( anchor for anchor, keyframe in (("first", block_state.image), ("last", block_state.last_image)) if keyframe is not None ) if block_state.height is None: - block_state.height, block_state.width = resolve_canvas_size(*(keyframes[0].size if keyframes else (16, 9))) + block_state.height, block_state.width = resolve_canvas_size(*keyframes[0].size) + + prepared = [] + for index, keyframe in enumerate(keyframes): + if keyframe.size == (block_state.width, block_state.height): + prepared.append(keyframe) + elif index == 0: + # The geometry anchor is stretched onto the canvas. `resize_mode="default"` is exactly PIL's + # `resize((width, height), LANCZOS)`, verified pixel-identical across aspect ratios. + prepared.append( + components.image_processor.resize(keyframe, height=block_state.height, width=block_state.width) + ) + else: + # The follower is cover-cropped. `VaeImageProcessor`'s `resize_mode="crop"` is *not* a drop-in here: + # it sizes with floor division and centres with `w // 2 - src_w // 2`, where MiniMax-H3 rounds and + # centres with `(src_w - w) // 2`. The two agree on some aspect ratios and differ by a pixel on + # others (106 of 218 sampled), which would move the conditioning latents off the reference + # implementation, so the released model's arithmetic is kept. + scale = max(block_state.width / keyframe.size[0], block_state.height / keyframe.size[1]) + resized_size = ( + max(block_state.width, round(keyframe.size[0] * scale)), + max(block_state.height, round(keyframe.size[1] * scale)), + ) + left = max(0, (resized_size[0] - block_state.width) // 2) + top = max(0, (resized_size[1] - block_state.height) // 2) + resized = keyframe.resize(resized_size, Image.Resampling.LANCZOS) + prepared.append(resized.crop((left, top, left + block_state.width, top + block_state.height))) + block_state.keyframes = prepared - aligned_num_frames = align_num_frames(block_state.num_frames) - if aligned_num_frames != block_state.num_frames: - logger.warning( - f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {block_state.num_frames} " - f"up to {aligned_num_frames}." - ) - block_state.num_frames = aligned_num_frames - - ( - block_state.num_latent_frames, - block_state.latent_height, - block_state.latent_width, - block_state.num_audio_latents, - ) = _latent_geometry(components, block_state.height, block_state.width, block_state.num_frames) - - block_state.keyframes = [ - prepare_keyframe_image(keyframe, block_state.height, block_state.width, stretch=index == 0) - for index, keyframe in enumerate(keyframes) - ] self.set_block_state(state, block_state) return components, state @@ -288,7 +252,6 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."), OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."), OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), - *_latent_geometry_outputs(), OutputParam( "prepared_references", type_hint=list[MiniMaxH3PreparedReference], @@ -397,12 +360,5 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt f"to {block_state.num_frames}." ) - ( - block_state.num_latent_frames, - block_state.latent_height, - block_state.latent_width, - block_state.num_audio_latents, - ) = _latent_geometry(components, block_state.height, block_state.width, block_state.num_frames) - self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py b/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py index b1a7b4f69724..90c2c8621fec 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py +++ b/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py @@ -22,7 +22,7 @@ MiniMaxH3Ref2VAPrepareLayoutStep, MiniMaxH3SetTimestepsStep, ) -from .before_encoder import MiniMaxH3Ref2VASetupStep, MiniMaxH3SetupStep +from .before_encoder import MiniMaxH3Ref2VASetupStep, MiniMaxH3ResizeStep from .decoders import MiniMaxH3AudioDecodeStep, MiniMaxH3VideoDecodeStep from .denoise import MiniMaxH3DenoiseStep, MiniMaxH3Ref2VADenoiseStep from .encoders import ( @@ -46,6 +46,57 @@ def _generation_outputs() -> list[OutputParam]: ] +# auto_docstring +class MiniMaxH3AutoResizeStep(ConditionalPipelineBlocks): + """ + Keyframe canvas block. + - MiniMaxH3ResizeStep runs for the `fl2va` task, whichever of the two keyframes is given. + - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped and the layout step falls + back to MiniMax-H3's own 16:9 canvas. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image`, *optional*): + Keyframe the video starts from. + last_image (`Image`, *optional*): + Keyframe the video ends on. + height (`int`, *optional*): + Height of the generated video in pixels, a multiple of 32. + width (`int`, *optional*): + Width of the generated video in pixels, a multiple of 32. + + Outputs: + height (`int`), width (`int`): + The resolved canvas. + keyframes (`list`): + The keyframes put onto that canvas, in packed order. + keyframe_anchors (`tuple`): + Which end of the video every keyframe is anchored to, in packed order. + """ + + model_name = "minimax-h3" + block_classes = [MiniMaxH3ResizeStep] + block_names = ["keyframes"] + block_trigger_inputs = ["image", "last_image"] + default_block_name = None + + def select_block(self, **kwargs) -> str | None: + if kwargs.get("image") is not None or kwargs.get("last_image") is not None: + return "keyframes" + return None + + @property + def description(self): + return ( + "Keyframe canvas block.\n" + + " - MiniMaxH3ResizeStep runs for the `fl2va` task, whichever of the two keyframes is given.\n" + + " - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped and the layout " + "step falls back to MiniMax-H3's own 16:9 canvas." + ) + + # auto_docstring class MiniMaxH3AutoKeyframeVaeEncoderStep(ConditionalPipelineBlocks): """ @@ -54,22 +105,15 @@ class MiniMaxH3AutoKeyframeVaeEncoderStep(ConditionalPipelineBlocks): - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped. Components: - vae (`AutoencoderKLMiniMaxH3`) scheduler (`MiniMaxH3Scheduler`) + vae (`AutoencoderKLMiniMaxH3`) Inputs: - keyframes (`list`, *optional*): + keyframes (`list`): The keyframes put onto the target canvas, in packed order. - latent_height (`int`, *optional*): - Height of the video latents. - latent_width (`int`, *optional*): - Width of the video latents. - generator (`Generator`, *optional*): - The generator of the request. The conditioning noise is drawn from it before the target noise of the - prepare-latents step. Outputs: - condition_latents (`Tensor`): - The noise-augmented video conditioning rows, in packed order. + condition_latents (`list`): + The normalized video conditioning latents, one tensor per keyframe, in packed order. """ model_name = "minimax-h3" @@ -205,7 +249,7 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): model_name = "minimax-h3" block_classes = [ - MiniMaxH3SetupStep, + MiniMaxH3AutoResizeStep, MiniMaxH3TextEncoderStep, MiniMaxH3AutoKeyframeVaeEncoderStep, MiniMaxH3PrepareLayoutStep, @@ -215,7 +259,7 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): MiniMaxH3DecodeStep, ] block_names = [ - "setup", + "resize", "text_encoder", "vae_encoder", "prepare_layout", diff --git a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py index 6e1bf9f4236a..9a785ff125ca 100644 --- a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py +++ b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py @@ -52,7 +52,7 @@ T2VA_WORKFLOW = [ - ("setup", "MiniMaxH3SetupStep"), + ("resize", "MiniMaxH3AutoResizeStep"), ("text_encoder", "MiniMaxH3TextEncoderStep"), ("prepare_layout", "MiniMaxH3PrepareLayoutStep"), ("prepare_latents", "MiniMaxH3PrepareLatentsStep"), From 2b89de7657355f8e6a420a253e5f6f10ee8ecff4 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 2 Aug 2026 17:18:03 +0000 Subject: [PATCH 04/16] Resolve the MiniMax-H3 geometry where the layout is built MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The frame alignment and the latent shapes were resolved in the setup step, ahead of the encoders, though the layout step is their first reader. They move to `prepare_layout` for both blocksets, along with the canvas fallback and the checks that guard them — the multiple-of-32 canvas and the 5-to-15 second ceiling now sit in the block that computes what they protect. `prepare_latents` picks up the other half: it draws the conditioning noise the encoders used to draw, mixes it in and packs every condition, so all three noise streams of a request are drawn in one place, in the order a seed reproduces. Each condition is packed on its own because `ref2va` references are encoded at their own resolutions and do not share a shape. `_layout_inputs`, `_layout_outputs`, `_set_layout_state`, `_latent_geometry` and `_latent_geometry_outputs` are inlined into the two blocks that used them, and `keyframe_condition_noise` and `prepare_keyframe_image` are gone with their last callers. Co-Authored-By: Claude Opus 5 (1M context) --- .../minimax_h3/before_denoise.py | 405 +++++++++++------- .../modular_pipelines/minimax_h3/packing.py | 73 ---- 2 files changed, 252 insertions(+), 226 deletions(-) diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py index ef874b1e0100..c0a1a0811dea 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -22,11 +22,19 @@ from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline from .packing import ( MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_CANVAS_MULTIPLE, + MINIMAX_H3_FPS, MINIMAX_H3_KEYFRAME_NOISE_AUG, + MINIMAX_H3_MAX_DURATION, + MINIMAX_H3_MIN_DURATION, MiniMaxH3PackedSequence, + align_num_frames, + audio_latent_num_frames, build_packed_sequence, build_row_timesteps, patchify_video_latents, + resolve_canvas_size, + video_latent_num_frames, ) from .packing_ref2va import MiniMaxH3PreparedReference, build_ref2va_packed_sequence @@ -34,93 +42,39 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def _layout_inputs() -> list[InputParam]: - r"""What both packed layouts are built from, beyond the conditioning of the task itself.""" - return [ - InputParam( - name="text_token_tags", - type_hint=torch.Tensor, - required=True, - description="The per-row modality tag of every row of `prompt_embeds`.", - ), - InputParam( - name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames." - ), - InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), - InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), - InputParam( - name="num_audio_latents", - type_hint=int, - required=True, - description="Number of audio latents per channel.", - ), - ] - - -def _layout_outputs() -> list[OutputParam]: - r"""The row layout of the packed sequence, shared by the two tasks.""" - return [ - OutputParam( - "layout", - type_hint=MiniMaxH3PackedSequence, - description="The structural description of the packed sequence.", - ), - OutputParam( - "position_ids", - type_hint=torch.Tensor, - description="The `(t, h, w)` rotary coordinate of every row, in float64.", - ), - OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."), - OutputParam( - "video_indices", - type_hint=torch.Tensor, - description="Sequence positions of the video rows, conditioning rows first.", - ), - OutputParam( - "audio_indices", - type_hint=torch.Tensor, - description="Sequence positions of the audio rows, reference rows first.", - ), - OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."), - OutputParam( - "num_condition_video_rows", - type_hint=int, - description="How many leading video rows are conditioning rows rather than generated rows.", - ), - OutputParam( - "num_condition_audio_rows", - type_hint=int, - description="How many leading audio rows are reference rows rather than generated rows.", - ), - ] - - -def _set_layout_state(block_state, layout: MiniMaxH3PackedSequence, device: torch.device) -> None: - block_state.layout = layout - block_state.position_ids = layout.position_ids.to(device) - block_state.token_tags = layout.token_tags.to(device) - block_state.video_indices = layout.video_indices.to(device) - block_state.audio_indices = layout.audio_indices.to(device) - block_state.text_indices = layout.text_indices.to(device) - block_state.num_condition_video_rows = layout.num_condition_video_rows - block_state.num_condition_audio_rows = layout.num_condition_audio_rows - - class MiniMaxH3PrepareLayoutStep(ModularPipelineBlocks): model_name = "minimax-h3" @property def description(self) -> str: return ( - "Builds the packed layout of a `t2va` / `fl2va` request — `[text | keyframe conditions | target audio | " - "target video]` — and its fp64 rotary grid. MiniMax-H3 runs full self-attention over this one sequence, " - "so the layout is what every later block addresses rows through." + "Resolves the geometry of a `t2va` / `fl2va` request — the canvas, the `17 * n + 5` frame count the video " + "VAE can decode and the latent shapes every later block keys off — and builds the packed layout from it: " + "`[text | keyframe conditions | target audio | target video]` plus its fp64 rotary grid. MiniMax-H3 runs " + "full self-attention over this one sequence, so the layout is what every later block addresses rows " + "through." ) @property def inputs(self) -> list[InputParam]: return [ - *_layout_inputs(), + InputParam( + name="text_token_tags", + type_hint=torch.Tensor, + required=True, + description="The per-row modality tag of every row of `prompt_embeds`.", + ), + InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."), + InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."), + InputParam( + name="num_frames", + type_hint=int, + default=124, + description=( + "Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video " + "VAE can decode; the resulting duration must stay between 5 and 15 seconds." + ), + ), InputParam( name="keyframe_anchors", type_hint=tuple, @@ -131,11 +85,87 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return _layout_outputs() + return [ + OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."), + OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."), + OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), + OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."), + OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."), + OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."), + OutputParam( + "num_audio_latents", type_hint=int, description="Number of generated audio latents per channel." + ), + OutputParam( + "layout", + type_hint=MiniMaxH3PackedSequence, + description="The structural description of the packed sequence.", + ), + OutputParam( + "position_ids", + type_hint=torch.Tensor, + description="The `(t, h, w)` rotary coordinate of every row, in float64.", + ), + OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."), + OutputParam( + "video_indices", + type_hint=torch.Tensor, + description="Sequence positions of the video rows, conditioning rows first.", + ), + OutputParam( + "audio_indices", + type_hint=torch.Tensor, + description="Sequence positions of the audio rows, reference rows first.", + ), + OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."), + OutputParam( + "num_condition_video_rows", + type_hint=int, + description="How many leading video rows are conditioning rows rather than generated rows.", + ), + OutputParam( + "num_condition_audio_rows", + type_hint=int, + description="How many leading audio rows are reference rows rather than generated rows.", + ), + ] @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + device = components._execution_device + + # Without a keyframe to take the aspect ratio from, MiniMax-H3 generates on its own 16:9 canvas. + if block_state.height is None: + block_state.height, block_state.width = resolve_canvas_size(16, 9) + if block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE: + raise ValueError( + f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got " + f"{block_state.height}x{block_state.width}." + ) + + aligned_num_frames = align_num_frames(block_state.num_frames) + if aligned_num_frames != block_state.num_frames: + logger.warning( + f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {block_state.num_frames} " + f"up to {aligned_num_frames}." + ) + block_state.num_frames = aligned_num_frames + # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has + # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds. + duration = block_state.num_frames / MINIMAX_H3_FPS + if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + raise ValueError( + f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at " + f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " + f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and " + f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames}." + ) + + ratio = components.vae_spatial_compression_ratio + block_state.num_latent_frames = video_latent_num_frames(block_state.num_frames) + block_state.latent_height = block_state.height // ratio + block_state.latent_width = block_state.width // ratio + block_state.num_audio_latents = audio_latent_num_frames(block_state.num_frames) layout = build_packed_sequence( block_state.text_token_tags, @@ -146,7 +176,14 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - components.patch_size, block_state.keyframe_anchors, ) - _set_layout_state(block_state, layout, components._execution_device) + block_state.layout = layout + block_state.position_ids = layout.position_ids.to(device) + block_state.token_tags = layout.token_tags.to(device) + block_state.video_indices = layout.video_indices.to(device) + block_state.audio_indices = layout.audio_indices.to(device) + block_state.text_indices = layout.text_indices.to(device) + block_state.num_condition_video_rows = layout.num_condition_video_rows + block_state.num_condition_audio_rows = layout.num_condition_audio_rows self.set_block_state(state, block_state) return components, state @@ -158,30 +195,94 @@ class MiniMaxH3Ref2VAPrepareLayoutStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Builds the packed layout of a `ref2va` request — `[text | reference blocks | target audio | target " - "video]` — and its fp64 rotary grid. The reference order advances the shared audio/video rotary clock, so " - "it is part of the layout rather than a detail of the presentation." + "Resolves the latent shapes of a `ref2va` request and builds its packed layout — `[text | reference " + "blocks | target audio | target video]` — plus its fp64 rotary grid. The reference order advances the " + "shared audio/video rotary clock, so it is part of the layout rather than a detail of the presentation." ) @property def inputs(self) -> list[InputParam]: return [ - *_layout_inputs(), + InputParam( + name="text_token_tags", + type_hint=torch.Tensor, + required=True, + description="The per-row modality tag of every row of `prompt_embeds`.", + ), InputParam( name="prepared_references", type_hint=list[MiniMaxH3PreparedReference], required=True, description="The prepared references, in packed order, with their latent geometry filled in.", ), + InputParam.template("height", required=True, description="Height of the generated video in pixels."), + InputParam.template("width", required=True, description="Width of the generated video in pixels."), + InputParam( + name="num_frames", + type_hint=int, + required=True, + description="Resolved number of frames, of the form 17 * n + 5.", + ), ] @property def intermediate_outputs(self) -> list[OutputParam]: - return _layout_outputs() + return [ + OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."), + OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."), + OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), + OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."), + OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."), + OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."), + OutputParam( + "num_audio_latents", type_hint=int, description="Number of generated audio latents per channel." + ), + OutputParam( + "layout", + type_hint=MiniMaxH3PackedSequence, + description="The structural description of the packed sequence.", + ), + OutputParam( + "position_ids", + type_hint=torch.Tensor, + description="The `(t, h, w)` rotary coordinate of every row, in float64.", + ), + OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."), + OutputParam( + "video_indices", + type_hint=torch.Tensor, + description="Sequence positions of the video rows, conditioning rows first.", + ), + OutputParam( + "audio_indices", + type_hint=torch.Tensor, + description="Sequence positions of the audio rows, reference rows first.", + ), + OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."), + OutputParam( + "num_condition_video_rows", + type_hint=int, + description="How many leading video rows are conditioning rows rather than generated rows.", + ), + OutputParam( + "num_condition_audio_rows", + type_hint=int, + description="How many leading audio rows are reference rows rather than generated rows.", + ), + ] @torch.no_grad() def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + device = components._execution_device + + # The canvas and the frame count are settled by the setup step: a `ref2va` soundtrack is truncated to the + # generated duration as the references are prepared, so `num_frames` has to be final before that runs. + ratio = components.vae_spatial_compression_ratio + block_state.num_latent_frames = video_latent_num_frames(block_state.num_frames) + block_state.latent_height = block_state.height // ratio + block_state.latent_width = block_state.width // ratio + block_state.num_audio_latents = audio_latent_num_frames(block_state.num_frames) layout = build_ref2va_packed_sequence( block_state.text_token_tags, @@ -192,7 +293,14 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt block_state.num_audio_latents, components.patch_size, ) - _set_layout_state(block_state, layout, components._execution_device) + block_state.layout = layout + block_state.position_ids = layout.position_ids.to(device) + block_state.token_tags = layout.token_tags.to(device) + block_state.video_indices = layout.video_indices.to(device) + block_state.audio_indices = layout.audio_indices.to(device) + block_state.text_indices = layout.text_indices.to(device) + block_state.num_condition_video_rows = layout.num_condition_video_rows + block_state.num_condition_audio_rows = layout.num_condition_audio_rows self.set_block_state(state, block_state) return components, state @@ -204,11 +312,15 @@ class MiniMaxH3PrepareLatentsStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Draws the initial noise of the generated rows and prepends the conditioning rows. MiniMax-H3 draws the " - "video noise as a latent tensor and patchifies it afterwards, then the audio noise directly in row " - "layout — both off the request's generator, after the conditioning noise of the encoder step." + "Draws every noise stream of the request and packs the video rows. MiniMax-H3 draws one condition at a " + "time first — noising the encoded anchors to its conditioning level — then the video noise as a latent " + "tensor, then the audio noise directly in row layout, all off the request's generator, in that order." ) + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", MiniMaxH3Scheduler)] + @property def inputs(self) -> list[InputParam]: return [ @@ -244,8 +356,12 @@ def inputs(self) -> list[InputParam]: ), InputParam( name="condition_latents", - type_hint=torch.Tensor, - description="The video conditioning rows to prepend, or None for a request that has none.", + type_hint=list[torch.Tensor], + description=( + "The encoded video conditioning latents, one `(1, latent_channels, num_latent_frames, " + "latent_height, latent_width)` tensor per condition in packed order, or None for a request that " + "has none. Noised and packed here." + ), ), InputParam( name="audio_condition_latents", @@ -269,83 +385,66 @@ def intermediate_outputs(self) -> list[OutputParam]: ), ] - @staticmethod - def prepare_latents( - components, - num_latent_frames: int, - latent_height: int, - latent_width: int, - num_audio_latents: int, - device: torch.device, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Draw the initial noise of both modalities and pack it into transformer rows. - - A request draws every stream from the one generator it is given, and the order is part of what that generator - reproduces: the conditioning noise of the keyframes or references first (one draw per condition, in - [`~modular_pipelines.minimax_h3.packing.keyframe_condition_noise`]), then the video noise here, as a latent tensor - that is patchified afterwards, then the audio noise, directly in row layout. Passing `latents` or - `audio_latents` skips its draw and shifts the ones after it. - - Args: - num_latent_frames (`int`): Number of video latent frames. - latent_height (`int`): Latent height. - latent_width (`int`): Latent width. - num_audio_latents (`int`): Number of audio latents per channel. - device (`torch.device`): The device the rows are drawn on. - generator (`torch.Generator`, *optional*): The generator of the request. - latents (`torch.Tensor`, *optional*): - Pre-generated video noise of shape `(1, latent_channels, num_latent_frames, latent_height, - latent_width)`, used instead of the draw. - audio_latents (`torch.Tensor`, *optional*): - Pre-generated audio noise of shape `(2, audio_latent_channels, num_audio_latents)`. - - Returns: - `tuple[torch.Tensor, torch.Tensor]`: the video rows and the channel-major audio rows. - """ + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + patch_size = components.patch_size + + # A request draws every stream from the one generator it is given, and the order is part of what that + # generator reproduces: one draw per condition first, then the video noise as a latent tensor, then the audio + # noise directly in row layout. Passing `latents` or `audio_latents` skips its draw and shifts the ones after + # it. + condition_rows = None + if block_state.condition_latents is not None: + # One draw per condition, in packed order. Each is packed on its own because `ref2va` references are + # encoded at their own resolutions, so their latents do not share a shape. + packed = [] + for condition in block_state.condition_latents: + noise = randn_tensor( + condition.shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + # The anchors are not fully clean: the released model noises them to `t = 0.999` and holds them there + # for every step. Mixing before the patchify is the same arithmetic, since patchify only permutes. + noised = components.scheduler.scale_noise(condition.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise) + packed.append(patchify_video_latents(noised, patch_size)) + condition_rows = torch.cat(packed) + + latents = block_state.latents if latents is None: latents = randn_tensor( - (1, components.vae_latent_channels, num_latent_frames, latent_height, latent_width), - generator=generator, + ( + 1, + components.vae_latent_channels, + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + ), + generator=block_state.generator, device=device, dtype=torch.float32, ) - video_rows = patchify_video_latents(latents.to(torch.float32), components.patch_size) + video_rows = patchify_video_latents(latents.to(device, torch.float32), patch_size) - if audio_latents is None: + if block_state.audio_latents is None: audio_rows = randn_tensor( - (num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, components.audio_latent_channels), - generator=generator, + (block_state.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, components.audio_latent_channels), + generator=block_state.generator, device=device, dtype=torch.float32, ) else: - audio_rows = audio_latents.to(torch.float32).permute(0, 2, 1).reshape(-1, components.audio_latent_channels) - return video_rows.to(device), audio_rows.to(device) - - @torch.no_grad() - def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) + audio_rows = ( + block_state.audio_latents.to(device, torch.float32) + .permute(0, 2, 1) + .reshape(-1, components.audio_latent_channels) + ) - latents, audio_latents = self.prepare_latents( - components, - block_state.num_latent_frames, - block_state.latent_height, - block_state.latent_width, - block_state.num_audio_latents, - components._execution_device, - block_state.generator, - block_state.latents, - block_state.audio_latents, - ) - if block_state.condition_latents is not None: - latents = torch.cat([block_state.condition_latents, latents]) + if condition_rows is not None: + video_rows = torch.cat([condition_rows, video_rows]) if block_state.audio_condition_latents is not None: - audio_latents = torch.cat([block_state.audio_condition_latents, audio_latents]) - block_state.latents, block_state.audio_latents = latents, audio_latents + audio_rows = torch.cat([block_state.audio_condition_latents.to(device), audio_rows]) + block_state.latents, block_state.audio_latents = video_rows, audio_rows self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/minimax_h3/packing.py b/src/diffusers/modular_pipelines/minimax_h3/packing.py index c84f7e7f9aa8..94e14226c778 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/packing.py +++ b/src/diffusers/modular_pipelines/minimax_h3/packing.py @@ -40,9 +40,6 @@ import numpy as np import torch -from PIL import Image - -from ...utils.torch_utils import randn_tensor # Per-row modality tags. They index the transformer's AdaLN table, so the values are a checkpoint contract. @@ -213,36 +210,6 @@ def audio_latent_num_frames(num_frames: int) -> int: return int(round(num_frames / MINIMAX_H3_FPS * MINIMAX_H3_AUDIO_LATENTS_PER_SECOND)) -def prepare_keyframe_image(image, height: int, width: int, stretch: bool): - r""" - Put a keyframe onto the target canvas. - - The first keyframe of a request is the geometry anchor and is *stretched* onto the canvas, while a second - keyframe follows that canvas and is cover-cropped (aspect-preserving max-scale LANCZOS resize plus a centre - crop). An image that already is the canvas is returned untouched, without a resampling pass. - - Args: - image (`PIL.Image.Image`): The keyframe, in RGB and already EXIF-transposed. - height (`int`): Canvas height. - width (`int`): Canvas width. - stretch (`bool`): Whether to stretch (geometry anchor) instead of cover-cropping (follower). - - Returns: - `PIL.Image.Image`: The prepared keyframe. - """ - if image.size == (width, height): - return image - if stretch: - return image.resize((width, height), Image.Resampling.LANCZOS) - - scale = max(width / image.size[0], height / image.size[1]) - resized_size = (max(width, round(image.size[0] * scale)), max(height, round(image.size[1] * scale))) - left = max(0, (resized_size[0] - width) // 2) - top = max(0, (resized_size[1] - height) // 2) - resized = image.resize(resized_size, Image.Resampling.LANCZOS) - return resized.crop((left, top, left + width, top + height)) - - def patchify_video_latents(latents: torch.Tensor, patch_size: tuple[int, int, int]) -> torch.Tensor: r""" Pack video latents into transformer rows. @@ -496,43 +463,3 @@ def build_row_timesteps( row_timesteps[layout.audio_indices[layout.num_condition_audio_rows :]] = audio_timestep row_timesteps[layout.audio_indices[: layout.num_condition_audio_rows]] = condition_audio_timestep return torch.unique(row_timesteps, sorted=True, return_inverse=True) - - -def keyframe_condition_noise( - condition_latent_shapes: tuple[tuple[int, int, int], ...], - patch_size: tuple[int, int, int], - latent_channels: int, - generator: torch.Generator | list[torch.Generator] | None = None, - device: torch.device | None = None, - dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - r""" - Draw the noise that the keyframe (or reference) conditioning rows are mixed with. - - One draw per condition, in packed order, off the request's generator. The conditioning rows are prepared before - the target rows, so these are the *first* draws of a request, ahead of the video and audio noise of - [`~MiniMaxH3PrepareLatentsStep.prepare_latents`] — the order is part of what a generator reproduces. - - Args: - condition_latent_shapes (`tuple[tuple[int, int, int], ...]`): - The `(num_latent_frames, latent_height, latent_width)` of every condition, in packed order. - patch_size (`tuple[int, int, int]`): The transformer's `(t, h, w)` patch. - latent_channels (`int`): Number of video latent channels. - generator (`torch.Generator`, *optional*): The generator of the request. - device (`torch.device`, *optional*): The device the noise is drawn on. - dtype (`torch.dtype`, defaults to `torch.float32`): The dtype of the noise. - - Returns: - `torch.Tensor` of shape `(num_condition_rows, latent_channels * prod(patch_size))`: the noise rows, - concatenated in packed order. - """ - rows = [] - for num_latent_frames, latent_height, latent_width in condition_latent_shapes: - noise = randn_tensor( - (1, latent_channels, num_latent_frames, latent_height, latent_width), - generator=generator, - device=device, - dtype=dtype, - ) - rows.append(patchify_video_latents(noise, patch_size)) - return torch.cat(rows) From 3f001c650c76b874c4ea18cfcb27d4c3e581c988 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 2 Aug 2026 18:53:15 +0000 Subject: [PATCH 05/16] Inline the MiniMax-H3 audio VAE's weight-norm conv helper `_wn_conv1d` was a one-line alias for `weight_norm(nn.Conv1d(...))` behind ten call sites, so reading any of them meant a detour. Co-Authored-By: Claude Opus 5 (1M context) --- .../autoencoder_kl_minimax_h3_audio.py | 37 ++++++++++--------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py index c35c62e28ec3..0b24768642bb 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py @@ -93,10 +93,6 @@ class MiniMaxH3AudioEncoderOutput(BaseOutput): latent_dist: MiniMaxH3AudioDiagonalGaussianDistribution -def _wn_conv1d(*args, **kwargs) -> nn.Module: - return weight_norm(nn.Conv1d(*args, **kwargs)) - - def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: r"""Kaiser-windowed sinc low-pass filter of shape `[1, 1, kernel_size]`. @@ -234,9 +230,9 @@ def __init__(self, dim: int, dilation: int): super().__init__() self.block = nn.Sequential( MiniMaxH3AudioSnake1d(dim), - _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=((7 - 1) * dilation) // 2), + weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=((7 - 1) * dilation) // 2)), MiniMaxH3AudioSnake1d(dim), - _wn_conv1d(dim, dim, kernel_size=1), + weight_norm(nn.Conv1d(dim, dim, kernel_size=1)), ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -257,12 +253,14 @@ def __init__(self, dim: int, stride: int): MiniMaxH3AudioResidualUnit(dim // 2, dilation=3), MiniMaxH3AudioResidualUnit(dim // 2, dilation=9), MiniMaxH3AudioSnake1d(dim // 2), - _wn_conv1d( - dim // 2, - dim, - kernel_size=2 * stride, - stride=stride, - padding=math.ceil(stride / 2), + weight_norm( + nn.Conv1d( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) ), ) @@ -275,13 +273,13 @@ class MiniMaxH3AudioEncoder(nn.Module): def __init__(self, d_model: int, strides: tuple[int, ...], d_latent: int): super().__init__() - block: list[nn.Module] = [_wn_conv1d(1, d_model, kernel_size=7, padding=3)] + block: list[nn.Module] = [weight_norm(nn.Conv1d(1, d_model, kernel_size=7, padding=3))] for stride in strides: d_model *= 2 block.append(MiniMaxH3AudioEncoderBlock(d_model, stride=stride)) block += [ MiniMaxH3AudioSnake1d(d_model), - _wn_conv1d(d_model, d_latent, kernel_size=3, padding=1), + weight_norm(nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1)), ] self.block = nn.Sequential(*block) @@ -403,12 +401,15 @@ def __init__(self, channels: int, kernel_size: int, dilation: tuple[int, ...]): super().__init__() self.convs1 = nn.ModuleList( [ - _wn_conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2) + weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2)) for d in dilation ] ) self.convs2 = nn.ModuleList( - [_wn_conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2) for _ in dilation] + [ + weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2)) + for _ in dilation + ] ) self.activations = nn.ModuleList( [ @@ -442,7 +443,7 @@ def __init__( self.num_kernels = len(resblock_kernel_sizes) self.num_upsamples = len(upsample_rates) - self.conv_pre = _wn_conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3) + self.conv_pre = weight_norm(nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)) # Each upsampler is wrapped in a one-element `ModuleList` in the original checkpoint # (`ups..0`); the extra nesting is kept so the state dict stays a passthrough. @@ -471,7 +472,7 @@ def __init__( self.resblocks.append(MiniMaxH3AudioAMPBlock(channels, kernel, tuple(dilation))) self.activation_post = MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels)) - self.conv_post = _wn_conv1d(channels, 1, 7, 1, padding=3, bias=False) + self.conv_post = weight_norm(nn.Conv1d(channels, 1, 7, 1, padding=3, bias=False)) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv_pre(hidden_states) From d6050f4217aeebf543aaa402c982d6350eb2c44e Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 2 Aug 2026 18:53:33 +0000 Subject: [PATCH 06/16] Give the MiniMax-H3 blocks their latents back MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A second pass over the blockset, all of it structural — the generated video and soundtrack are unchanged. The decoders used to take packed rows plus five geometry scalars and unpack them on the way in, so they could not accept latents from anywhere else and an injected step between denoise and decode would have had to work in row space. An after-denoise step now drops the conditioning rows and unpacks both modalities, following `Flux2UnpackLatentsStep` and `QwenImageAfterDenoiseStep`, leaving the decoders to denormalize and decode. `output_type="latent"` goes with it: popping the decoder is how a modular pipeline hands back latents, and the branch had the declared outputs lying about their own types. `MiniMaxH3PackedSequence` is gone. It was the only internal dataclass in the blockset and it made the layout step publish its index tensors twice, once as CPU copies inside the dataclass and once as device copies on the block state. The builders return a tuple and `set_timesteps` declares the four values it actually reads instead of taking the whole layout. The denoiser stops enumerating the rows of the packed sequence. The layout tags `token_tags`, `position_ids` and the three index tensors as `denoiser_input_fields` — their names are already the transformer's own — so the denoise block declares one input and forwards what the signature accepts. `encode_prompt` takes the conditioner it needs rather than the whole component bag, and the modality tags it writes are arguments with defaults. The `17` and `5` of the chunking were module constants duplicating the video VAE's `clip_length` and `tokens_chunk_size`; they are read off the component now, so a checkpoint that chunks differently cannot silently disagree with the packing. Every remaining helper with one or two callers is inlined into them. Two notes left in place: a TODO on firing the conditioner's offload hook by hand, which needs a real answer rather than a workaround, and a check that the conditioning rows the layout reserved match the conditioning that was encoded — in a hand-assembled chain the canvas is user input, and the mismatch used to surface as an `index_copy` shape error inside the transformer. Co-Authored-By: Claude Opus 5 (1M context) --- .../minimax_h3/before_denoise.py | 154 +++++++++---- .../minimax_h3/before_encoder.py | 14 +- .../modular_pipelines/minimax_h3/decoders.py | 195 ++++++++++------ .../modular_pipelines/minimax_h3/denoise.py | 213 ++++++++++-------- .../modular_pipelines/minimax_h3/encoders.py | 183 ++++++++------- .../minimax_h3/modular_blocks_minimax_h3.py | 152 +++++++------ .../minimax_h3/modular_pipeline.py | 12 + .../modular_pipelines/minimax_h3/packing.py | 151 +++---------- .../minimax_h3/packing_ref2va.py | 33 ++- 9 files changed, 613 insertions(+), 494 deletions(-) diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py index c0a1a0811dea..cecb04598e37 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -27,7 +27,6 @@ MINIMAX_H3_KEYFRAME_NOISE_AUG, MINIMAX_H3_MAX_DURATION, MINIMAX_H3_MIN_DURATION, - MiniMaxH3PackedSequence, align_num_frames, audio_latent_num_frames, build_packed_sequence, @@ -95,28 +94,36 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "num_audio_latents", type_hint=int, description="Number of generated audio latents per channel." ), - OutputParam( - "layout", - type_hint=MiniMaxH3PackedSequence, - description="The structural description of the packed sequence.", - ), OutputParam( "position_ids", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="The `(t, h, w)` rotary coordinate of every row, in float64.", ), - OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."), + OutputParam( + "token_tags", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="The modality tag of every row.", + ), OutputParam( "video_indices", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="Sequence positions of the video rows, conditioning rows first.", ), OutputParam( "audio_indices", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="Sequence positions of the audio rows, reference rows first.", ), - OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."), + OutputParam( + "text_indices", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Sequence positions of the text rows.", + ), OutputParam( "num_condition_video_rows", type_hint=int, @@ -143,7 +150,9 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - f"{block_state.height}x{block_state.width}." ) - aligned_num_frames = align_num_frames(block_state.num_frames) + frames_per_chunk = components.vae_frames_per_chunk + latents_per_chunk = components.vae_latents_per_chunk + aligned_num_frames = align_num_frames(block_state.num_frames, frames_per_chunk, latents_per_chunk) if aligned_num_frames != block_state.num_frames: logger.warning( f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {block_state.num_frames} " @@ -162,12 +171,22 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - ) ratio = components.vae_spatial_compression_ratio - block_state.num_latent_frames = video_latent_num_frames(block_state.num_frames) + block_state.num_latent_frames = video_latent_num_frames( + block_state.num_frames, frames_per_chunk, latents_per_chunk + ) block_state.latent_height = block_state.height // ratio block_state.latent_width = block_state.width // ratio block_state.num_audio_latents = audio_latent_num_frames(block_state.num_frames) - layout = build_packed_sequence( + ( + position_ids, + token_tags, + video_indices, + audio_indices, + text_indices, + block_state.num_condition_video_rows, + block_state.num_condition_audio_rows, + ) = build_packed_sequence( block_state.text_token_tags, block_state.num_latent_frames, block_state.latent_height, @@ -176,14 +195,11 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - components.patch_size, block_state.keyframe_anchors, ) - block_state.layout = layout - block_state.position_ids = layout.position_ids.to(device) - block_state.token_tags = layout.token_tags.to(device) - block_state.video_indices = layout.video_indices.to(device) - block_state.audio_indices = layout.audio_indices.to(device) - block_state.text_indices = layout.text_indices.to(device) - block_state.num_condition_video_rows = layout.num_condition_video_rows - block_state.num_condition_audio_rows = layout.num_condition_audio_rows + block_state.position_ids = position_ids.to(device) + block_state.token_tags = token_tags.to(device) + block_state.video_indices = video_indices.to(device) + block_state.audio_indices = audio_indices.to(device) + block_state.text_indices = text_indices.to(device) self.set_block_state(state, block_state) return components, state @@ -237,28 +253,36 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "num_audio_latents", type_hint=int, description="Number of generated audio latents per channel." ), - OutputParam( - "layout", - type_hint=MiniMaxH3PackedSequence, - description="The structural description of the packed sequence.", - ), OutputParam( "position_ids", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="The `(t, h, w)` rotary coordinate of every row, in float64.", ), - OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."), + OutputParam( + "token_tags", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="The modality tag of every row.", + ), OutputParam( "video_indices", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="Sequence positions of the video rows, conditioning rows first.", ), OutputParam( "audio_indices", type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", description="Sequence positions of the audio rows, reference rows first.", ), - OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."), + OutputParam( + "text_indices", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Sequence positions of the text rows.", + ), OutputParam( "num_condition_video_rows", type_hint=int, @@ -279,12 +303,22 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt # The canvas and the frame count are settled by the setup step: a `ref2va` soundtrack is truncated to the # generated duration as the references are prepared, so `num_frames` has to be final before that runs. ratio = components.vae_spatial_compression_ratio - block_state.num_latent_frames = video_latent_num_frames(block_state.num_frames) + block_state.num_latent_frames = video_latent_num_frames( + block_state.num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk + ) block_state.latent_height = block_state.height // ratio block_state.latent_width = block_state.width // ratio block_state.num_audio_latents = audio_latent_num_frames(block_state.num_frames) - layout = build_ref2va_packed_sequence( + ( + position_ids, + token_tags, + video_indices, + audio_indices, + text_indices, + block_state.num_condition_video_rows, + block_state.num_condition_audio_rows, + ) = build_ref2va_packed_sequence( block_state.text_token_tags, block_state.prepared_references, block_state.num_latent_frames, @@ -293,14 +327,11 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt block_state.num_audio_latents, components.patch_size, ) - block_state.layout = layout - block_state.position_ids = layout.position_ids.to(device) - block_state.token_tags = layout.token_tags.to(device) - block_state.video_indices = layout.video_indices.to(device) - block_state.audio_indices = layout.audio_indices.to(device) - block_state.text_indices = layout.text_indices.to(device) - block_state.num_condition_video_rows = layout.num_condition_video_rows - block_state.num_condition_audio_rows = layout.num_condition_audio_rows + block_state.position_ids = position_ids.to(device) + block_state.token_tags = token_tags.to(device) + block_state.video_indices = video_indices.to(device) + block_state.audio_indices = audio_indices.to(device) + block_state.text_indices = text_indices.to(device) self.set_block_state(state, block_state) return components, state @@ -354,6 +385,12 @@ def inputs(self) -> list[InputParam]: type_hint=torch.Tensor, description="Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.", ), + InputParam( + name="num_condition_video_rows", + type_hint=int, + default=0, + description="How many conditioning rows the layout reserved, which the packed conditioning must match.", + ), InputParam( name="condition_latents", type_hint=list[torch.Tensor], @@ -409,6 +446,15 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - noised = components.scheduler.scale_noise(condition.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise) packed.append(patchify_video_latents(noised, patch_size)) condition_rows = torch.cat(packed) + # In a hand-assembled chain the canvas reaching the layout is user input, so it can disagree with the + # keyframes that were actually encoded. Left alone the mismatch first surfaces as an `index_copy` shape + # error inside the transformer, 50 layers deep. + if condition_rows.shape[0] != block_state.num_condition_video_rows: + raise ValueError( + f"The layout reserved {block_state.num_condition_video_rows} conditioning rows but the encoded " + f"conditioning latents pack into {condition_rows.shape[0]}. The canvas the layout was built from " + "and the one the conditioning was encoded at do not agree." + ) latents = block_state.latents if latents is None: @@ -474,10 +520,34 @@ def inputs(self) -> list[InputParam]: return [ InputParam.template("num_inference_steps", required=True), InputParam( - name="layout", - type_hint=MiniMaxH3PackedSequence, + name="video_indices", + type_hint=torch.Tensor, + required=True, + description="Sequence positions of the video rows, conditioning rows first.", + ), + InputParam( + name="audio_indices", + type_hint=torch.Tensor, required=True, - description="The structural description of the packed sequence.", + description="Sequence positions of the audio rows, reference rows first.", + ), + InputParam( + name="text_indices", + type_hint=torch.Tensor, + required=True, + description="Sequence positions of the text rows.", + ), + InputParam( + name="num_condition_video_rows", + type_hint=int, + default=0, + description="How many leading video rows are conditioning rows.", + ), + InputParam( + name="num_condition_audio_rows", + type_hint=int, + default=0, + description="How many leading audio rows are reference rows.", ), ] @@ -510,7 +580,11 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - tuple( tensor.to(device) for tensor in build_row_timesteps( - block_state.layout, + block_state.video_indices, + block_state.audio_indices, + block_state.num_condition_video_rows, + block_state.num_condition_audio_rows, + block_state.text_indices.numel(), float(timestep), float(audio_timestep), max(float(timestep), MINIMAX_H3_KEYFRAME_NOISE_AUG), diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py b/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py index f32e521002cd..86c045be683e 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py @@ -183,7 +183,13 @@ def _check_inputs(components, block_state) -> None: ) # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds. - aligned_num_frames = None if block_state.num_frames is None else align_num_frames(block_state.num_frames) + aligned_num_frames = ( + None + if block_state.num_frames is None + else align_num_frames( + block_state.num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk + ) + ) duration = None if aligned_num_frames is None else aligned_num_frames / MINIMAX_H3_FPS if duration is not None and not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: raise ValueError( @@ -310,7 +316,9 @@ def prepare_references( f"`references[{index}]` is {duration:g} seconds long, outside the " f"{MINIMAX_H3_MIN_DURATION} to {MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates." ) - num_frames = align_num_frames(round(duration * MINIMAX_H3_FPS)) + num_frames = align_num_frames( + round(duration * MINIMAX_H3_FPS), components.vae_frames_per_chunk, components.vae_latents_per_chunk + ) # The duration the request generates is the one of the *aligned* frame count, so that is what the # ceiling has to hold for: a 14.99 second soundtrack rounds up to 362 frames, i.e. 15.083 seconds. if num_frames / MINIMAX_H3_FPS > MINIMAX_H3_MAX_DURATION: @@ -320,7 +328,7 @@ def prepare_references( f"{MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates. Pass `num_frames` to generate a " "shorter video from this soundtrack." ) - num_frames = align_num_frames(num_frames) + num_frames = align_num_frames(num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk) for reference, entry in zip(resolved, references): if reference.kind == "image": diff --git a/src/diffusers/modular_pipelines/minimax_h3/decoders.py b/src/diffusers/modular_pipelines/minimax_h3/decoders.py index fc4ee359267b..2eec15d15638 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/decoders.py +++ b/src/diffusers/modular_pipelines/minimax_h3/decoders.py @@ -21,27 +21,123 @@ from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline -from .packing import ( - MINIMAX_H3_PIXEL_MEAN, - MINIMAX_H3_PIXEL_STD, - unpack_audio_tokens, - unpatchify_video_tokens, -) +from .packing import MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD logger = logging.get_logger(__name__) # pylint: disable=invalid-name +class MiniMaxH3AfterDenoiseStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Turns the denoised rows of the packed sequence back into latents: drops the conditioning rows the loop " + "never wrote, then unpacks the video rows into `(1, latent_channels, num_latent_frames, latent_height, " + "latent_width)` and the channel-major audio rows into the `(2, audio_latent_channels, num_audio_latents)` " + "the mono audio VAE consumes. The decoders then take latents from any source, and popping them leaves " + "latents in hand." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="The denoised video rows of the packed sequence, conditioning rows first.", + ), + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + required=True, + description="The denoised audio rows of the packed sequence, reference rows first.", + ), + InputParam( + name="num_condition_video_rows", + type_hint=int, + default=0, + description="How many leading video rows are conditioning rows and are dropped here.", + ), + InputParam( + name="num_condition_audio_rows", + type_hint=int, + default=0, + description="How many leading audio rows are reference rows and are dropped here.", + ), + InputParam( + name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames." + ), + InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), + InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), + InputParam( + name="num_audio_latents", + type_hint=int, + required=True, + description="Number of audio latents per channel.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The generated video latents, of shape `(1, latent_channels, num_latent_frames, " + "latent_height, latent_width)`.", + ), + OutputParam( + "audio_latents", + type_hint=torch.Tensor, + description="The generated audio latents, one batch item per stereo channel.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + patch_t, patch_h, patch_w = components.patch_size + channels = components.vae_latent_channels + + # The inverse of the patchify in the prepare-latents step: rows are frame-major then row-major. + rows = block_state.latents[block_state.num_condition_video_rows :] + rows = rows.reshape( + -1, + block_state.num_latent_frames // patch_t, + block_state.latent_height // patch_h, + block_state.latent_width // patch_w, + channels, + patch_t, + patch_h, + patch_w, + ) + rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7) + block_state.latents = rows.reshape( + -1, channels, block_state.num_latent_frames, block_state.latent_height, block_state.latent_width + ).contiguous() + + # Audio rows are channel-major, and the mono audio VAE takes the two stereo channels as two batch items. + audio_rows = block_state.audio_latents[block_state.num_condition_audio_rows :] + audio_rows = audio_rows.reshape(MINIMAX_H3_AUDIO_CHANNELS, block_state.num_audio_latents, audio_rows.shape[-1]) + block_state.audio_latents = audio_rows.permute(0, 2, 1).contiguous() + + self.set_block_state(state, block_state) + return components, state + + class MiniMaxH3VideoDecodeStep(ModularPipelineBlocks): model_name = "minimax-h3" @property def description(self) -> str: return ( - "Unpacks the generated video rows back into latents, denormalizes them and decodes them into video. The " - "spatial tiling of the video VAE covers the canvas exactly, so the decoded frames need no crop back, but " - "the decode itself runs under float16 autocast even though the VAE weights are float32, and the VAE " - "produces ImageNet-normalized RGB that is reverted here." + "Denormalizes the generated video latents and decodes them into video. The spatial tiling of the video " + "VAE covers the canvas exactly, so the decoded frames need no crop back, but the decode itself runs under " + "float16 autocast even though the VAE weights are float32, and the VAE produces ImageNet-normalized RGB " + "that is reverted here." ) @property @@ -65,22 +161,9 @@ def inputs(self) -> list[InputParam]: name="latents", type_hint=torch.Tensor, required=True, - description="The denoised video rows of the packed sequence, conditioning rows first.", - ), - InputParam( - name="num_condition_video_rows", - type_hint=int, - default=0, - description="How many leading video rows are conditioning rows and are dropped here.", - ), - InputParam( - name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames." - ), - InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), - InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), - InputParam.template( - "output_type", description="Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents." + description="The generated video latents.", ), + InputParam.template("output_type", description="Output format: 'pil', 'np' or 'pt'."), ] @property @@ -92,29 +175,16 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - block_state = self.get_block_state(state) device = components._execution_device - latents = unpatchify_video_tokens( - block_state.latents[block_state.num_condition_video_rows :], - block_state.num_latent_frames, - block_state.latent_height, - block_state.latent_width, - components.vae_latent_channels, - components.patch_size, - ) latents_mean = torch.tensor(components.vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1) latents_std = torch.tensor(components.vae.config.latents_std, device=device).view(1, -1, 1, 1, 1) - latents = latents * latents_std + latents_mean - - if block_state.output_type == "latent": - block_state.videos = latents - else: - with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"): - video = components.vae.decode(latents, return_dict=False)[0] - pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) - pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) - video = (video.float() * pixel_std + pixel_mean).clamp(0, 1) - block_state.videos = components.video_processor.postprocess_video( - video, output_type=block_state.output_type - ) + latents = block_state.latents * latents_std + latents_mean + + with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"): + video = components.vae.decode(latents, return_dict=False)[0] + pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) + video = (video.float() * pixel_std + pixel_mean).clamp(0, 1) + block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) self.set_block_state(state, block_state) return components, state @@ -126,8 +196,8 @@ class MiniMaxH3AudioDecodeStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Unpacks the generated audio rows back into latents, denormalizes them and decodes them into a stereo " - "waveform. The audio VAE is mono and takes the two stereo channels as two batch items." + "Denormalizes the generated audio latents and decodes them into a stereo waveform. The audio VAE is mono " + "and takes the two stereo channels as two batch items." ) @property @@ -141,22 +211,7 @@ def inputs(self) -> list[InputParam]: name="audio_latents", type_hint=torch.Tensor, required=True, - description="The denoised audio rows of the packed sequence, reference rows first.", - ), - InputParam( - name="num_condition_audio_rows", - type_hint=int, - default=0, - description="How many leading audio rows are reference rows and are dropped here.", - ), - InputParam( - name="num_audio_latents", - type_hint=int, - required=True, - description="Number of audio latents per channel.", - ), - InputParam.template( - "output_type", description="Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents." + description="The generated audio latents, one batch item per stereo channel.", ), ] @@ -180,18 +235,12 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - block_state = self.get_block_state(state) device = components._execution_device - audio_latents = unpack_audio_tokens( - block_state.audio_latents[block_state.num_condition_audio_rows :], block_state.num_audio_latents - ) audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean, device=device).view(1, -1, 1) audio_latents_std = torch.tensor(components.audio_vae.config.latents_std, device=device).view(1, -1, 1) - audio_latents = audio_latents * audio_latents_std + audio_latents_mean + audio_latents = block_state.audio_latents * audio_latents_std + audio_latents_mean - if block_state.output_type == "latent": - block_state.audio = audio_latents - else: - audio = components.audio_vae.decode(audio_latents, return_dict=False)[0] - block_state.audio = audio.float().permute(1, 0, 2) + audio = components.audio_vae.decode(audio_latents, return_dict=False)[0] + block_state.audio = audio.float().permute(1, 0, 2) block_state.sampling_rate = components.audio_sampling_rate self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/minimax_h3/denoise.py b/src/diffusers/modular_pipelines/minimax_h3/denoise.py index d149273371bb..156242907339 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/denoise.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect + import torch from ...models import MiniMaxH3Transformer3DModel @@ -30,91 +32,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def _denoiser_inputs() -> list[InputParam]: - r"""Everything one MiniMax-H3 forward reads, beyond the transformer itself.""" - return [ - InputParam( - name="latents", - type_hint=torch.Tensor, - required=True, - description="The video rows of the packed sequence, conditioning rows first.", - ), - InputParam( - name="audio_latents", - type_hint=torch.Tensor, - required=True, - description="The channel-major audio rows of the packed sequence, reference rows first.", - ), - InputParam.template("prompt_embeds"), - InputParam( - name="row_timestep_plan", - type_hint=list, - required=True, - description="One `(timestep, timestep_indices)` pair per step.", - ), - InputParam( - name="token_tags", type_hint=torch.Tensor, required=True, description="The modality tag of every row." - ), - InputParam( - name="position_ids", - type_hint=torch.Tensor, - required=True, - description="The `(t, h, w)` rotary coordinate of every row.", - ), - InputParam( - name="video_indices", - type_hint=torch.Tensor, - required=True, - description="Sequence positions of the video rows.", - ), - InputParam( - name="audio_indices", - type_hint=torch.Tensor, - required=True, - description="Sequence positions of the audio rows.", - ), - InputParam( - name="text_indices", - type_hint=torch.Tensor, - required=True, - description="Sequence positions of the text rows.", - ), - InputParam.template("attention_kwargs"), - ] - - -def _denoiser_outputs() -> list[OutputParam]: - return [ - OutputParam( - "noise_pred", type_hint=torch.Tensor, description="Predicted velocity of the video rows of the sequence." - ), - OutputParam( - "audio_noise_pred", - type_hint=torch.Tensor, - description="Predicted velocity of the audio rows of the sequence.", - ), - ] - - -def _predict_velocity(transformer: MiniMaxH3Transformer3DModel, block_state: BlockState, i: int): - r"""One MiniMax-H3 forward pass: every row of the packed sequence, at its own noise level, at once.""" - unique_timesteps, timestep_indices = block_state.row_timestep_plan[i] - return transformer( - hidden_states=block_state.latents[None], - audio_hidden_states=block_state.audio_latents[None], - encoder_hidden_states=block_state.prompt_embeds, - timestep=unique_timesteps, - timestep_indices=timestep_indices, - token_tags=block_state.token_tags, - position_ids=block_state.position_ids, - video_indices=block_state.video_indices, - audio_indices=block_state.audio_indices, - text_indices=block_state.text_indices, - attention_kwargs=block_state.attention_kwargs, - return_dict=False, - ) - - class MiniMaxH3LoopDenoiser(ModularPipelineBlocks): model_name = "minimax-h3" @@ -132,16 +49,71 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: - return _denoiser_inputs() + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="The video rows of the packed sequence, conditioning rows first.", + ), + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + required=True, + description="The channel-major audio rows of the packed sequence, reference rows first.", + ), + InputParam.template("prompt_embeds"), + InputParam( + name="row_timestep_plan", + type_hint=list, + required=True, + description="One `(timestep, timestep_indices)` pair per step.", + ), + InputParam( + kwargs_type="denoiser_input_fields", + description=( + "The structural description of the packed sequence the transformer reads by name: `token_tags`, " + "`position_ids` and the three row-index tensors." + ), + ), + InputParam.template("attention_kwargs"), + ] @property def intermediate_outputs(self) -> list[OutputParam]: - return _denoiser_outputs() + return [ + OutputParam( + "noise_pred", + type_hint=torch.Tensor, + description="Predicted velocity of the video rows of the sequence.", + ), + OutputParam( + "audio_noise_pred", + type_hint=torch.Tensor, + description="Predicted velocity of the audio rows of the sequence.", + ), + ] @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): - block_state.noise_pred, block_state.audio_noise_pred = _predict_velocity( - components.transformer, block_state, i + transformer = components.transformer + unique_timesteps, timestep_indices = block_state.row_timestep_plan[i] + # The layout tags its outputs `denoiser_input_fields`, and their names are the transformer's own parameter + # names, so the rows of the packed sequence are described to it without this block enumerating them. + layout_kwargs = { + name: value + for name, value in block_state.denoiser_input_fields.items() + if name in inspect.signature(transformer.forward).parameters + } + block_state.noise_pred, block_state.audio_noise_pred = transformer( + hidden_states=block_state.latents[None], + audio_hidden_states=block_state.audio_latents[None], + encoder_hidden_states=block_state.prompt_embeds, + timestep=unique_timesteps, + timestep_indices=timestep_indices, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **layout_kwargs, ) return components, block_state @@ -162,16 +134,71 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: - return _denoiser_inputs() + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="The video rows of the packed sequence, conditioning rows first.", + ), + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + required=True, + description="The channel-major audio rows of the packed sequence, reference rows first.", + ), + InputParam.template("prompt_embeds"), + InputParam( + name="row_timestep_plan", + type_hint=list, + required=True, + description="One `(timestep, timestep_indices)` pair per step.", + ), + InputParam( + kwargs_type="denoiser_input_fields", + description=( + "The structural description of the packed sequence the transformer reads by name: `token_tags`, " + "`position_ids` and the three row-index tensors." + ), + ), + InputParam.template("attention_kwargs"), + ] @property def intermediate_outputs(self) -> list[OutputParam]: - return _denoiser_outputs() + return [ + OutputParam( + "noise_pred", + type_hint=torch.Tensor, + description="Predicted velocity of the video rows of the sequence.", + ), + OutputParam( + "audio_noise_pred", + type_hint=torch.Tensor, + description="Predicted velocity of the audio rows of the sequence.", + ), + ] @torch.no_grad() def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): - block_state.noise_pred, block_state.audio_noise_pred = _predict_velocity( - components.transformer_ref, block_state, i + transformer = components.transformer_ref + unique_timesteps, timestep_indices = block_state.row_timestep_plan[i] + # The layout tags its outputs `denoiser_input_fields`, and their names are the transformer's own parameter + # names, so the rows of the packed sequence are described to it without this block enumerating them. + layout_kwargs = { + name: value + for name, value in block_state.denoiser_input_fields.items() + if name in inspect.signature(transformer.forward).parameters + } + block_state.noise_pred, block_state.audio_noise_pred = transformer( + hidden_states=block_state.latents[None], + audio_hidden_states=block_state.audio_latents[None], + encoder_hidden_states=block_state.prompt_embeds, + timestep=unique_timesteps, + timestep_indices=timestep_indices, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **layout_kwargs, ) return components, block_state diff --git a/src/diffusers/modular_pipelines/minimax_h3/encoders.py b/src/diffusers/modular_pipelines/minimax_h3/encoders.py index 26e3e2b55e5c..e70a40fde61b 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/encoders.py +++ b/src/diffusers/modular_pipelines/minimax_h3/encoders.py @@ -40,40 +40,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def _check_prompt(prompt) -> None: - r"""MiniMax-H3 packs one request into one sequence, so a batch of prompts is not a thing.""" - if not isinstance(prompt, str): - raise ValueError( - f"MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got {type(prompt)}." - ) - - -def _conditioner_components() -> list[ComponentSpec]: - r"""MiniMax-H3's conditioner: a Qwen3-VL read at its 50th decoder layer, with its language-model head unused.""" - return [ - ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration), - ComponentSpec("tokenizer", Qwen2TokenizerFast), - ComponentSpec("processor", Qwen3VLProcessor), - ] - - -def _conditioner_outputs() -> list[OutputParam]: - return [ - OutputParam.template( - "prompt_embeds", - description=( - "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the " - "50th decoder layer of the Qwen3-VL conditioner." - ), - ), - OutputParam( - "text_token_tags", - type_hint=torch.Tensor, - description="The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video.", - ), - ] - - class MiniMaxH3TextEncoderStep(ModularPipelineBlocks): model_name = "minimax-h3" @@ -87,7 +53,11 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: - return _conditioner_components() + return [ + ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ComponentSpec("processor", Qwen3VLProcessor), + ] @property def inputs(self) -> list[InputParam]: @@ -102,15 +72,34 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return _conditioner_outputs() + return [ + OutputParam.template( + "prompt_embeds", + description=( + "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the " + "50th decoder layer of the Qwen3-VL conditioner." + ), + ), + OutputParam( + "text_token_tags", + type_hint=torch.Tensor, + description=( + "The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video." + ), + ), + ] @staticmethod def encode_prompt( - components, + text_encoder, + tokenizer, + processor, prompt: str, images: list | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, + text_tag: int = MINIMAX_H3_TEXT_TAG, + video_tag: int = MINIMAX_H3_VIDEO_TAG, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Build MiniMax-H3's presentation of a request and encode it. @@ -131,10 +120,8 @@ def encode_prompt( `tuple[torch.Tensor, torch.Tensor]`: the `(1, num_text_tokens, 5120)` hidden states and the `(num_text_tokens,)` per-row modality tags. """ - device = device or components._execution_device - dtype = dtype or components.transformer.dtype - num_layers = components.text_encoder.config.text_config.num_hidden_layers + num_layers = text_encoder.config.text_config.num_hidden_layers if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: raise ValueError( f"MiniMax-H3 conditions on `hidden_states[{MINIMAX_H3_TEXT_ENCODER_LAYER}]` of its Qwen3-VL " @@ -146,42 +133,46 @@ def encode_prompt( pixel_values, image_grid_thw = None, None token_ids, token_tags = [], [] if images: - vision = components.processor.image_processor(images=images, return_tensors="pt") + vision = processor.image_processor(images=images, return_tensors="pt") pixel_values, image_grid_thw = vision["pixel_values"], vision["image_grid_thw"] - merge_size = components.processor.image_processor.merge_size**2 + merge_size = processor.image_processor.merge_size**2 for index in range(len(images)): num_image_tokens = int(image_grid_thw[index].prod()) // merge_size - label_ids = components.tokenizer(f": ", add_special_tokens=False)["input_ids"] + label_ids = tokenizer(f": ", add_special_tokens=False)["input_ids"] vision_ids = ( - [components.tokenizer.convert_tokens_to_ids("<|vision_start|>")] - + [components.tokenizer.convert_tokens_to_ids("<|image_pad|>")] * num_image_tokens - + [components.tokenizer.convert_tokens_to_ids("<|vision_end|>")] + [tokenizer.convert_tokens_to_ids("<|vision_start|>")] + + [tokenizer.convert_tokens_to_ids("<|image_pad|>")] * num_image_tokens + + [tokenizer.convert_tokens_to_ids("<|vision_end|>")] ) token_ids += label_ids + vision_ids - token_tags += [MINIMAX_H3_TEXT_TAG] * len(label_ids) + [MINIMAX_H3_VIDEO_TAG] * len(vision_ids) - prompt_ids = components.tokenizer(prompt, add_special_tokens=False)["input_ids"] + token_tags += [text_tag] * len(label_ids) + [video_tag] * len(vision_ids) + prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] token_ids += prompt_ids - token_tags += [MINIMAX_H3_TEXT_TAG] * len(prompt_ids) + token_tags += [text_tag] * len(prompt_ids) input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) # Qwen3-VL lays its 3D rotary positions out per modality run, which it reads off the token type ids the # processor derives from the vision pad ids (`0` text, `1` image, `2` video). mm_token_type_ids = torch.tensor( - components.processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device + processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device ) # `text_encoder.model` is a submodule, and a CPU-offload hook — accelerate's or the one the # `ComponentsManager` attaches — wraps the *top-level* module's `forward` alone, so calling the submodule # directly would leave the conditioner on the CPU. Fire the hook by hand instead of routing through # `text_encoder(...)`: MiniMax-H3 reads `hidden_states[50]` and never uses the language-model head, whose # vocabulary-wide projection over every token is all the top-level forward would add. - hook = getattr(components.text_encoder, "_hf_hook", None) + # TODO: firing another module's offload hook by hand is not something a block should do. It is here + # because MiniMax-H3 reads `hidden_states[50]` off `text_encoder.model` while the hook wraps only the + # top-level `forward`. Needs a real answer — an opt-in on the conditioner, or a hook that follows + # submodule calls. + hook = getattr(text_encoder, "_hf_hook", None) if hook is not None and hasattr(hook, "pre_forward"): - hook.pre_forward(components.text_encoder) - outputs = components.text_encoder.model( + hook.pre_forward(text_encoder) + outputs = text_encoder.model( input_ids=input_ids, attention_mask=torch.ones_like(input_ids), mm_token_type_ids=mm_token_type_ids, - pixel_values=None if pixel_values is None else pixel_values.to(device, components.text_encoder.dtype), + pixel_values=None if pixel_values is None else pixel_values.to(device, text_encoder.dtype), image_grid_thw=None if image_grid_thw is None else image_grid_thw.to(device), use_cache=False, output_hidden_states=True, @@ -192,12 +183,18 @@ def encode_prompt( @torch.no_grad() def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - _check_prompt(block_state.prompt) + if not isinstance(block_state.prompt, str): + raise ValueError( + "MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got " + f"{type(block_state.prompt)}." + ) # `encode_prompt` defaults the embedding dtype to the denoiser's; a text encoder block has no denoiser of # its own — it is meant to run on its own — so it emits the conditioner's dtype, as every other model does. block_state.prompt_embeds, block_state.text_token_tags = self.encode_prompt( - components, + components.text_encoder, + components.tokenizer, + components.processor, block_state.prompt, block_state.keyframes, device=components._execution_device, @@ -310,7 +307,11 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: - return _conditioner_components() + return [ + ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ComponentSpec("processor", Qwen3VLProcessor), + ] @property def inputs(self) -> list[InputParam]: @@ -326,11 +327,28 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return _conditioner_outputs() + return [ + OutputParam.template( + "prompt_embeds", + description=( + "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the " + "50th decoder layer of the Qwen3-VL conditioner." + ), + ), + OutputParam( + "text_token_tags", + type_hint=torch.Tensor, + description=( + "The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video." + ), + ), + ] @staticmethod def encode_prompt( - components, + text_encoder, + tokenizer, + processor, prompt: str, references: list[MiniMaxH3PreparedReference], device: torch.device | None = None, @@ -358,10 +376,8 @@ def encode_prompt( `tuple[torch.Tensor, torch.Tensor]`: the `(1, num_text_tokens, 5120)` hidden states and the `(num_text_tokens,)` per-row modality tags. """ - device = device or components._execution_device - dtype = dtype or components.transformer_ref.dtype - num_layers = components.text_encoder.config.text_config.num_hidden_layers + num_layers = text_encoder.config.text_config.num_hidden_layers if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: raise ValueError( f"MiniMax-H3 conditions on `hidden_states[{MINIMAX_H3_TEXT_ENCODER_LAYER}]` of its Qwen3-VL " @@ -370,11 +386,11 @@ def encode_prompt( f"{MINIMAX_H3_TEXT_ENCODER_LAYER} layers is post-norm and is not the conditioning MiniMax-H3 expects." ) - merge_size = components.processor.image_processor.merge_size**2 + merge_size = processor.image_processor.merge_size**2 pixel_values, image_grid_thw, image_token_counts = None, None, [] images = [reference.image for reference in references if reference.kind == "image"] if images: - vision = components.processor.image_processor(images=images, return_tensors="pt") + vision = processor.image_processor(images=images, return_tensors="pt") pixel_values, image_grid_thw = vision["pixel_values"], vision["image_grid_thw"] image_token_counts = [int(grid.prod()) // merge_size for grid in image_grid_thw] @@ -384,7 +400,7 @@ def encode_prompt( sampled = [sample_reference_video_frames(reference.frames) for reference in videos] for reference, (_, block_timestamps) in zip(videos, sampled): reference.block_timestamps = block_timestamps - vision = components.processor.video_processor( + vision = processor.video_processor( videos=[np.stack(frames) for frames, _ in sampled], do_sample_frames=False, return_tensors="pt" ) pixel_values_videos, video_grid_thw = vision["pixel_values_videos"], vision["video_grid_thw"] @@ -397,30 +413,34 @@ def encode_prompt( ) token_ids, token_tags = build_ref2va_presentation( - components.tokenizer, prompt, references, image_token_counts, video_block_token_counts + tokenizer, prompt, references, image_token_counts, video_block_token_counts ) input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) # Qwen3-VL lays its 3D rotary positions out per modality run, which it reads off the token type ids the # processor derives from the vision pad ids (`0` text, `1` image, `2` video). mm_token_type_ids = torch.tensor( - components.processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device + processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device ) # `text_encoder.model` is a submodule, and a CPU-offload hook — accelerate's or the one the # `ComponentsManager` attaches — wraps the *top-level* module's `forward` alone, so calling the submodule # directly would leave the conditioner on the CPU. Fire the hook by hand instead of routing through # `text_encoder(...)`: MiniMax-H3 reads `hidden_states[50]` and never uses the language-model head, whose # vocabulary-wide projection over every token is all the top-level forward would add. - hook = getattr(components.text_encoder, "_hf_hook", None) + # TODO: firing another module's offload hook by hand is not something a block should do. It is here + # because MiniMax-H3 reads `hidden_states[50]` off `text_encoder.model` while the hook wraps only the + # top-level `forward`. Needs a real answer — an opt-in on the conditioner, or a hook that follows + # submodule calls. + hook = getattr(text_encoder, "_hf_hook", None) if hook is not None and hasattr(hook, "pre_forward"): - hook.pre_forward(components.text_encoder) - outputs = components.text_encoder.model( + hook.pre_forward(text_encoder) + outputs = text_encoder.model( input_ids=input_ids, attention_mask=torch.ones_like(input_ids), mm_token_type_ids=mm_token_type_ids, - pixel_values=None if pixel_values is None else pixel_values.to(device, components.text_encoder.dtype), + pixel_values=None if pixel_values is None else pixel_values.to(device, text_encoder.dtype), image_grid_thw=None if image_grid_thw is None else image_grid_thw.to(device), pixel_values_videos=( - None if pixel_values_videos is None else pixel_values_videos.to(device, components.text_encoder.dtype) + None if pixel_values_videos is None else pixel_values_videos.to(device, text_encoder.dtype) ), video_grid_thw=None if video_grid_thw is None else video_grid_thw.to(device), use_cache=False, @@ -432,12 +452,18 @@ def encode_prompt( @torch.no_grad() def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - _check_prompt(block_state.prompt) + if not isinstance(block_state.prompt, str): + raise ValueError( + "MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got " + f"{type(block_state.prompt)}." + ) # `encode_prompt` defaults the embedding dtype to the denoiser's; a text encoder block has no denoiser of # its own — it is meant to run on its own — so it emits the conditioner's dtype, as every other model does. block_state.prompt_embeds, block_state.text_token_tags = self.encode_prompt( - components, + components.text_encoder, + components.tokenizer, + components.processor, block_state.prompt, block_state.prepared_references, device=components._execution_device, @@ -524,7 +550,6 @@ def encode_references( video rows and the `(num_condition_audio_rows, audio_latent_channels)` audio rows, both float32 on CPU and both `None` when the references carry no such rows. """ - device = device or components._execution_device latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1) latents_std = torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1) pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) @@ -538,7 +563,13 @@ def encode_references( if reference.kind == "image": pixels = torch.from_numpy(np.array(reference.image)).to(device).permute(2, 0, 1)[None, :, None] else: - frames = reference.frames[: trim_reference_num_frames(reference.frames.shape[0])] + frames = reference.frames[ + : trim_reference_num_frames( + reference.frames.shape[0], + components.vae_frames_per_chunk, + components.vae_latents_per_chunk, + ) + ] pixels = torch.from_numpy(frames.copy()).to(device).permute(3, 0, 1, 2)[None] pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std # A single frame is encoded by the (tiled) spatial encoder alone; a video goes through the temporal diff --git a/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py b/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py index 90c2c8621fec..bfe300c4f09e 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py +++ b/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py @@ -23,7 +23,7 @@ MiniMaxH3SetTimestepsStep, ) from .before_encoder import MiniMaxH3Ref2VASetupStep, MiniMaxH3ResizeStep -from .decoders import MiniMaxH3AudioDecodeStep, MiniMaxH3VideoDecodeStep +from .decoders import MiniMaxH3AfterDenoiseStep, MiniMaxH3AudioDecodeStep, MiniMaxH3VideoDecodeStep from .denoise import MiniMaxH3DenoiseStep, MiniMaxH3Ref2VADenoiseStep from .encoders import ( MiniMaxH3KeyframeVaeEncoderStep, @@ -51,29 +51,33 @@ class MiniMaxH3AutoResizeStep(ConditionalPipelineBlocks): """ Keyframe canvas block. - MiniMaxH3ResizeStep runs for the `fl2va` task, whichever of the two keyframes is given. - - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped and the layout step falls - back to MiniMax-H3's own 16:9 canvas. + - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped and the layout step falls back to MiniMax-H3's own 16:9 canvas. Components: image_processor (`VaeImageProcessor`) Inputs: image (`Image`, *optional*): - Keyframe the video starts from. + Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is derived from its own + aspect ratio. last_image (`Image`, *optional*): - Keyframe the video ends on. + Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with `image` it is the + follower of the two and is cover-cropped onto the canvas. height (`int`, *optional*): Height of the generated video in pixels, a multiple of 32. width (`int`, *optional*): Width of the generated video in pixels, a multiple of 32. Outputs: - height (`int`), width (`int`): - The resolved canvas. + height (`int`): + Resolved height of the generated video in pixels. + width (`int`): + Resolved width of the generated video in pixels. keyframes (`list`): - The keyframes put onto that canvas, in packed order. + The keyframes put onto the target canvas, in packed order. keyframe_anchors (`tuple`): - Which end of the video every keyframe is anchored to, in packed order. + Which end of the video every keyframe is anchored to, in packed order. Positional with `keyframes`, so both are + resolved here. """ model_name = "minimax-h3" @@ -108,12 +112,13 @@ class MiniMaxH3AutoKeyframeVaeEncoderStep(ConditionalPipelineBlocks): vae (`AutoencoderKLMiniMaxH3`) Inputs: - keyframes (`list`): + keyframes (`list`, *optional*): The keyframes put onto the target canvas, in packed order. Outputs: condition_latents (`list`): - The normalized video conditioning latents, one tensor per keyframe, in packed order. + The normalized video conditioning latents, one `(1, latent_channels, 1, latent_height, latent_width)` tensor per + keyframe, in packed order. """ model_name = "minimax-h3" @@ -142,27 +147,17 @@ class MiniMaxH3DecodeStep(SequentialPipelineBlocks): Decodes the denoised rows of the packed sequence into the generated video and its soundtrack. Components: - vae (`AutoencoderKLMiniMaxH3`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLMiniMaxH3Audio`) + vae (`AutoencoderKLMiniMaxH3`) + video_processor (`VideoProcessor`) + audio_vae (`AutoencoderKLMiniMaxH3Audio`) Inputs: latents (`Tensor`): - The denoised video rows of the packed sequence, conditioning rows first. - num_condition_video_rows (`int`, *optional*, defaults to 0): - How many leading video rows are conditioning rows and are dropped here. - num_latent_frames (`int`): - Number of video latent frames. - latent_height (`int`): - Height of the video latents. - latent_width (`int`): - Width of the video latents. + The generated video latents. output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents. + Output format: 'pil', 'np' or 'pt'. audio_latents (`Tensor`): - The denoised audio rows of the packed sequence, reference rows first. - num_condition_audio_rows (`int`, *optional*, defaults to 0): - How many leading audio rows are reference rows and are dropped here. - num_audio_latents (`int`): - Number of audio latents per channel. + The generated audio latents, one batch item per stereo channel. Outputs: videos (`list`): @@ -189,8 +184,7 @@ def outputs(self): # auto_docstring class MiniMaxH3Blocks(SequentialPipelineBlocks): """ - Modular pipeline blocks for joint video + audio generation with MiniMax-H3, covering the `t2va` (text only) and - `fl2va` (first and/or last keyframe) tasks of the FL2VA checkpoint. + Modular pipeline blocks for joint video + audio generation with MiniMax-H3, covering the `t2va` (text only) and `fl2va` (first and/or last keyframe) tasks of the FL2VA checkpoint. Supported workflows: - `t2va`: requires `prompt` @@ -198,45 +192,58 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): - `fl2va_last_frame`: requires `prompt`, `last_image` Components: - text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`Qwen2Tokenizer`) processor (`Qwen3VLProcessor`) - vae (`AutoencoderKLMiniMaxH3`) scheduler (`MiniMaxH3Scheduler`) audio_scheduler (`MiniMaxH3Scheduler`) - transformer (`MiniMaxH3Transformer3DModel`) video_processor (`VideoProcessor`) audio_vae - (`AutoencoderKLMiniMaxH3Audio`) + image_processor (`VaeImageProcessor`) + text_encoder (`Qwen3VLForConditionalGeneration`) + tokenizer (`Qwen2Tokenizer`) + processor (`Qwen3VLProcessor`) + vae (`AutoencoderKLMiniMaxH3`) + scheduler (`MiniMaxH3Scheduler`) + audio_scheduler (`MiniMaxH3Scheduler`) + transformer (`MiniMaxH3Transformer3DModel`) + video_processor (`VideoProcessor`) + audio_vae (`AutoencoderKLMiniMaxH3Audio`) Inputs: image (`Image`, *optional*): - Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is derived - from its own aspect ratio. + Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is derived from its own + aspect ratio. last_image (`Image`, *optional*): - Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with `image` - it is the follower of the two and is cover-cropped onto the canvas. + Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with `image` it is the + follower of the two and is cover-cropped onto the canvas. height (`int`, *optional*): Height of the generated video in pixels, a multiple of 32. width (`int`, *optional*): Width of the generated video in pixels, a multiple of 32. - num_frames (`int`, *optional*, defaults to 124): - Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can - decode; the resulting duration must stay between 5 and 15 seconds. prompt (`str`): The prompt to guide generation, a single string. + keyframes (`list`, *optional*): + The keyframes put onto the target canvas, in packed order (empty or None for `t2va`). + num_frames (`int`, *optional*, defaults to 124): + Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode; + the resulting duration must stay between 5 and 15 seconds. + keyframe_anchors (`tuple`, *optional*, defaults to ()): + Which end of the video every keyframe is anchored to, in packed order. generator (`Generator`, *optional*): - The generator of the request. The conditioning noise is drawn from it before the target noise of the - prepare-latents step. + The generator of the request. The video noise is drawn from it first, then the audio noise. latents (`Tensor`, *optional*): - Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used - instead of the draw. + Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the + draw. audio_latents (`Tensor`, *optional*): Pre-generated audio noise of shape `(2, 32, num_audio_latents)`. - condition_latents (`Tensor`, *optional*): - The video conditioning rows to prepend, or None for a request that has none. + condition_latents (`list`, *optional*): + The encoded video conditioning latents, one `(1, latent_channels, num_latent_frames, latent_height, latent_width)` + tensor per condition in packed order, or None for a request that has none. Noised and packed here. audio_condition_latents (`Tensor`, *optional*): The audio conditioning rows to prepend, or None for a request that has none. num_inference_steps (`int`): The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + The structural description of the packed sequence the transformer reads by name: `token_tags`, `position_ids` and + the three row-index tensors. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents. + Output format: 'pil', 'np' or 'pt'. Outputs: videos (`list`): @@ -256,6 +263,7 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): MiniMaxH3PrepareLatentsStep, MiniMaxH3SetTimestepsStep, MiniMaxH3DenoiseStep, + MiniMaxH3AfterDenoiseStep, MiniMaxH3DecodeStep, ] block_names = [ @@ -266,6 +274,7 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): "prepare_latents", "set_timesteps", "denoise", + "after_denoise", "decode", ] # One repository holds both checkpoint partitions, so the two blocksets are two workflows over one shared @@ -293,51 +302,56 @@ def outputs(self): # auto_docstring class MiniMaxH3Ref2VABlocks(SequentialPipelineBlocks): """ - Modular pipeline blocks for joint video + audio generation from an ordered list of image, video and audio - references with MiniMax-H3, the `ref2va` task of the Ref2VA checkpoint. + Modular pipeline blocks for joint video + audio generation from an ordered list of image, video and audio references with MiniMax-H3, the `ref2va` task of the Ref2VA checkpoint. Supported workflows: - `ref2va`: requires `prompt`, `references` Components: - text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`Qwen2Tokenizer`) processor (`Qwen3VLProcessor`) - vae (`AutoencoderKLMiniMaxH3`) audio_vae (`AutoencoderKLMiniMaxH3Audio`) scheduler (`MiniMaxH3Scheduler`) - audio_scheduler (`MiniMaxH3Scheduler`) transformer_ref (`MiniMaxH3Transformer3DModel`) video_processor - (`VideoProcessor`) + text_encoder (`Qwen3VLForConditionalGeneration`) + tokenizer (`Qwen2Tokenizer`) + processor (`Qwen3VLProcessor`) + vae (`AutoencoderKLMiniMaxH3`) + audio_vae (`AutoencoderKLMiniMaxH3Audio`) + scheduler (`MiniMaxH3Scheduler`) + audio_scheduler (`MiniMaxH3Scheduler`) + transformer_ref (`MiniMaxH3Transformer3DModel`) + video_processor (`VideoProcessor`) Inputs: references (`list`): - The references to condition on, **in the order the model should read them**: the order labels them in the - prompt presentation and lays them out on the shared rotary clock, so a different order is a different - request. Every [`MiniMaxH3Reference`] carries exactly one medium, a path or in-memory media — `image` (at - most 9), `video` at its own `fps` (at most 3, whose `audio` soundtrack is conditioned on as well), or - `audio` at its own `sample_rate` (at most 3) — for at most 12 references in total, and audio references - cannot be the only ones. A path is decoded when the reference is built, so these blocks only ever see - pixels and samples. + The references to condition on, **in the order the model should read them**: the order labels them in the prompt + presentation and lays them out on the shared rotary clock, so a different order is a different request. Every + [`MiniMaxH3Reference`] carries exactly one medium, a path or in-memory media — `image` (at most 9), `video` at its + own `fps` (at most 3, whose `audio` soundtrack is conditioned on as well), or `audio` at its own `sample_rate` (at + most 3) — for at most 12 references in total, and audio references cannot be the only ones. A path is decoded when + the reference is built, so these blocks only ever see pixels and samples. height (`int`, *optional*): Height of the generated video in pixels, a multiple of 32. width (`int`, *optional*): Width of the generated video in pixels, a multiple of 32. num_frames (`int`, *optional*): - Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can - decode. May be left out, but only when exactly one reference carries audio, in which case the duration is - that soundtrack's. + Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode. + May be left out, but only when exactly one reference carries audio, in which case the duration is that + soundtrack's. prompt (`str`): The prompt to guide generation, a single string. generator (`Generator`, *optional*): - The generator of the request. The conditioning noise is drawn from it before the target noise of the - prepare-latents step. + The generator of the request. The video noise is drawn from it first, then the audio noise. latents (`Tensor`, *optional*): - Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used - instead of the draw. + Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the + draw. audio_latents (`Tensor`, *optional*): Pre-generated audio noise of shape `(2, 32, num_audio_latents)`. num_inference_steps (`int`): The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + The structural description of the packed sequence the transformer reads by name: `token_tags`, `position_ids` and + the three row-index tensors. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents. + Output format: 'pil', 'np' or 'pt'. Outputs: videos (`list`): @@ -359,6 +373,7 @@ class MiniMaxH3Ref2VABlocks(SequentialPipelineBlocks): MiniMaxH3PrepareLatentsStep, MiniMaxH3SetTimestepsStep, MiniMaxH3Ref2VADenoiseStep, + MiniMaxH3AfterDenoiseStep, MiniMaxH3DecodeStep, ] block_names = [ @@ -369,6 +384,7 @@ class MiniMaxH3Ref2VABlocks(SequentialPipelineBlocks): "prepare_latents", "set_timesteps", "denoise", + "after_denoise", "decode", ] # The `ref2va` task name, i.e. the value a `workflow=` argument to `ModularPipeline.from_pretrained` would take diff --git a/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py b/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py index 97ab2314beaa..125b002d1305 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py @@ -62,6 +62,18 @@ def vae_latent_channels(self): return self.vae.config.latent_channels return 24 + @property + def vae_frames_per_chunk(self): + if getattr(self, "vae", None) is not None: + return self.vae.config.clip_length + return 17 + + @property + def vae_latents_per_chunk(self): + if getattr(self, "vae", None) is not None: + return self.vae.tokens_chunk_size + return 5 + @property def audio_sampling_rate(self): if getattr(self, "audio_vae", None) is not None: diff --git a/src/diffusers/modular_pipelines/minimax_h3/packing.py b/src/diffusers/modular_pipelines/minimax_h3/packing.py index 94e14226c778..a76badbc9960 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/packing.py +++ b/src/diffusers/modular_pipelines/minimax_h3/packing.py @@ -36,8 +36,6 @@ backends available. """ -from dataclasses import dataclass - import numpy as np import torch @@ -58,11 +56,6 @@ MINIMAX_H3_MIN_DURATION = 5.0 MINIMAX_H3_MAX_DURATION = 15.0 -# The video VAE encodes 17 pixel frames per chunk and drops the 3 trailing latent frames of every chunk, so -# `17 * n + 5` pixel frames map to `5 * n + 2` latent frames. -MINIMAX_H3_FRAMES_PER_CHUNK = 17 -MINIMAX_H3_LATENTS_PER_CHUNK = 5 - # The pixel convention of the video VAE: ImageNet-normalized RGB over a `[0, 1]` base range. MINIMAX_H3_PIXEL_MEAN = (0.485, 0.456, 0.406) MINIMAX_H3_PIXEL_STD = (0.229, 0.224, 0.225) @@ -91,40 +84,6 @@ _ROPE_SPATIAL_SCALE = 32 -@dataclass -class MiniMaxH3PackedSequence: - r""" - The structural description of one packed MiniMax-H3 sequence. - - Attributes: - sequence_length (`int`): - Total number of rows, `L + C + A + V`. - position_ids (`torch.Tensor` of shape `(sequence_length, 3)`, float64): - The `(t, h, w)` rotary coordinate of every row. - token_tags (`torch.Tensor` of shape `(sequence_length,)`): - The modality tag of every row. - video_indices (`torch.Tensor`): - Sequence positions of the video rows: the keyframe conditioning rows first, then the target rows. - audio_indices (`torch.Tensor`): - Sequence positions of the audio rows: reference rows first (`ref2va` only), then the target rows. - text_indices (`torch.Tensor`): - Sequence positions of the text rows. - num_condition_video_rows (`int`): - How many leading entries of `video_indices` are conditioning rows rather than generated rows. - num_condition_audio_rows (`int`): - How many leading entries of `audio_indices` are reference rows rather than generated rows. - """ - - sequence_length: int - position_ids: torch.Tensor - token_tags: torch.Tensor - video_indices: torch.Tensor - audio_indices: torch.Tensor - text_indices: torch.Tensor - num_condition_video_rows: int - num_condition_audio_rows: int - - def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: r""" Resolve a display aspect ratio into a MiniMax-H3 canvas. @@ -163,24 +122,26 @@ def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, return max(multiple, round(height / multiple) * multiple), max(multiple, round(width / multiple) * multiple) -def align_num_frames(num_frames: int) -> int: +def align_num_frames(num_frames: int, frames_per_chunk: int, latents_per_chunk: int) -> int: r""" - Snap a frame count up to the next `17 * n + 5` the video VAE can encode. + Snap a frame count up to the next `frames_per_chunk * n + latents_per_chunk` the video VAE can encode. Args: num_frames (`int`): The requested number of frames. + frames_per_chunk (`int`): Pixel frames the video VAE encodes per chunk, its `clip_length`. + latents_per_chunk (`int`): Latent frames a chunk keeps, the VAE's `tokens_chunk_size`. Returns: `int`: The aligned number of frames. """ if num_frames < 1: raise ValueError(f"`num_frames` must be positive, got {num_frames}.") - while num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: + while num_frames % frames_per_chunk != latents_per_chunk: num_frames += 1 return num_frames -def video_latent_num_frames(num_frames: int) -> int: +def video_latent_num_frames(num_frames: int, frames_per_chunk: int, latents_per_chunk: int) -> int: r""" The number of latent frames the video VAE produces for a `17 * n + 5` frame count. @@ -190,11 +151,11 @@ def video_latent_num_frames(num_frames: int) -> int: Returns: `int`: The number of latent frames, `5 * n + 2`. """ - if num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: - raise ValueError(f"`num_frames` must be of the form 17 * n + 5, got {num_frames}.") - return ( - num_frames - MINIMAX_H3_LATENTS_PER_CHUNK - ) // MINIMAX_H3_FRAMES_PER_CHUNK * MINIMAX_H3_LATENTS_PER_CHUNK + 2 + if num_frames % frames_per_chunk != latents_per_chunk: + raise ValueError( + f"`num_frames` must be of the form {frames_per_chunk} * n + {latents_per_chunk}, got {num_frames}." + ) + return (num_frames - latents_per_chunk) // frames_per_chunk * latents_per_chunk + 2 def audio_latent_num_frames(num_frames: int) -> int: @@ -242,59 +203,6 @@ def patchify_video_latents(latents: torch.Tensor, patch_size: tuple[int, int, in return latents.reshape(-1, channels * patch_t * patch_h * patch_w).contiguous() -def unpatchify_video_tokens( - rows: torch.Tensor, - num_latent_frames: int, - latent_height: int, - latent_width: int, - channels: int, - patch_size: tuple[int, int, int], -) -> torch.Tensor: - r""" - Unpack transformer rows back into video latents. The inverse of [`patchify_video_latents`]. - - Args: - rows (`torch.Tensor` of shape `(num_patches, channels * prod(patch_size))`): The packed rows. - num_latent_frames (`int`): Number of latent frames. - latent_height (`int`): Latent height. - latent_width (`int`): Latent width. - channels (`int`): Number of latent channels. - patch_size (`tuple[int, int, int]`): The `(t, h, w)` patch. - - Returns: - `torch.Tensor` of shape `(batch_size, channels, num_latent_frames, latent_height, latent_width)`. - """ - patch_t, patch_h, patch_w = patch_size - rows = rows.reshape( - -1, - num_latent_frames // patch_t, - latent_height // patch_h, - latent_width // patch_w, - channels, - patch_t, - patch_h, - patch_w, - ) - rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7) - return rows.reshape(-1, channels, num_latent_frames, latent_height, latent_width).contiguous() - - -def unpack_audio_tokens(rows: torch.Tensor, num_audio_latents: int) -> torch.Tensor: - r""" - Unpack the channel-major audio rows into audio VAE latents. - - Args: - rows (`torch.Tensor` of shape `(num_audio_latents * 2, latent_channels)`): The packed audio rows. - num_audio_latents (`int`): Number of audio latents per channel. - - Returns: - `torch.Tensor` of shape `(2, latent_channels, num_audio_latents)`: One batch item per stereo channel, which - is what the mono audio VAE consumes. - """ - rows = rows.reshape(MINIMAX_H3_AUDIO_CHANNELS, num_audio_latents, rows.shape[-1]) - return rows.permute(0, 2, 1).contiguous() - - def _spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: r""" One aspect-normalized spatial rotary axis: `dim // patch` coordinates centred on the unit interval, scaled up by @@ -341,7 +249,7 @@ def build_packed_sequence( num_audio_latents: int, patch_size: tuple[int, int, int], keyframe_anchors: tuple[str, ...] = (), -) -> MiniMaxH3PackedSequence: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: r""" Build the `[text | keyframe conditions | target audio | target video]` layout used by the `t2va` and `fl2va` tasks. @@ -360,7 +268,8 @@ def build_packed_sequence( latent frame, `"last"` at the last one. Returns: - [`MiniMaxH3PackedSequence`] + `tuple`: `position_ids`, `token_tags`, `video_indices`, `audio_indices`, `text_indices`, and the number of + leading video and audio rows that are conditioning rather than generated. """ _, patch_h, patch_w = patch_size rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) @@ -421,20 +330,15 @@ def build_packed_sequence( token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG - return MiniMaxH3PackedSequence( - sequence_length=sequence_length, - position_ids=position_ids, - token_tags=token_tags, - video_indices=video_indices, - audio_indices=audio_indices, - text_indices=text_indices, - num_condition_video_rows=num_condition_rows, - num_condition_audio_rows=0, - ) + return position_ids, token_tags, video_indices, audio_indices, text_indices, num_condition_rows, 0 def build_row_timesteps( - layout: MiniMaxH3PackedSequence, + video_indices: torch.Tensor, + audio_indices: torch.Tensor, + num_condition_video_rows: int, + num_condition_audio_rows: int, + num_text_tokens: int, video_timestep: float, audio_timestep: float, condition_video_timestep: float, @@ -449,7 +353,11 @@ def build_row_timesteps( output head and inherit the video timestep. Args: - layout ([`MiniMaxH3PackedSequence`]): The packed layout. + video_indices (`torch.Tensor`): Sequence positions of the video rows, conditioning rows first. + audio_indices (`torch.Tensor`): Sequence positions of the audio rows, reference rows first. + num_condition_video_rows (`int`): How many leading video rows are conditioning rows. + num_condition_audio_rows (`int`): How many leading audio rows are reference rows. + num_text_tokens (`int`): Number of text rows, which never reach an output head. video_timestep (`float`): Timestep of the generated video rows. audio_timestep (`float`): Timestep of the generated audio rows. condition_video_timestep (`float`): Timestep of the video conditioning rows. @@ -458,8 +366,9 @@ def build_row_timesteps( Returns: `tuple[torch.Tensor, torch.Tensor]`: the distinct timesteps, sorted, and the index of every row into them. """ - row_timesteps = torch.full((layout.sequence_length,), video_timestep, dtype=torch.float32) - row_timesteps[layout.video_indices[: layout.num_condition_video_rows]] = condition_video_timestep - row_timesteps[layout.audio_indices[layout.num_condition_audio_rows :]] = audio_timestep - row_timesteps[layout.audio_indices[: layout.num_condition_audio_rows]] = condition_audio_timestep + sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens) + row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32) + row_timesteps[video_indices[:num_condition_video_rows]] = condition_video_timestep + row_timesteps[audio_indices[num_condition_audio_rows:]] = audio_timestep + row_timesteps[audio_indices[:num_condition_audio_rows]] = condition_audio_timestep return torch.unique(row_timesteps, sorted=True, return_inverse=True) diff --git a/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py b/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py index 3518c9558a0c..f68fd0217028 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py +++ b/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py @@ -59,11 +59,8 @@ MINIMAX_H3_AUDIO_TAG, MINIMAX_H3_CANVAS_MULTIPLE, MINIMAX_H3_FPS, - MINIMAX_H3_FRAMES_PER_CHUNK, - MINIMAX_H3_LATENTS_PER_CHUNK, MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG, - MiniMaxH3PackedSequence, _spatial_position_grid, _temporal_position_grid, resolve_canvas_size, @@ -440,7 +437,7 @@ def build_ref2va_packed_sequence( latent_width: int, num_audio_latents: int, patch_size: tuple[int, int, int], -) -> MiniMaxH3PackedSequence: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: r""" Build the `[text | reference blocks | target audio | target video]` layout of the `ref2va` task. @@ -457,7 +454,8 @@ def build_ref2va_packed_sequence( patch_size (`tuple[int, int, int]`): The transformer's `(t, h, w)` patch. Returns: - [`MiniMaxH3PackedSequence`] + `tuple`: `position_ids`, `token_tags`, `video_indices`, `audio_indices`, `text_indices`, and the number of + leading video and audio rows that are references rather than generated. """ _, patch_h, patch_w = patch_size num_text_tokens = text_token_tags.shape[0] @@ -539,15 +537,14 @@ def build_ref2va_packed_sequence( token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG - return MiniMaxH3PackedSequence( - sequence_length=sequence_length, - position_ids=position_ids, - token_tags=token_tags, - video_indices=video_indices, - audio_indices=audio_indices, - text_indices=text_indices, - num_condition_video_rows=num_reference_video_rows, - num_condition_audio_rows=num_reference_audio_rows, + return ( + position_ids, + token_tags, + video_indices, + audio_indices, + text_indices, + num_reference_video_rows, + num_reference_audio_rows, ) @@ -819,7 +816,7 @@ def emit(segment: tuple[list[int], list[int]]) -> None: return token_ids, token_tags -def trim_reference_num_frames(num_frames: int) -> int: +def trim_reference_num_frames(num_frames: int, frames_per_chunk: int, latents_per_chunk: int) -> int: r""" Snap a reference video's frame count *down* to a `17 * n + 5` the video VAE encodes without padding. @@ -834,8 +831,4 @@ def trim_reference_num_frames(num_frames: int) -> int: """ if num_frames < 1: raise ValueError(f"A reference video must have at least one frame, got {num_frames}.") - return ( - max(1, (num_frames - MINIMAX_H3_LATENTS_PER_CHUNK) // MINIMAX_H3_FRAMES_PER_CHUNK) - * MINIMAX_H3_FRAMES_PER_CHUNK - + MINIMAX_H3_LATENTS_PER_CHUNK - ) + return max(1, (num_frames - latents_per_chunk) // frames_per_chunk) * frames_per_chunk + latents_per_chunk From 5e846542656a47ffc0d9c01967d9a73bb59df8d6 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 3 Aug 2026 16:32:15 +0000 Subject: [PATCH 07/16] Give each MiniMax-H3 reference its own type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `MiniMaxH3Reference` was one dataclass standing in for three unrelated shapes, with a `__post_init__` policing which of its five fields may co-occur and four path-decoding branches inside the constructor. It splits into `MiniMaxH3ImageReference`, `MiniMaxH3VideoReference` (frames, fps, and the video's own soundtrack) and `MiniMaxH3AudioReference` under an empty base, so the types say what the check used to. Opening media files moves out of the dataclass entirely — the blocks' own contract says they never open files — into an opt-in `MiniMaxH3Ref2VALoadReferencesStep` plus exported `decode_reference_video` / `decode_reference_audio` helpers, so a path becomes a request in exactly one place. The module constants stop leaking into the blocks. Everything read by more than one block — fps, the duration bounds, the canvas multiple, the text-encoder layer, the pixel stats, the encode seed, the modality tags — becomes a property on the pipeline, read as `components.fps` the way the blocks already read `components.patch_size`; `canvas_multiple` genuinely derives, as the VAE's spatial compression times the patch width. The per-modality reference limits are read by the setup step alone and become its `__init__` config. `MiniMaxH3PreparedReference` is deleted. The reference encoder used to write latent geometry into objects it declared as an input, and the layout step read it back — an undeclared boundary, and a hand-assembled chain that skipped the encoder built a packed sequence with zero reference rows and no error. The layout now reads geometry off the shapes of the `condition_latents` it already receives, and `audio_condition_latents` becomes a list so the per-reference counts survive. The setup step emits normalized media as the same public reference types, the `image` -> `keyframes` pattern the resize step already uses for the prepared/raw boundary. The generated frames and audio are bit-exact on the documented cases: the reference split swept alone; the constants move is verified property-by- property against the values it replaced, with the layout builder byte-identical old-vs-new on synthetic layouts, and the combined tree swept after the follow-up commit. Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/minimax_h3/__init__.py | 24 +- .../minimax_h3/before_denoise.py | 68 +-- .../minimax_h3/before_encoder.py | 186 +++++--- .../modular_pipelines/minimax_h3/decoders.py | 7 +- .../modular_pipelines/minimax_h3/encoders.py | 134 +++--- .../minimax_h3/modular_pipeline.py | 82 ++++ .../modular_pipelines/minimax_h3/packing.py | 29 +- .../minimax_h3/packing_ref2va.py | 439 ++++++------------ .../minimax_h3/reference_loading.py | 272 +++++++++++ .../test_modular_pipeline_minimax_h3.py | 266 +++++++---- 10 files changed, 933 insertions(+), 574 deletions(-) create mode 100644 src/diffusers/modular_pipelines/minimax_h3/reference_loading.py diff --git a/src/diffusers/modular_pipelines/minimax_h3/__init__.py b/src/diffusers/modular_pipelines/minimax_h3/__init__.py index 6f492f17eff1..eac92574b60c 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/__init__.py +++ b/src/diffusers/modular_pipelines/minimax_h3/__init__.py @@ -23,7 +23,17 @@ else: _import_structure["modular_blocks_minimax_h3"] = ["MiniMaxH3Blocks", "MiniMaxH3Ref2VABlocks"] _import_structure["modular_pipeline"] = ["MiniMaxH3ModularPipeline", "MiniMaxH3Ref2VAModularPipeline"] - _import_structure["packing_ref2va"] = ["MiniMaxH3Reference"] + _import_structure["packing_ref2va"] = [ + "MiniMaxH3AudioReference", + "MiniMaxH3ImageReference", + "MiniMaxH3Reference", + "MiniMaxH3VideoReference", + ] + _import_structure["reference_loading"] = [ + "MiniMaxH3Ref2VALoadReferencesStep", + "decode_reference_audio", + "decode_reference_video", + ] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -34,7 +44,17 @@ else: from .modular_blocks_minimax_h3 import MiniMaxH3Blocks, MiniMaxH3Ref2VABlocks from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline - from .packing_ref2va import MiniMaxH3Reference + from .packing_ref2va import ( + MiniMaxH3AudioReference, + MiniMaxH3ImageReference, + MiniMaxH3Reference, + MiniMaxH3VideoReference, + ) + from .reference_loading import ( + MiniMaxH3Ref2VALoadReferencesStep, + decode_reference_audio, + decode_reference_video, + ) else: import sys diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py index cecb04598e37..95bd786cca97 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -21,12 +21,6 @@ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline from .packing import ( - MINIMAX_H3_AUDIO_CHANNELS, - MINIMAX_H3_CANVAS_MULTIPLE, - MINIMAX_H3_FPS, - MINIMAX_H3_KEYFRAME_NOISE_AUG, - MINIMAX_H3_MAX_DURATION, - MINIMAX_H3_MIN_DURATION, align_num_frames, audio_latent_num_frames, build_packed_sequence, @@ -35,7 +29,7 @@ resolve_canvas_size, video_latent_num_frames, ) -from .packing_ref2va import MiniMaxH3PreparedReference, build_ref2va_packed_sequence +from .packing_ref2va import MiniMaxH3Reference, build_ref2va_packed_sequence logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -143,10 +137,10 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - # Without a keyframe to take the aspect ratio from, MiniMax-H3 generates on its own 16:9 canvas. if block_state.height is None: - block_state.height, block_state.width = resolve_canvas_size(16, 9) - if block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE: + block_state.height, block_state.width = resolve_canvas_size(16, 9, components.canvas_multiple) + if block_state.height % components.canvas_multiple or block_state.width % components.canvas_multiple: raise ValueError( - f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got " + f"`height` and `width` must be multiples of {components.canvas_multiple}, got " f"{block_state.height}x{block_state.width}." ) @@ -161,13 +155,13 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - block_state.num_frames = aligned_num_frames # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds. - duration = block_state.num_frames / MINIMAX_H3_FPS - if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + duration = block_state.num_frames / components.fps + if not components.min_duration <= duration <= components.max_duration: raise ValueError( - f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at " - f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " - f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and " - f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames}." + f"MiniMax-H3 generates between {components.min_duration} and {components.max_duration} seconds " + f"at {components.fps} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE " + f"can encode, must be between {int(components.min_duration * components.fps)} and " + f"{int(components.max_duration * components.fps)}, got {block_state.num_frames}." ) ratio = components.vae_spatial_compression_ratio @@ -227,9 +221,24 @@ def inputs(self) -> list[InputParam]: ), InputParam( name="prepared_references", - type_hint=list[MiniMaxH3PreparedReference], + type_hint=list[MiniMaxH3Reference], required=True, - description="The prepared references, in packed order, with their latent geometry filled in.", + description="The references normalized by the setup step, in packed order.", + ), + InputParam( + name="condition_latents", + type_hint=list[torch.Tensor], + required=True, + description=( + "The encoded video conditioning latents, one per image and video reference in packed order. " + "Their shape is where every reference block's geometry comes from." + ), + ), + InputParam( + name="audio_condition_latents", + type_hint=list[torch.Tensor], + required=True, + description="The encoded audio conditioning rows, one per audio-bearing reference in packed order.", ), InputParam.template("height", required=True, description="Height of the generated video in pixels."), InputParam.template("width", required=True, description="Width of the generated video in pixels."), @@ -321,6 +330,8 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt ) = build_ref2va_packed_sequence( block_state.text_token_tags, block_state.prepared_references, + block_state.condition_latents, + block_state.audio_condition_latents, block_state.num_latent_frames, block_state.latent_height, block_state.latent_width, @@ -402,8 +413,11 @@ def inputs(self) -> list[InputParam]: ), InputParam( name="audio_condition_latents", - type_hint=torch.Tensor, - description="The audio conditioning rows to prepend, or None for a request that has none.", + type_hint=list[torch.Tensor], + description=( + "The audio conditioning rows to prepend, one tensor per audio-bearing reference in packed " + "order. Empty for a request that has none, which is every `t2va` and `fl2va` one." + ), ), ] @@ -433,7 +447,7 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - # noise directly in row layout. Passing `latents` or `audio_latents` skips its draw and shifts the ones after # it. condition_rows = None - if block_state.condition_latents is not None: + if block_state.condition_latents: # One draw per condition, in packed order. Each is packed on its own because `ref2va` references are # encoded at their own resolutions, so their latents do not share a shape. packed = [] @@ -443,7 +457,7 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - ) # The anchors are not fully clean: the released model noises them to `t = 0.999` and holds them there # for every step. Mixing before the patchify is the same arithmetic, since patchify only permutes. - noised = components.scheduler.scale_noise(condition.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise) + noised = components.scheduler.scale_noise(condition.to(device), components.keyframe_noise_aug, noise) packed.append(patchify_video_latents(noised, patch_size)) condition_rows = torch.cat(packed) # In a hand-assembled chain the canvas reaching the layout is user input, so it can disagree with the @@ -474,7 +488,7 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - if block_state.audio_latents is None: audio_rows = randn_tensor( - (block_state.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, components.audio_latent_channels), + (block_state.num_audio_latents * components.audio_channels, components.audio_latent_channels), generator=block_state.generator, device=device, dtype=torch.float32, @@ -488,8 +502,10 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - if condition_rows is not None: video_rows = torch.cat([condition_rows, video_rows]) - if block_state.audio_condition_latents is not None: - audio_rows = torch.cat([block_state.audio_condition_latents.to(device), audio_rows]) + if block_state.audio_condition_latents: + audio_rows = torch.cat( + [rows.to(device) for rows in block_state.audio_condition_latents] + [audio_rows] + ) block_state.latents, block_state.audio_latents = video_rows, audio_rows self.set_block_state(state, block_state) @@ -587,7 +603,7 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - block_state.text_indices.numel(), float(timestep), float(audio_timestep), - max(float(timestep), MINIMAX_H3_KEYFRAME_NOISE_AUG), + max(float(timestep), components.keyframe_noise_aug), 1.0, ) ) diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py b/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py index 86c045be683e..41eee332faf4 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py @@ -23,24 +23,17 @@ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline from .packing import ( - MINIMAX_H3_CANVAS_MULTIPLE, - MINIMAX_H3_FPS, - MINIMAX_H3_MAX_DURATION, - MINIMAX_H3_MIN_DURATION, align_num_frames, resolve_canvas_size, ) from .packing_ref2va import ( - MINIMAX_H3_MAX_REFERENCE_AUDIOS, - MINIMAX_H3_MAX_REFERENCE_IMAGES, - MINIMAX_H3_MAX_REFERENCE_VIDEOS, - MINIMAX_H3_MAX_REFERENCES, - MiniMaxH3PreparedReference, + MiniMaxH3AudioReference, + MiniMaxH3ImageReference, MiniMaxH3Reference, + MiniMaxH3VideoReference, prepare_reference_frames, prepare_reference_image, prepare_reference_waveform, - reference_kind, reference_media_to_uint8, resample_reference_frames, resolve_reference_image_size, @@ -126,7 +119,9 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - if keyframe is not None ) if block_state.height is None: - block_state.height, block_state.width = resolve_canvas_size(*keyframes[0].size) + block_state.height, block_state.width = resolve_canvas_size( + *keyframes[0].size, components.canvas_multiple + ) prepared = [] for index, keyframe in enumerate(keyframes): @@ -162,6 +157,27 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - class MiniMaxH3Ref2VASetupStep(ModularPipelineBlocks): model_name = "minimax-h3-ref2va" + def __init__( + self, max_images: int = 9, max_videos: int = 3, max_audios: int = 3, max_references: int = 12 + ): + r""" + Resolve a `ref2va` request's plan. + + Args: + max_images (`int`, defaults to 9): Image references a request may carry. + max_videos (`int`, defaults to 3): Video references a request may carry. + max_audios (`int`, defaults to 3): Audio references a request may carry. + max_references (`int`, defaults to 12): References of any modality a request may carry in total. + + The four limits are what MiniMax-H3 documents for the released checkpoint; they bound nothing but this + block's own validation, so a fine-tune that packs more can raise them. + """ + self.max_images = max_images + self.max_videos = max_videos + self.max_audios = max_audios + self.max_references = max_references + super().__init__() + @property def description(self) -> str: return ( @@ -170,15 +186,13 @@ def description(self) -> str: "imply when it was left open, and the latent geometry every later block keys off." ) - @staticmethod - def _check_inputs(components, block_state) -> None: + def _check_inputs(self, components, block_state) -> None: if (block_state.height is None) != (block_state.width is None): raise ValueError("`height` and `width` have to be passed together, or neither of them.") - if block_state.height is not None and ( - block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE - ): + multiple = components.canvas_multiple + if block_state.height is not None and (block_state.height % multiple or block_state.width % multiple): raise ValueError( - f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got " + f"`height` and `width` must be multiples of {multiple}, got " f"{block_state.height}x{block_state.width}." ) # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has @@ -190,13 +204,13 @@ def _check_inputs(components, block_state) -> None: block_state.num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk ) ) - duration = None if aligned_num_frames is None else aligned_num_frames / MINIMAX_H3_FPS - if duration is not None and not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + duration = None if aligned_num_frames is None else aligned_num_frames / components.fps + if duration is not None and not components.min_duration <= duration <= components.max_duration: raise ValueError( - f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at " - f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " - f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and " - f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames} (rounded up to " + f"MiniMax-H3 generates between {components.min_duration} and {components.max_duration} seconds at " + f"{components.fps} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " + f"encode, must be between {int(components.min_duration * components.fps)} and " + f"{int(components.max_duration * components.fps)}, got {block_state.num_frames} (rounded up to " f"{aligned_num_frames})." ) @@ -204,17 +218,27 @@ def _check_inputs(components, block_state) -> None: raise ValueError( "`ref2va` needs at least one reference; use `MiniMaxH3ModularPipeline` for text-only requests." ) - kinds = [reference_kind(index, entry) for index, entry in enumerate(block_state.references)] + for index, entry in enumerate(block_state.references): + if not isinstance(entry, MiniMaxH3Reference): + raise ValueError( + f"`references[{index}]` must be a [`MiniMaxH3ImageReference`], [`MiniMaxH3VideoReference`] or " + f"[`MiniMaxH3AudioReference`], got {type(entry)}. MiniMax-H3 blocks never open media files, so a " + "request that holds paths decodes them first — with " + "[`~modular_pipelines.minimax_h3.decode_reference_video`] and " + "[`~modular_pipelines.minimax_h3.decode_reference_audio`], or by putting a " + "[`MiniMaxH3Ref2VALoadReferencesStep`] in front of these blocks." + ) + kinds = [entry.kind for entry in block_state.references] for kind, limit in ( - ("image", MINIMAX_H3_MAX_REFERENCE_IMAGES), - ("video", MINIMAX_H3_MAX_REFERENCE_VIDEOS), - ("audio", MINIMAX_H3_MAX_REFERENCE_AUDIOS), + ("image", self.max_images), + ("video", self.max_videos), + ("audio", self.max_audios), ): if kinds.count(kind) > limit: raise ValueError(f"MiniMax-H3 accepts at most {limit} {kind} references, got {kinds.count(kind)}.") - if len(kinds) > MINIMAX_H3_MAX_REFERENCES: + if len(kinds) > self.max_references: raise ValueError( - f"MiniMax-H3 accepts at most {MINIMAX_H3_MAX_REFERENCES} references in total, got {len(kinds)}." + f"MiniMax-H3 accepts at most {self.max_references} references in total, got {len(kinds)}." ) if set(kinds) == {"audio"}: raise ValueError( @@ -232,11 +256,14 @@ def inputs(self) -> list[InputParam]: description=( "The references to condition on, **in the order the model should read them**: the order labels " "them in the prompt presentation and lays them out on the shared rotary clock, so a different " - "order is a different request. Every [`MiniMaxH3Reference`] carries exactly one medium, a path or " - "in-memory media — `image` (at most 9), `video` at its own `fps` (at most 3, whose `audio` " - "soundtrack is conditioned on as well), or `audio` at its own `sample_rate` (at most 3) — for at " - "most 12 references in total, and audio references cannot be the only ones. A path is decoded " - "when the reference is built, so these blocks only ever see pixels and samples." + "order is a different request. One dataclass per modality, all holding in-memory media — a " + "[`MiniMaxH3ImageReference`] (at most 9), a [`MiniMaxH3VideoReference`] at its own `fps` (at most " + "3, whose `audio` soundtrack is conditioned on as well), or a [`MiniMaxH3AudioReference`] at its " + "own `sample_rate` (at most 3) — for at most 12 references in total, and audio references cannot " + "be the only ones. These blocks never open a media file: decode with " + "[`~modular_pipelines.minimax_h3.decode_reference_video`] and " + "[`~modular_pipelines.minimax_h3.decode_reference_audio`], which bring the rates along, or put a " + "[`MiniMaxH3Ref2VALoadReferencesStep`] in front of these blocks." ), ), InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."), @@ -260,15 +287,21 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), OutputParam( "prepared_references", - type_hint=list[MiniMaxH3PreparedReference], - description="The references prepared at their own resolutions, in packed order.", + type_hint=list[MiniMaxH3Reference], + description=( + "The references normalized onto MiniMax-H3's own rates and resolutions, in packed order: the " + "same public reference types the request passed in, with an image resized to its own 2048 pixel " + "short edge, a video resampled onto 24 fps and onto the canvas its own aspect ratio resolves " + "to, and a soundtrack put on the audio VAE's sample rate and truncated to the generated " + "duration." + ), ), ] @staticmethod def prepare_references( components, references: list[MiniMaxH3Reference], num_frames: int | None - ) -> tuple[list[MiniMaxH3PreparedReference], int]: + ) -> tuple[list[MiniMaxH3Reference], int]: r""" Resolve the references and, if it was left open, the duration they imply. @@ -278,7 +311,9 @@ def prepare_references( the generated duration. None of this touches the target canvas. A reference that left its `fps` or its `sample_rate` out is taken to already be at MiniMax-H3's own rate, and - its frames or its samples then flow through untouched. + its frames or its samples then flow through untouched. Decoding a file with + [`~modular_pipelines.minimax_h3.decode_reference_video`] or + [`~modular_pipelines.minimax_h3.decode_reference_audio`] fills both in from the container. A video reference goes through the two passes the reference implementation's `ffmpeg` decode applied, in the same order: the constant frame rate resample of `resample_reference_frames` and the LANCZOS rescale of @@ -292,63 +327,74 @@ def prepare_references( The requested frame count, or `None` to derive it from the single audio-bearing reference. Returns: - `tuple[list[MiniMaxH3PreparedReference], int]`: the prepared references, in packed order, and the frame - count. + `tuple[list[MiniMaxH3Reference], int]`: the references normalized onto MiniMax-H3's own rates and + resolutions, in packed order and of the same public types they came in as, and the frame count. """ - resolved = [ - MiniMaxH3PreparedReference(kind=reference_kind(index, entry), has_audio=entry.has_audio) - for index, entry in enumerate(references) - ] - # The duration may be left open, but then exactly one reference may carry audio, or the request is ambiguous. if num_frames is None: - audio_bearing = [index for index, reference in enumerate(resolved) if reference.has_audio] + audio_bearing = [index for index, reference in enumerate(references) if reference.has_audio] if len(audio_bearing) != 1: raise ValueError( "`num_frames` may only be left to the references when exactly one of them carries audio, got " f"{len(audio_bearing)}." ) index = audio_bearing[0] - sample_rate = references[index].sample_rate or components.audio_sampling_rate + sample_rate = references[index].sample_rate + if sample_rate is None: + sample_rate = components.audio_sampling_rate duration = references[index].audio.shape[-1] / sample_rate - if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + if not components.min_duration <= duration <= components.max_duration: raise ValueError( f"`references[{index}]` is {duration:g} seconds long, outside the " - f"{MINIMAX_H3_MIN_DURATION} to {MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates." + f"{components.min_duration} to {components.max_duration} seconds MiniMax-H3 generates." ) num_frames = align_num_frames( - round(duration * MINIMAX_H3_FPS), components.vae_frames_per_chunk, components.vae_latents_per_chunk + round(duration * components.fps), components.vae_frames_per_chunk, components.vae_latents_per_chunk ) # The duration the request generates is the one of the *aligned* frame count, so that is what the # ceiling has to hold for: a 14.99 second soundtrack rounds up to 362 frames, i.e. 15.083 seconds. - if num_frames / MINIMAX_H3_FPS > MINIMAX_H3_MAX_DURATION: + if num_frames / components.fps > components.max_duration: raise ValueError( f"`references[{index}]` is {duration:g} seconds long, which rounds up to {num_frames} frames " - f"(`17 * n + 5`), i.e. {num_frames / MINIMAX_H3_FPS:g} seconds — past the " - f"{MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates. Pass `num_frames` to generate a " + f"(`17 * n + 5`), i.e. {num_frames / components.fps:g} seconds — past the " + f"{components.max_duration} seconds MiniMax-H3 generates. Pass `num_frames` to generate a " "shorter video from this soundtrack." ) num_frames = align_num_frames(num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk) - for reference, entry in zip(resolved, references): - if reference.kind == "image": - image = entry.image - if not isinstance(image, Image.Image): - image = Image.fromarray(reference_media_to_uint8(image)) - image = ImageOps.exif_transpose(image).convert("RGB") - height, width = resolve_reference_image_size(*image.size) - reference.image = prepare_reference_image(image, height, width) - elif reference.kind == "video": - frames = resample_reference_frames(reference_media_to_uint8(entry.video), float(entry.fps)) - reference.frames = prepare_reference_frames(frames, num_frames) - if reference.has_audio: - reference.waveform = prepare_reference_waveform( + prepared = [] + for entry in references: + waveform = None + if entry.has_audio: + sample_rate = entry.sample_rate + if sample_rate is None: + sample_rate = components.audio_sampling_rate + waveform = prepare_reference_waveform( entry.audio, - entry.sample_rate or components.audio_sampling_rate, + sample_rate, components.audio_sampling_rate, - max_duration=num_frames / MINIMAX_H3_FPS, + max_duration=num_frames / components.fps, + ) + + if entry.kind == "image": + image = ImageOps.exif_transpose(entry.image).convert("RGB") + height, width = resolve_reference_image_size(*image.size, components.canvas_multiple) + prepared.append(MiniMaxH3ImageReference(image=prepare_reference_image(image, height, width))) + elif entry.kind == "video": + frames = resample_reference_frames(reference_media_to_uint8(entry.frames), float(entry.fps)) + prepared.append( + MiniMaxH3VideoReference( + frames=prepare_reference_frames(frames, num_frames, components.canvas_multiple), + fps=float(components.fps), + audio=waveform, + sample_rate=None if waveform is None else components.audio_sampling_rate, + ) + ) + else: + prepared.append( + MiniMaxH3AudioReference(audio=waveform, sample_rate=components.audio_sampling_rate) ) - return resolved, num_frames + return prepared, num_frames @torch.no_grad() def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: @@ -356,7 +402,7 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt self._check_inputs(components, block_state) if block_state.height is None: - block_state.height, block_state.width = resolve_canvas_size(16, 9) + block_state.height, block_state.width = resolve_canvas_size(16, 9, components.canvas_multiple) requested_num_frames = block_state.num_frames block_state.prepared_references, block_state.num_frames = self.prepare_references( diff --git a/src/diffusers/modular_pipelines/minimax_h3/decoders.py b/src/diffusers/modular_pipelines/minimax_h3/decoders.py index 2eec15d15638..6620d8a36996 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/decoders.py +++ b/src/diffusers/modular_pipelines/minimax_h3/decoders.py @@ -21,7 +21,6 @@ from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline -from .packing import MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -121,7 +120,7 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - # Audio rows are channel-major, and the mono audio VAE takes the two stereo channels as two batch items. audio_rows = block_state.audio_latents[block_state.num_condition_audio_rows :] - audio_rows = audio_rows.reshape(MINIMAX_H3_AUDIO_CHANNELS, block_state.num_audio_latents, audio_rows.shape[-1]) + audio_rows = audio_rows.reshape(components.audio_channels, block_state.num_audio_latents, audio_rows.shape[-1]) block_state.audio_latents = audio_rows.permute(0, 2, 1).contiguous() self.set_block_state(state, block_state) @@ -181,8 +180,8 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"): video = components.vae.decode(latents, return_dict=False)[0] - pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) - pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) + pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1) video = (video.float() * pixel_std + pixel_mean).clamp(0, 1) block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) diff --git a/src/diffusers/modular_pipelines/minimax_h3/encoders.py b/src/diffusers/modular_pipelines/minimax_h3/encoders.py index e70a40fde61b..a6ae5c7f9a07 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/encoders.py +++ b/src/diffusers/modular_pipelines/minimax_h3/encoders.py @@ -21,16 +21,8 @@ from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline -from .packing import ( - MINIMAX_H3_KEYFRAME_ENCODE_SEED, - MINIMAX_H3_PIXEL_MEAN, - MINIMAX_H3_PIXEL_STD, - MINIMAX_H3_TEXT_ENCODER_LAYER, - MINIMAX_H3_TEXT_TAG, - MINIMAX_H3_VIDEO_TAG, -) from .packing_ref2va import ( - MiniMaxH3PreparedReference, + MiniMaxH3Reference, build_ref2va_presentation, sample_reference_video_frames, trim_reference_num_frames, @@ -96,10 +88,12 @@ def encode_prompt( processor, prompt: str, images: list | None = None, + *, + text_encoder_layer: int,# can you set the default? + text_tag: int,# can you set the default? + video_tag: int, # can you set the default? device: torch.device | None = None, dtype: torch.dtype | None = None, - text_tag: int = MINIMAX_H3_TEXT_TAG, - video_tag: int = MINIMAX_H3_VIDEO_TAG, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Build MiniMax-H3's presentation of a request and encode it. @@ -122,12 +116,12 @@ def encode_prompt( """ num_layers = text_encoder.config.text_config.num_hidden_layers - if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: + if num_layers <= text_encoder_layer: raise ValueError( - f"MiniMax-H3 conditions on `hidden_states[{MINIMAX_H3_TEXT_ENCODER_LAYER}]` of its Qwen3-VL " - f"conditioner, which needs more than {MINIMAX_H3_TEXT_ENCODER_LAYER} decoder layers, but " + f"MiniMax-H3 conditions on `hidden_states[{text_encoder_layer}]` of its Qwen3-VL " + f"conditioner, which needs more than {text_encoder_layer} decoder layers, but " f"`text_encoder` has {num_layers}. The last hidden state of a stack truncated to exactly " - f"{MINIMAX_H3_TEXT_ENCODER_LAYER} layers is post-norm and is not the conditioning MiniMax-H3 expects." + f"{text_encoder_layer} layers is post-norm and is not the conditioning MiniMax-H3 expects." ) pixel_values, image_grid_thw = None, None @@ -177,7 +171,7 @@ def encode_prompt( use_cache=False, output_hidden_states=True, ) - prompt_embeds = outputs.hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER].to(device=device, dtype=dtype) + prompt_embeds = outputs.hidden_states[text_encoder_layer].to(device=device, dtype=dtype) return prompt_embeds, torch.tensor(token_tags, dtype=torch.long) @torch.no_grad() @@ -197,6 +191,9 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - components.processor, block_state.prompt, block_state.keyframes, + text_encoder_layer=components.text_encoder_layer, + text_tag=components.text_tag, + video_tag=components.video_tag, device=components._execution_device, dtype=components.text_encoder.dtype, ) @@ -245,7 +242,9 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @staticmethod - def encode_keyframes(vae, images: list, device: torch.device) -> list[torch.Tensor]: + def encode_keyframes( + vae, images: list, pixel_mean: tuple, pixel_std: tuple, encode_seed: int, device: torch.device + ) -> list[torch.Tensor]: r""" Encode the `fl2va` keyframes into normalized conditioning latents. @@ -258,6 +257,9 @@ def encode_keyframes(vae, images: list, device: torch.device) -> list[torch.Tens vae (`AutoencoderKLMiniMaxH3`): The video VAE. images (`list[PIL.Image.Image]`): The keyframes, already prepared onto the target canvas, in packed order. + pixel_mean (`tuple[float, float, float]`), pixel_std (`tuple[float, float, float]`): + The video VAE's pixel convention, i.e. `components.pixel_mean` / `components.pixel_std`. + encode_seed (`int`): Seed the posterior is sampled under, i.e. `components.keyframe_encode_seed`. device (`torch.device`): The device to run the VAE on. Returns: @@ -267,15 +269,15 @@ def encode_keyframes(vae, images: list, device: torch.device) -> list[torch.Tens """ latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1) latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1) - pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) - pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) + pixel_mean = torch.tensor(pixel_mean, device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(pixel_std, device=device).view(1, -1, 1, 1, 1) keyframe_latents = [] for image in images: pixels = torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None] pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std posterior = vae.encode(pixels, return_dict=False)[0] - latents = posterior.sample(generator=torch.Generator().manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) + latents = posterior.sample(generator=torch.Generator().manual_seed(encode_seed)) # The sampled latent is rounded to float16 before it is normalized: ~11 bits of every conditioning # latent, so the released model's conditioning cannot be reproduced without it. latents = latents.to(torch.float16).float().cpu() @@ -287,7 +289,14 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - block_state = self.get_block_state(state) device = components._execution_device - block_state.condition_latents = self.encode_keyframes(components.vae, block_state.keyframes, device) + block_state.condition_latents = self.encode_keyframes( + components.vae, + block_state.keyframes, + components.pixel_mean, + components.pixel_std, + components.keyframe_encode_seed, + device, + ) self.set_block_state(state, block_state) return components, state @@ -319,7 +328,7 @@ def inputs(self) -> list[InputParam]: InputParam.template("prompt", description="The prompt to guide generation, a single string."), InputParam( name="prepared_references", - type_hint=list[MiniMaxH3PreparedReference], + type_hint=list[MiniMaxH3Reference], required=True, description="The prepared references, in packed order.", ), @@ -350,7 +359,11 @@ def encode_prompt( tokenizer, processor, prompt: str, - references: list[MiniMaxH3PreparedReference], + references: list[MiniMaxH3Reference], + *, + text_encoder_layer: int, # can you add default + text_tag: int, # can you add default + video_tag: int, # can you add default device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -366,7 +379,7 @@ def encode_prompt( Args: prompt (`str`): The prompt to encode. - references (`list[MiniMaxH3PreparedReference]`): + references (`list[MiniMaxH3Reference]`): The prepared references, in packed order, as returned by [`~MiniMaxH3Ref2VASetupStep.prepare_references`]. device (`torch.device`, *optional*): The device to run the conditioner on. @@ -378,12 +391,12 @@ def encode_prompt( """ num_layers = text_encoder.config.text_config.num_hidden_layers - if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: + if num_layers <= text_encoder_layer: raise ValueError( - f"MiniMax-H3 conditions on `hidden_states[{MINIMAX_H3_TEXT_ENCODER_LAYER}]` of its Qwen3-VL " - f"conditioner, which needs more than {MINIMAX_H3_TEXT_ENCODER_LAYER} decoder layers, but " + f"MiniMax-H3 conditions on `hidden_states[{text_encoder_layer}]` of its Qwen3-VL " + f"conditioner, which needs more than {text_encoder_layer} decoder layers, but " f"`text_encoder` has {num_layers}. The last hidden state of a stack truncated to exactly " - f"{MINIMAX_H3_TEXT_ENCODER_LAYER} layers is post-norm and is not the conditioning MiniMax-H3 expects." + f"{text_encoder_layer} layers is post-norm and is not the conditioning MiniMax-H3 expects." ) merge_size = processor.image_processor.merge_size**2 @@ -394,26 +407,26 @@ def encode_prompt( pixel_values, image_grid_thw = vision["pixel_values"], vision["image_grid_thw"] image_token_counts = [int(grid.prod()) // merge_size for grid in image_grid_thw] - pixel_values_videos, video_grid_thw, video_block_token_counts = None, None, [] + pixel_values_videos, video_grid_thw = None, None + video_block_token_counts, video_block_timestamps = [], [] videos = [reference for reference in references if reference.kind == "video"] if videos: sampled = [sample_reference_video_frames(reference.frames) for reference in videos] - for reference, (_, block_timestamps) in zip(videos, sampled): - reference.block_timestamps = block_timestamps + video_block_timestamps = [timestamps for _, timestamps in sampled] vision = processor.video_processor( videos=[np.stack(frames) for frames, _ in sampled], do_sample_frames=False, return_tensors="pt" ) pixel_values_videos, video_grid_thw = vision["pixel_values_videos"], vision["video_grid_thw"] video_block_token_counts = [int(grid[1]) * int(grid[2]) // merge_size for grid in video_grid_thw] - for reference, grid in zip(videos, video_grid_thw): - if int(grid[0]) != len(reference.block_timestamps): + for timestamps, grid in zip(video_block_timestamps, video_grid_thw): + if int(grid[0]) != len(timestamps): raise ValueError( f"The processor merged a reference video into {int(grid[0])} vision blocks, but MiniMax-H3 " - f"labels {len(reference.block_timestamps)} of them." + f"labels {len(timestamps)} of them." ) token_ids, token_tags = build_ref2va_presentation( - tokenizer, prompt, references, image_token_counts, video_block_token_counts + tokenizer, prompt, references, image_token_counts, video_block_token_counts, video_block_timestamps ) input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) # Qwen3-VL lays its 3D rotary positions out per modality run, which it reads off the token type ids the @@ -446,7 +459,7 @@ def encode_prompt( use_cache=False, output_hidden_states=True, ) - prompt_embeds = outputs.hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER].to(device=device, dtype=dtype) + prompt_embeds = outputs.hidden_states[text_encoder_layer].to(device=device, dtype=dtype) return prompt_embeds, torch.tensor(token_tags, dtype=torch.long) @torch.no_grad() @@ -466,6 +479,9 @@ def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineSt components.processor, block_state.prompt, block_state.prepared_references, + text_encoder_layer=components.text_encoder_layer, + text_tag=components.text_tag, + video_tag=components.video_tag, device=components._execution_device, dtype=components.text_encoder.dtype, ) @@ -499,9 +515,9 @@ def inputs(self) -> list[InputParam]: return [ InputParam( name="prepared_references", - type_hint=list[MiniMaxH3PreparedReference], + type_hint=list[MiniMaxH3Reference], required=True, - description="The prepared references, in packed order. Their latent geometry is filled in here.", + description="The references normalized by the setup step, in packed order.", ), ] @@ -519,20 +535,22 @@ def intermediate_outputs(self) -> list[OutputParam]: ), OutputParam( "audio_condition_latents", - type_hint=torch.Tensor, + type_hint=list[torch.Tensor], description=( - "The clean audio conditioning rows of the reference soundtracks, in packed order, or None when " - "the references carry none." + "The clean audio conditioning rows of the reference soundtracks, one `(num_audio_latents * 2, " + "audio_latent_channels)` tensor per audio-bearing reference in packed order. One entry per " + "reference rather than one concatenated block, because the packed layout is built from the row " + "count of each." ), ), ] @staticmethod def encode_references( - components, references: list[MiniMaxH3PreparedReference], device: torch.device | None = None - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + components, references: list[MiniMaxH3Reference], device: torch.device | None = None + ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: r""" - Encode the references into packed conditioning rows, and resolve their latent geometry. + Encode the references into conditioning latents. Image and video references go through the video VAE with the same recipe the `fl2va` keyframes use: the posterior is *sampled* under a generator seeded with 42 independently of the request seed, and the sampled @@ -540,20 +558,23 @@ def encode_references( encoder alone, while a video reference goes through the 17-frames-per-5-latents temporal chunking. Reference soundtracks instead take the posterior *mean*, and are never sampled. + The latent geometry every later block keys off is the shape of what this returns, so nothing has to be + written back onto the references. + Args: - references (`list[MiniMaxH3PreparedReference]`): - The prepared references, in packed order. Their latent geometry is filled in here. + references (`list[MiniMaxH3Reference]`): + The references normalized by the setup step, in packed order. device (`torch.device`, *optional*): The device to run the VAEs on. Returns: - `tuple[torch.Tensor, torch.Tensor]`: the `(num_condition_video_rows, latent_channels * prod(patch_size))` - video rows and the `(num_condition_audio_rows, audio_latent_channels)` audio rows, both float32 on CPU and - both `None` when the references carry no such rows. + `tuple[list[torch.Tensor], list[torch.Tensor]]`: one `(1, latent_channels, num_latent_frames, + latent_height, latent_width)` tensor per image and video reference, and one `(num_audio_latents * 2, + audio_latent_channels)` tensor per audio-bearing reference, both in packed order and float32 on CPU. """ latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1) latents_std = torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1) - pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) - pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) + pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1) audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean).view(1, 1, -1) audio_latents_std = torch.tensor(components.audio_vae.config.latents_std).view(1, 1, -1) @@ -575,23 +596,22 @@ def encode_references( # A single frame is encoded by the (tiled) spatial encoder alone; a video goes through the temporal # chunking, which is what turns `17 * n + 5` frames into `5 * n + 2` latent frames. posterior = components.vae.encode(pixels, return_dict=False)[0] - latents = posterior.sample(generator=torch.Generator().manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) + latents = posterior.sample( + generator=torch.Generator().manual_seed(components.keyframe_encode_seed) + ) # The sampled latent is rounded to float16 before it is normalized: ~11 bits of every conditioning # latent, so the released model's conditioning cannot be reproduced without it. latents = latents.to(torch.float16).float().cpu() - reference.num_latent_frames = latents.shape[2] - reference.latent_height, reference.latent_width = latents.shape[3], latents.shape[4] video_latents.append((latents - latents_mean) / latents_std) if reference.has_audio: - posterior = components.audio_vae.encode(reference.waveform.to(device)[:, None], return_dict=False)[0] + posterior = components.audio_vae.encode(reference.audio.to(device)[:, None], return_dict=False)[0] # Channel-major rows: the two stereo channels are two batch items of the mono audio VAE. latents = posterior.mode().float().cpu().transpose(1, 2) - reference.num_audio_latents = latents.shape[1] normalized = (latents - audio_latents_mean) / audio_latents_std audio_rows.append(normalized.reshape(-1, components.audio_latent_channels)) - return video_latents or None, torch.cat(audio_rows) if audio_rows else None + return video_latents, audio_rows @torch.no_grad() def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: diff --git a/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py b/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py index 125b002d1305..f4084d03b0e7 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py @@ -14,6 +14,12 @@ from ...utils import logging from ..modular_pipeline import ModularPipeline +from .packing import ( + MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_FPS, + MINIMAX_H3_TEXT_TAG, + MINIMAX_H3_VIDEO_TAG, +) logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -92,6 +98,82 @@ def patch_size(self): return tuple(self.transformer.config.patch_size) return (1, 2, 2) + @property + def canvas_multiple(self): + r"""What the generated height and width have to be a multiple of, 32 for the released checkpoint.""" + # A canvas has to survive the VAE's spatial compression and still be a whole number of patch rows wide, so + # the multiple is the product of the two. + return self.vae_spatial_compression_ratio * self.patch_size[2] + + @property + def fps(self): + r"""MiniMax-H3's own frame rate. Everything it generates and conditions on is resampled onto it.""" + return MINIMAX_H3_FPS + + @property + def min_duration(self): + r"""Shortest video MiniMax-H3 generates, in seconds.""" + return 5.0 + + @property + def max_duration(self): + r"""Longest video MiniMax-H3 generates, in seconds.""" + return 15.0 + + @property + def audio_channels(self): + r"""Channels of the generated soundtrack: MiniMax-H3 is stereo, packed channel-major.""" + return MINIMAX_H3_AUDIO_CHANNELS + + @property + def text_encoder_layer(self): + r""" + Which Qwen3-VL hidden state conditions the transformer. + + MiniMax-H3 reads `hidden_states[50]`, not the final one: the last layer is post-norm and is not the + conditioning the released weights were trained against. + """ + return 50 + + @property + def pixel_mean(self): + r"""Per-channel mean the video VAE's input is normalized by, ImageNet's.""" + return (0.485, 0.456, 0.406) + + @property + def pixel_std(self): + r"""Per-channel standard deviation the video VAE's input is normalized by, ImageNet's.""" + return (0.229, 0.224, 0.225) + + @property + def keyframe_encode_seed(self): + r""" + Seed the conditioning posterior is sampled under, independently of the request's own generator. + + Fixed at 42 in the reference implementation, so the same keyframe always encodes to the same anchor. + """ + return 42 + + @property + def keyframe_noise_aug(self): + r""" + The `t` a visual conditioning anchor is held at: 0.999, just short of clean. + + The released model was trained with its anchors very slightly noised, so conditioning on exactly `t = 1.0` + is off-distribution. + """ + return 0.999 + + @property + def text_tag(self): + r"""The modality tag of a text row of the packed sequence.""" + return MINIMAX_H3_TEXT_TAG + + @property + def video_tag(self): + r"""The modality tag of a video row of the packed sequence, which a vision block's rows also carry.""" + return MINIMAX_H3_VIDEO_TAG + class MiniMaxH3Ref2VAModularPipeline(MiniMaxH3ModularPipeline): """ diff --git a/src/diffusers/modular_pipelines/minimax_h3/packing.py b/src/diffusers/modular_pipelines/minimax_h3/packing.py index a76badbc9960..681bc73caab5 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/packing.py +++ b/src/diffusers/modular_pipelines/minimax_h3/packing.py @@ -46,36 +46,19 @@ MINIMAX_H3_AUDIO_TAG = 2 # MiniMax-H3 generates at a fixed 24 fps and was released for a 768 pixel short edge only, with a soft area cap of -# 768x1344 and both axes rounded to a multiple of 32. +# 768x1344. The multiple both axes round to is not here: it follows from the VAE and the transformer, so it is +# `MiniMaxH3ModularPipeline.canvas_multiple` and reaches these helpers as an argument. MINIMAX_H3_FPS = 24 MINIMAX_H3_SHORT_EDGE = 768 MINIMAX_H3_MAX_PIXELS = 768 * 1344 -MINIMAX_H3_CANVAS_MULTIPLE = 32 MINIMAX_H3_MIN_ASPECT_RATIO = 1 / 4 MINIMAX_H3_MAX_ASPECT_RATIO = 4 -MINIMAX_H3_MIN_DURATION = 5.0 -MINIMAX_H3_MAX_DURATION = 15.0 - -# The pixel convention of the video VAE: ImageNet-normalized RGB over a `[0, 1]` base range. -MINIMAX_H3_PIXEL_MEAN = (0.485, 0.456, 0.406) -MINIMAX_H3_PIXEL_STD = (0.229, 0.224, 0.225) - -# MiniMax-H3 conditions on the *unnormalized* hidden state its Qwen3-VL conditioner produces after the 50th of its 64 -# decoder layers, i.e. `hidden_states[50]` (`hidden_states[0]` being the embedding output). -MINIMAX_H3_TEXT_ENCODER_LAYER = 50 # The audio VAE hops 800 samples at 32 kHz, i.e. 40 latents per second. Stereo is carried as two channel-major # blocks of audio rows (and as two batch items at the audio VAE boundary, which is mono). MINIMAX_H3_AUDIO_LATENTS_PER_SECOND = 40 MINIMAX_H3_AUDIO_CHANNELS = 2 -# Conditioning rows are not fully clean: the released model noises keyframe latents to `t = 0.999` and runs them at -# that timestep for every denoising step. -MINIMAX_H3_KEYFRAME_NOISE_AUG = 0.999 - -# The seeded posterior sample of the keyframe VAE encode. Fixed at 42 independently of the request seed. -MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 - # Rotary-time constants. One latent frame spans `5/3 * frames_per_latent` rotary units, where the pattern # `(1, 4, 4, 4, 4)` mirrors the VAE's 17-pixel-frames-to-5-latent-frames grouping; the spatial axes are normalized # by the square root of the latent area and scaled by 32. @@ -84,17 +67,19 @@ _ROPE_SPATIAL_SCALE = 32 -def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: +def resolve_canvas_size(aspect_width: float, aspect_height: float, canvas_multiple: int) -> tuple[int, int]: r""" Resolve a display aspect ratio into a MiniMax-H3 canvas. The short edge starts at 768, the area is capped at `768 * 1344` and both axes are then rounded to the nearest - multiple of 32 — so the final area may end up slightly above the pre-rounding budget. Only the ratio of the two + `canvas_multiple` — so the final area may end up slightly above the pre-rounding budget. Only the ratio of the two arguments matters; pass either the aspect ratio (`16, 9`) or the source dimensions of a keyframe. Args: aspect_width (`float`): Width of the target ratio. aspect_height (`float`): Height of the target ratio. + canvas_multiple (`int`): + What both axes round to, i.e. `components.canvas_multiple` — 32 for the released checkpoint. Returns: `tuple[int, int]`: the `(height, width)` of the canvas. @@ -118,7 +103,7 @@ def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, scale = (MINIMAX_H3_MAX_PIXELS / area) ** 0.5 width, height = width * scale, height * scale - multiple = MINIMAX_H3_CANVAS_MULTIPLE + multiple = canvas_multiple return max(multiple, round(height / multiple) * multiple), max(multiple, round(width / multiple) * multiple) diff --git a/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py b/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py index f68fd0217028..8f0646a4588e 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py +++ b/src/diffusers/modular_pipelines/minimax_h3/packing_ref2va.py @@ -32,32 +32,24 @@ resolution (2048 pixel short edge for images, MiniMax-H3's 768 pixel canvas for videos) and carries its own aspect-normalized spatial grid. -References reach the blocks as in-memory media: a video is decoded frames plus the `fps` they carry, and a soundtrack -a waveform plus its sample rate. A [`MiniMaxH3Reference`] takes a path or a URL too, but it decodes it when it is -built, so that no block of this model ever opens a media file. +References reach the blocks as in-memory media, one dataclass per modality: a [`MiniMaxH3VideoReference`] is decoded +frames plus the `fps` they carry, a [`MiniMaxH3AudioReference`] a waveform plus its sample rate, and a +[`MiniMaxH3ImageReference`] a single image. No block of this model opens a media file; decoding a path is the caller's +job, and [`~modular_pipelines.minimax_h3.reference_loading`] is the convenience for doing it. """ -import contextlib import math -import os -import tempfile -from dataclasses import dataclass, field -from typing import Any -from urllib.parse import unquote, urlparse +from dataclasses import dataclass import numpy as np -import requests import torch from PIL import Image -from ...utils import is_av_available, load_image -from ...utils.constants import DIFFUSERS_REQUEST_TIMEOUT from .packing import ( _ROPE_FRAME_RESCALE, _ROPE_FRAMES_PER_LATENT, MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_AUDIO_TAG, - MINIMAX_H3_CANVAS_MULTIPLE, MINIMAX_H3_FPS, MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG, @@ -83,302 +75,121 @@ MINIMAX_H3_MAX_REFERENCES = 12 -@contextlib.contextmanager -def _local_media_file(media): - r"""The reference media as a local file: a URL is downloaded to a temporary file, removed on the way out.""" - path = str(media) - if not path.startswith(("http://", "https://")): - if not os.path.isfile(path): - raise ValueError( - f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {path} is not a valid path." - ) - yield path - return - - response = requests.get(path, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT) - if response.status_code != 200: - raise ValueError(f"Failed to download {path}. Status code: {response.status_code}") - suffix = os.path.splitext(os.path.basename(unquote(urlparse(path).path)))[1] - download = tempfile.NamedTemporaryFile(suffix=suffix, delete=False) - try: - with download as file: - for chunk in response.iter_content(chunk_size=8192): - file.write(chunk) - yield download.name - finally: - os.remove(download.name) - - -def _import_av(): - r"""PyAV, the soft dependency a reference decodes a media file with.""" - if not is_av_available(): - raise ImportError( - "Decoding a MiniMax-H3 reference from a file needs PyAV. You can install it with `pip install av`, or " - "pass the decoded media itself: frames and the `fps` they carry for a video, a `(channels, num_samples)` " - "waveform and its `sample_rate` for audio." - ) - - import av - - return av - - -def _decode_reference_soundtrack(av, container, stream) -> tuple[torch.Tensor, int]: +@dataclass +class MiniMaxH3Reference: r""" - An audio stream's samples as a `(channels, num_samples)` float32 waveform, at the rate the container carries them. - - Args: - av (`module`): PyAV. - container (`av.container.InputContainer`): The open container. - stream (`av.audio.stream.AudioStream`): The stream to decode. + Base class of the three references a [`MiniMaxH3Ref2VABlocks`] request conditions on: + [`MiniMaxH3ImageReference`], [`MiniMaxH3VideoReference`] and [`MiniMaxH3AudioReference`]. - Returns: - `tuple[torch.Tensor, int]`: the waveform and its sample rate. - """ - sample_rate = int(stream.codec_context.sample_rate) - # Planar float is a format conversion only: the sample rate and the channel layout stay the container's own, and a - # mono soundtrack is upmixed later, by `prepare_reference_waveform`. - resampler = av.audio.resampler.AudioResampler(format="fltp", layout=stream.layout, rate=sample_rate) - chunks = [] - for frame in container.decode(stream): - chunks += [torch.from_numpy(resampled.to_ndarray()) for resampled in resampler.resample(frame)] - # Whatever the resampler is still holding. - chunks += [torch.from_numpy(resampled.to_ndarray()) for resampled in resampler.resample(None)] - return torch.cat(chunks, dim=-1).to(torch.float32), sample_rate - - -def decode_reference_video(media) -> tuple[np.ndarray, float, tuple[torch.Tensor, int] | None]: - r""" - Decode a reference video file into `uint8` RGB frames, at the resolution and the frame rate it carries. + References are passed to the blocks as a list, **in the order the model should read them**: the order labels them + in the prompt presentation and lays them out on the shared rotary clock, so a different order is a different + request. - Args: - media (`str` or `os.PathLike`): Path or URL of the video. + Every reference holds in-memory media, and the rate that media carries where there is one — MiniMax-H3 resamples a + reference onto its own 24 fps and onto the audio VAE's sample rate, so a rate lost on the way in is a request + conditioned at the wrong speed. Decoding a file is the caller's job: + [`~modular_pipelines.minimax_h3.decode_reference_video`] and + [`~modular_pipelines.minimax_h3.decode_reference_audio`] do it along with the rates, and + [`MiniMaxH3Ref2VALoadReferencesStep`] wraps them in a block. - Returns: - `tuple[np.ndarray, float, tuple[torch.Tensor, int]]`: the `(num_frames, height, width, 3)` frames, the frame - rate the container reports, and its soundtrack with that soundtrack's own sample rate, `None` when the - container carries no audio stream. - """ - av = _import_av() - with _local_media_file(media) as path, av.open(path) as container: - stream = container.streams.video[0] - frames, rotation = [], 0.0 - for frame in container.decode(stream): - # The display matrix rotation belongs to the stream, and PyAV surfaces it on every frame of it. - rotation = frame.rotation - frames.append(frame.to_ndarray(format="rgb24")) - frame_rate = float(stream.average_rate or stream.guessed_rate) - soundtrack = None - if container.streams.audio: - # Decoding the frames drained the container, so the soundtrack is read in a second pass over it. - container.seek(0) - soundtrack = _decode_reference_soundtrack(av, container, container.streams.audio[0]) - - if not frames: - raise ValueError(f"No video frames to decode in {media}.") - frames = np.stack(frames) - # `ffmpeg` displays a frame upright by undoing the counterclockwise rotation the display matrix carries, which is - # what this reproduces, snapped to the nearest quarter turn. A non-square pixel aspect ratio is left alone: the - # reference implementation resolved a reference's canvas from its *display* geometry, so a stream that carries a - # sample aspect ratio is conditioned on at the wrong shape, and correcting it is untested guesswork here. - turns = round(rotation / 90.0) % 4 - if turns: - frames = np.ascontiguousarray(np.rot90(frames, k=-turns, axes=(1, 2))) - return frames, frame_rate, soundtrack - - -def decode_reference_audio(media) -> tuple[torch.Tensor, int]: - r""" - Decode a reference audio file into a waveform, at the sample rate it carries. + ```py + >>> import numpy as np + >>> from diffusers.utils import load_image + >>> from diffusers.modular_pipelines.minimax_h3 import ( + ... MiniMaxH3AudioReference, + ... MiniMaxH3ImageReference, + ... MiniMaxH3VideoReference, + ... decode_reference_audio, + ... decode_reference_video, + ... ) - Args: - media (`str` or `os.PathLike`): Path or URL of the audio, or of a video whose soundtrack is taken. + >>> references = [ + ... MiniMaxH3ImageReference(image=load_image("subject.png")), + ... decode_reference_video("motion_ref.mp4"), # frames, their `fps`, and the soundtrack + ... decode_reference_audio("voice.wav"), # waveform and its `sample_rate` + ... ] - Returns: - `tuple[torch.Tensor, int]`: the `(channels, num_samples)` float32 waveform and its sample rate. + >>> # Media a request produced itself declares the rate it was produced at. + >>> frames = np.zeros((30, 480, 854, 3), dtype="uint8") + >>> reference = MiniMaxH3VideoReference(frames=frames, fps=30.0) + ``` """ - av = _import_av() - with _local_media_file(media) as path, av.open(path) as container: - if not container.streams.audio: - raise ValueError(f"No audio stream to decode in {media}.") - return _decode_reference_soundtrack(av, container, container.streams.audio[0]) @dataclass -class MiniMaxH3Reference: +class MiniMaxH3ImageReference(MiniMaxH3Reference): r""" - One omni-reference of a [`MiniMaxH3Ref2VABlocks`] request: an image, a video, or an audio clip. + A subject, style or scene reference: at most 9 per request. - A reference carries exactly one medium — plus, for a video, the `audio` of its own soundtrack, which is then - conditioned on as that reference's own. References are passed to the blocks as a list, **in the order the model - should read them**: the order labels them in the prompt presentation and lays them out on the shared rotary clock, - so a different order is a different request. - - Every medium may be a path or a URL as well as in-memory media. A path is decoded here, when the reference is - built, and with it the rates that come with it: no MiniMax-H3 block opens a media file. Decoding a video or an - audio file needs [PyAV](https://github.com/PyAV-Org/PyAV). + Attributes: + image (`PIL.Image.Image`): + The reference image. It never binds the generated geometry — it is encoded at a 2048 pixel short edge of + its own aspect ratio, whatever canvas the request generates at. + """ - ```py - >>> from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3Reference + image: Image.Image - >>> # A file or a URL is decoded on the spot, at the rate the container reports. - >>> references = [ - ... MiniMaxH3Reference(video="motion_ref.mp4"), - ... MiniMaxH3Reference(image="subject.png"), - ... MiniMaxH3Reference(audio="voice.wav"), - ... ] + kind = "image" + has_audio = False - >>> # In-memory media instead carries the rates it was produced at, MiniMax-H3's own by default. - >>> import numpy as np - >>> frames = np.zeros((30, 480, 854, 3), dtype="uint8") - >>> reference = MiniMaxH3Reference(video=frames, fps=30.0) - ``` +@dataclass +class MiniMaxH3VideoReference(MiniMaxH3Reference): + r""" + A motion and camera reference: at most 3 per request, conditioned on together with its own soundtrack. Attributes: - image (`str`, `os.PathLike`, `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`, *optional*): - A subject, style or scene reference: at most 9 per request. A path or a URL, which is read with - [`~utils.load_image`], a `(height, width, 3)` array or a `(3, height, width)` tensor, `uint8` or floating - point over `[0, 1]`. Mutually exclusive with `video`. - video (`str`, `os.PathLike`, `list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`, *optional*): - A motion and camera reference: at most 3 per request. A path or a URL, which PyAV decodes into frames, a - list of images, a `(num_frames, height, width, 3)` array or a `(num_frames, 3, height, width)` tensor. - Mutually exclusive with `image`. A decoded file brings its soundtrack along, as this reference's own, so - conditioning on a file's motion alone means decoding its frames first, with [`~utils.load_video`]. - fps (`float`, *optional*): - The frame rate `video` carries its frames at, which is what places its vision blocks on the conditioner's - 2 fps grid. Left out, it is the rate the container reports for a decoded file and MiniMax-H3's own 24 fps - for in-memory frames, and `fps` holds that resolved rate once the reference is built. Passing it wins over - both, which is only needed when a container's metadata is wrong. MiniMax-H3's clock is 24 fps, so any other - rate is resampled onto it by dropping and duplicating whole frames. - audio (`str`, `os.PathLike` or `torch.Tensor` of shape `(channels, num_samples)`, *optional*): - A voice or music reference, mono or stereo: at most 3 per request, and never on its own — an audio - reference has to be paired with at least one image or video reference. A path or a URL, which PyAV decodes - into a waveform, or the waveform itself. Passed next to `video`, it is that video's soundtrack instead of a - reference of its own. An audio reference never reaches the conditioner and is encoded by the audio VAE - alone. + frames (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`): + The reference frames: a list of images, a `(num_frames, height, width, 3)` array or a `(num_frames, 3, + height, width)` tensor, `uint8` or floating point over `[0, 1]`. + fps (`float`, *optional*, defaults to 24.0): + The frame rate `frames` carries, which is what places the reference's vision blocks on the conditioner's 2 + fps grid. MiniMax-H3's own clock is 24 fps, so any other rate is resampled onto it by dropping and + duplicating whole frames — which makes this the field to get right when the frames came from a file. + audio (`torch.Tensor` of shape `(channels, num_samples)`, *optional*): + This video's soundtrack, mono or stereo, conditioned on as the reference's own rather than as a reference + of its own. Left out, the reference conditions on motion alone. sample_rate (`int`, *optional*): - The rate `audio` carries its samples at. Left out, it is the rate the container reports for a decoded file, - and for an in-memory waveform the audio VAE's own, which leaves the samples untouched. Passing it wins over - both, which is only needed when a container's metadata is wrong. Any other rate is resampled onto the audio - VAE's own. + The rate `audio` carries its samples at. Left out, it is the audio VAE's own, which leaves the samples + untouched; any other rate is resampled onto it. """ - image: str | os.PathLike | Image.Image | np.ndarray | torch.Tensor | None = None - video: str | os.PathLike | list[Image.Image] | np.ndarray | torch.Tensor | None = None + frames: list[Image.Image] | np.ndarray | torch.Tensor fps: float | None = None - audio: str | os.PathLike | torch.Tensor | None = None + audio: torch.Tensor | None = None sample_rate: int | None = None - def __post_init__(self): - # A video reference conditions on its soundtrack too, so `audio` is a second medium of a video reference. - media = [name for name in ("image", "video", "audio") if getattr(self, name) is not None] - if media not in (["image"], ["video"], ["audio"], ["video", "audio"]): - raise ValueError( - "A `MiniMaxH3Reference` must carry exactly one of `image`, `video` or `audio` — plus, for a video, " - f"the `audio` of its soundtrack — got {media if media else 'none of them'}." - ) + kind = "video" - # A path is decoded on the spot, so that the blocks only ever see in-memory media. A rate the request - # passed explicitly wins over the one the container reports, for a container whose metadata is wrong. - if isinstance(self.image, (str, os.PathLike)): - self.image = load_image(str(self.image)) - if isinstance(self.video, (str, os.PathLike)): - frames, frame_rate, soundtrack = decode_reference_video(self.video) - self.video = frames - self.fps = frame_rate if self.fps is None else self.fps - if soundtrack is not None and self.audio is None: - self.audio, soundtrack_sample_rate = soundtrack - self.sample_rate = soundtrack_sample_rate if self.sample_rate is None else self.sample_rate - if isinstance(self.audio, (str, os.PathLike)): - self.audio, sample_rate = decode_reference_audio(self.audio) - self.sample_rate = sample_rate if self.sample_rate is None else self.sample_rate + def __post_init__(self): if self.fps is None: self.fps = float(MINIMAX_H3_FPS) - @property - def kind(self) -> str: - r"""The modality this reference is packed as: `"image"`, `"video"` or `"audio"`.""" - if self.image is not None: - return "image" - return "video" if self.video is not None else "audio" - @property def has_audio(self) -> bool: - r"""Whether this reference contributes audio rows, i.e. whether it carries a waveform.""" + r"""Whether this reference contributes audio rows, i.e. whether it carries a soundtrack.""" return self.audio is not None -def reference_kind(index: int, entry: Any) -> str: - r""" - The modality of one `references` entry, which the [`MiniMaxH3Reference`] validated at construction. - """ - if not isinstance(entry, MiniMaxH3Reference): - raise ValueError( - f"`references[{index}]` must be a [`MiniMaxH3Reference`], got {type(entry)}. A request is built from " - "the dataclass: `MiniMaxH3Reference(image=...)`, `MiniMaxH3Reference(video=...)` or " - "`MiniMaxH3Reference(audio=...)`." - ) - # A reference decodes a path when it is built, so the blocks only ever see in-memory media. - for name in ("image", "video", "audio"): - if isinstance(getattr(entry, name), (str, os.PathLike)): - raise ValueError( - f"`references[{index}].{name}` is a path. MiniMax-H3 blocks never open media files: rebuild " - "the reference, which decodes a path as it is built." - ) - return entry.kind - - @dataclass -class MiniMaxH3PreparedReference: +class MiniMaxH3AudioReference(MiniMaxH3Reference): r""" - One `ref2va` reference prepared for packing, in packed order. - - A [`MiniMaxH3Reference`] is resolved in three passes: the blocks read the modality off the request (`kind`, - `has_audio`), prepares the pixels or samples (`image`, `frames`, `waveform`), and finally encodes them, which is - what fixes the latent geometry (`num_latent_frames`, `latent_height`, `latent_width`, `num_audio_latents`) the - packed layout is built from. + A voice or music reference: at most 3 per request, and never on its own — an audio reference has to be paired with + at least one image or video reference. It never reaches the conditioner and is encoded by the audio VAE alone. Attributes: - kind (`str`): - `"image"`, `"video"` or `"audio"`. - has_audio (`bool`): - Whether this reference contributes audio rows. Always `True` for `"audio"`, and `True` for a `"video"` the - request passed a soundtrack with. - image (`PIL.Image.Image`): - The prepared reference image. - frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): - The prepared reference video, `uint8` RGB at 24 fps. - waveform (`torch.Tensor` of shape `(2, num_samples)`): - The prepared soundtrack, stereo at the audio VAE's sample rate. - block_timestamps (`list[float]`): - The timestamp of every vision block the conditioner sees for a video reference. - num_latent_frames (`int`), latent_height (`int`), latent_width (`int`): - Latent geometry of the visual rows. - num_audio_latents (`int`): - Number of audio latents per channel. + audio (`torch.Tensor` of shape `(channels, num_samples)`): + The reference waveform, mono or stereo. + sample_rate (`int`, *optional*): + The rate `audio` carries its samples at. Left out, it is the audio VAE's own, which leaves the samples + untouched; any other rate is resampled onto it. """ - kind: str - has_audio: bool = False - image: Any = None - frames: Any = None - waveform: torch.Tensor | None = None - block_timestamps: list[float] = field(default_factory=list) - num_latent_frames: int = 1 - latent_height: int = 0 - latent_width: int = 0 - num_audio_latents: int = 0 - - @property - def num_video_rows(self) -> int: - r"""The number of packed video rows, for the `(1, 2, 2)` patch MiniMax-H3 packs video latents with.""" - return self.num_latent_frames * (self.latent_height // 2) * (self.latent_width // 2) + audio: torch.Tensor + sample_rate: int | None = None - @property - def num_audio_rows(self) -> int: - r"""The number of packed audio rows: one per latent and per stereo channel.""" - return self.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + kind = "audio" + has_audio = True def _temporal_position_span(num_latent_frames: int) -> float: @@ -431,7 +242,9 @@ def _fill_audio_positions( def build_ref2va_packed_sequence( text_token_tags: torch.Tensor, - references: list[MiniMaxH3PreparedReference], + references: list[MiniMaxH3Reference], + condition_latents: list[torch.Tensor], + audio_condition_latents: list[torch.Tensor], num_latent_frames: int, latent_height: int, latent_width: int, @@ -445,8 +258,13 @@ def build_ref2va_packed_sequence( text_token_tags (`torch.Tensor` of shape `(num_text_tokens,)`): The modality tag of every text row. Text is tagged `1`, except for the rows of a reference's vision block, which MiniMax-H3 tags `0` (video). - references (`list[MiniMaxH3PreparedReference]`): - The references, in packed order, with their latent geometry already resolved. + references (`list[MiniMaxH3Reference]`): + The references, in packed order. Only their modality is read here; the geometry comes from the latents. + condition_latents (`list[torch.Tensor]`): + One `(1, channels, num_latent_frames, latent_height, latent_width)` tensor per image and video reference, + in packed order, as [`MiniMaxH3Ref2VAReferenceEncoderStep`] produced them. + audio_condition_latents (`list[torch.Tensor]`): + One `(num_audio_latents * 2, audio_latent_channels)` tensor per audio-bearing reference, in packed order. num_latent_frames (`int`): Number of target latent frames. latent_height (`int`): Target latent height. latent_width (`int`): Target latent width. @@ -461,8 +279,17 @@ def build_ref2va_packed_sequence( num_text_tokens = text_token_tags.shape[0] num_target_video_rows = num_latent_frames * (latent_height // patch_h) * (latent_width // patch_w) num_target_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS - num_reference_video_rows = sum(reference.num_video_rows for reference in references if reference.kind != "audio") - num_reference_audio_rows = sum(reference.num_audio_rows for reference in references) + + # The geometry of every reference block is the shape of what the encoder produced for it, so the two can never + # disagree. Both lists are in packed order but skip the references they do not apply to, so they are consumed as + # iterators alongside the reference list rather than indexed by it. + visual_geometry = iter(tuple(latents.shape[2:5]) for latents in condition_latents) + audio_row_counts = iter(rows.shape[0] for rows in audio_condition_latents) + num_reference_video_rows = sum( + frames * (height // patch_h) * (width // patch_w) + for frames, height, width in (tuple(latents.shape[2:5]) for latents in condition_latents) + ) + num_reference_audio_rows = sum(rows.shape[0] for rows in audio_condition_latents) sequence_length = ( num_text_tokens + num_reference_video_rows @@ -482,39 +309,47 @@ def build_ref2va_packed_sequence( rotary_time = float(num_text_tokens) for reference in references: if reference.kind == "image": - rows = slice(cursor, cursor + reference.num_video_rows) + num_latent_frames_, reference_height, reference_width = next(visual_geometry) + num_video_rows = ( + num_latent_frames_ * (reference_height // patch_h) * (reference_width // patch_w) + ) + rows = slice(cursor, cursor + num_video_rows) cursor = rows.stop video_indices.append(torch.arange(rows.start, rows.stop)) - frame_grid, _ = _frame_position_grid(reference.latent_height, reference.latent_width, patch_h, patch_w) + frame_grid, _ = _frame_position_grid(reference_height, reference_width, patch_h, patch_w) position_ids[rows, 0] = rotary_time position_ids[rows, 1:] = frame_grid # An image is a single frame and takes a single integer rotary slot, not a latent frame's 5/3 units. rotary_time += 1.0 elif reference.kind == "audio": - rows = slice(cursor, cursor + reference.num_audio_rows) + num_audio_rows = next(audio_row_counts) + reference_audio_latents = num_audio_rows // MINIMAX_H3_AUDIO_CHANNELS + rows = slice(cursor, cursor + num_audio_rows) cursor = rows.stop audio_indices.append(torch.arange(rows.start, rows.stop)) - _fill_audio_positions(position_ids, rows, reference.num_audio_latents, rotary_time, target_width_grid) - rotary_time += float(reference.num_audio_latents) + _fill_audio_positions(position_ids, rows, reference_audio_latents, rotary_time, target_width_grid) + rotary_time += float(reference_audio_latents) elif reference.kind == "video": # A video reference's soundtrack rows are packed immediately before its video rows and share their # origin, so the two are rotary-aligned exactly as the generated audio and video are. - audio_rows = slice(cursor, cursor + reference.num_audio_rows) - video_rows = slice(audio_rows.stop, audio_rows.stop + reference.num_video_rows) + num_audio_rows = next(audio_row_counts) if reference.has_audio else 0 + reference_audio_latents = num_audio_rows // MINIMAX_H3_AUDIO_CHANNELS + num_latent_frames_, reference_height, reference_width = next(visual_geometry) + num_video_rows = ( + num_latent_frames_ * (reference_height // patch_h) * (reference_width // patch_w) + ) + audio_rows = slice(cursor, cursor + num_audio_rows) + video_rows = slice(audio_rows.stop, audio_rows.stop + num_video_rows) cursor = video_rows.stop audio_indices.append(torch.arange(audio_rows.start, audio_rows.stop)) video_indices.append(torch.arange(video_rows.start, video_rows.stop)) - frame_grid, width_grid = _frame_position_grid( - reference.latent_height, reference.latent_width, patch_h, patch_w - ) - _fill_audio_positions(position_ids, audio_rows, reference.num_audio_latents, rotary_time, width_grid) - frame_time = _temporal_position_grid(reference.num_latent_frames, rotary_time) + frame_grid, width_grid = _frame_position_grid(reference_height, reference_width, patch_h, patch_w) + _fill_audio_positions(position_ids, audio_rows, reference_audio_latents, rotary_time, width_grid) + frame_time = _temporal_position_grid(num_latent_frames_, rotary_time) position_ids[video_rows, 0] = frame_time.repeat_interleave(frame_grid.shape[0]) - position_ids[video_rows, 1:] = frame_grid.repeat(reference.num_latent_frames, 1) - rotary_time += max( - float(reference.num_audio_latents), _temporal_position_span(reference.num_latent_frames) - ) + position_ids[video_rows, 1:] = frame_grid.repeat(num_latent_frames_, 1) + rotary_time += max(float(reference_audio_latents), _temporal_position_span(num_latent_frames_)) else: raise ValueError(f"A reference must be an 'image', a 'video' or an 'audio', got {reference.kind!r}.") @@ -548,14 +383,16 @@ def build_ref2va_packed_sequence( ) -def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: +def resolve_reference_image_size(width: int, height: int, canvas_multiple: int) -> tuple[int, int]: r""" - Resolve the resolution a reference image is encoded at: a 2048 pixel short edge, both axes rounded to a multiple - of 32. Upscaling is intended, and unlike the target canvas there is no area cap. + Resolve the resolution a reference image is encoded at: a 2048 pixel short edge, both axes rounded to + `canvas_multiple`. Upscaling is intended, and unlike the target canvas there is no area cap. Args: width (`int`): Width of the source image. height (`int`): Height of the source image. + canvas_multiple (`int`): + What both axes round to, i.e. `components.canvas_multiple` — 32 for the released checkpoint. Returns: `tuple[int, int]`: the `(height, width)` the reference is resized to. @@ -566,7 +403,7 @@ def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: raise ValueError(f"A reference image must be within 1:4 and 4:1, got {width}x{height}.") scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height) - multiple = MINIMAX_H3_CANVAS_MULTIPLE + multiple = canvas_multiple return ( max(multiple, round(height * scale / multiple) * multiple), max(multiple, round(width * scale / multiple) * multiple), @@ -648,7 +485,7 @@ def resample_reference_frames(frames: np.ndarray, fps: float) -> np.ndarray: return np.repeat(frames, np.diff(slots, append=math.floor(frames.shape[0] * scale + 0.5)), axis=0) -def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray: +def prepare_reference_frames(frames: np.ndarray, num_frames: int, canvas_multiple: int) -> np.ndarray: r""" Put a reference video onto the canvas its own aspect ratio resolves to, and cap it at the generated frame count. @@ -661,6 +498,8 @@ def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray: frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): The reference video, `uint8` RGB at 24 fps, as returned by [`resample_reference_frames`]. num_frames (`int`): The frame count the reference is truncated to, i.e. the target's own frame count. + canvas_multiple (`int`): + What both axes of its own canvas round to, i.e. `components.canvas_multiple`. Returns: `np.ndarray` of shape `(num_frames, height, width, 3)`: The prepared reference. @@ -670,7 +509,7 @@ def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray: f"A reference video must be `(num_frames, height, width, 3)` RGB frames, got {tuple(frames.shape)}." ) frames = frames[:num_frames] - height, width = resolve_canvas_size(frames.shape[2], frames.shape[1]) + height, width = resolve_canvas_size(frames.shape[2], frames.shape[1], canvas_multiple) if frames.shape[1:3] == (height, width): return frames return np.stack( @@ -753,9 +592,10 @@ def prepare_reference_waveform( def build_ref2va_presentation( tokenizer, prompt: str, - references: list[MiniMaxH3PreparedReference], + references: list[MiniMaxH3Reference], image_token_counts: list[int], video_block_token_counts: list[int], + video_block_timestamps: list[list[float]], ) -> tuple[list[int], list[int]]: r""" Tokenize MiniMax-H3's presentation of a `ref2va` request. @@ -769,9 +609,10 @@ def build_ref2va_presentation( Args: tokenizer (`Qwen2TokenizerFast`): Tokenizer of the conditioner. prompt (`str`): The prompt, appended verbatim. - references (`list[MiniMaxH3PreparedReference]`): The prepared references, in packed order. + references (`list[MiniMaxH3Reference]`): The normalized references, in packed order. image_token_counts (`list[int]`): Number of vision tokens of every image reference's block. video_block_token_counts (`list[int]`): Number of vision tokens per block of every video reference. + video_block_timestamps (`list[list[float]]`): The timestamp of every vision block, per video reference. Returns: `tuple[list[int], list[int]]`: the token ids and their modality tags. A vision block is tagged `0` (video) and @@ -808,7 +649,7 @@ def emit(segment: tuple[list[int], list[int]]) -> None: elif reference.kind == "video": counts["video"] += 1 emit(text(f"