[WIP] Minimax h3 follow up (review & refactor) - #14371
Open
yiyixuxu wants to merge 10 commits into
Open
Conversation
`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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
`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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
`_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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
`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 <noreply@anthropic.com>
Both transformer partitions ship in one repo, so there is no reason for two pipeline classes. `MiniMaxH3Ref2VABlocks` and `MiniMaxH3Ref2VAModularPipeline` are gone; `MiniMaxH3Blocks` becomes the auto blockset — the before-encode, text-encoder, VAE-encoder and core-denoise slots each pick their branch on the `references` trigger, the way `z_image` and `krea2` pack theirs — with four workflows: `t2va`, `fl2va`, `fl2va_last_frame`, `ref2va`. `from_pretrained(..., workflow="ref2va")` prunes to the ref2va branch and loads only `transformer_ref`; the released repo's `modular_model_index.json` keeps working because the class name is reused. The two loop denoisers merge — they differed only in which transformer partition they read, which becomes block config — and `patch_size` on the one pipeline reads whichever partition is loaded. The rest is each block owning its own work. Validation moves out of `_check_inputs` into `__call__` — which also leaves ref2va's duration ceiling checked once instead of twice with two different messages. `num_frames` is now required on ref2va: with it passed, every audio-bearing reference already rides at its own truncated length, so deriving it from a soundtrack added no capability — it filled in one number, worked only in the single-soundtrack case, and cost a conditional-optional API rule plus two error paths. Media normalization becomes two staticmethods on the setup step, with the image branch inlined; the seeded VAE conditioning recipe and the Qwen3-VL embedding call are shared functions, `encode_vae_condition` and `get_qwen3vl_prompt_embeds`, used by both halves; the packed-sequence builders land on their layout steps with the rotary helpers module-level beside them. That empties `packing.py` and `packing_ref2va.py`; what remains of the latter — the four public reference dataclasses — is renamed `references.py`. Two guards for hand-assembled chains: `prepare_latents` verifies the packed condition rows match the counts the layout reserved for both modalities, and the ref2va layout validates the canvas multiple the way t2va's already did. All 11 documented cases — t2va, the three fl2va anchorings, the six ref2va configurations, and the defaults path — are bit-exact on frames and audio against the pre-refactor baseline. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
github-actions
Bot
added
documentation
yiyixuxu
changed the title
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Blocks already declare their workflows (`_workflow_map`, `get_workflow`), but loading could not use them: selecting a workflow meant pruning the blocks by hand and calling `init_pipeline` yourself. Both loading entry points now take the name directly. `ModularPipeline.from_pretrained(repo, workflow="ref2va")` prunes the blocks before the component specs are derived, so the pipeline only expects — and a plain `load_components()` only loads — the components that workflow uses. For MiniMax-H3 that is the difference between pulling one 61.7GB transformer partition and pulling both. `load_components(workflow=...)` does the same restriction at load time for a pipeline that kept the full blocks, which is the shape for switching workflows under one ComponentsManager; passing it together with `names` raises. An unknown workflow name now lists the available ones. `unload_components(names)` is the explicit counterpart for taking a component back out: the attribute returns to `None` while the spec survives, so `load_components` can restore it, and an attached ComponentsManager drops it too. The method releases its own reference before the manager's removal runs, and ends with an unconditional `gc.collect()` + `empty_device_cache()` — the manager's cleanup is skipped on its auto-offload path, so relying on it would leave hooked modules cycling in memory. Unloading is removal, not offloading: parking a component in cheaper memory is what the manager's offloading already does. The mixin tests pin the new arguments to the route that already existed: `from_pretrained(workflow=...)` must produce the same component names and byte-equal specs as `get_workflow(...).init_pipeline(repo)`, and `load_components(workflow=...)` on full blocks must load exactly what the pruned pipeline loads. The unload test checks the spec object survives the round trip and the ComponentsManager forgets the component. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
`decode_reference_video` and `decode_reference_audio` were module-level helpers
a user had to know existed next to the classes they construct. They become
`from_file` classmethods on the references themselves — the type names the
modality, so the loader no longer has to: `MiniMaxH3VideoReference.from_file`
returns the frames with the container's true frame rate and its soundtrack,
`MiniMaxH3AudioReference.from_file` a waveform with its sample rate (from an
audio file or a video's soundtrack), and `MiniMaxH3ImageReference.from_file`
goes through `load_image`. The rates riding along is still the point: a rate
lost on the way in — which is what a bare `load_video` does — is a request
conditioned at the wrong speed, with nothing to raise about it.
`MiniMaxH3Ref2VALoadReferencesStep` is deleted. It earned its place when the
dataclasses could not take paths and there was no ergonomic loader; with
`from_file` it was a worse spelling of the same thing, naming the modality in
a `{"video": path}` string key where the classmethod names it in a type. It
was never part of `MiniMaxH3Blocks` and nothing shipped. With the block gone,
`reference_loading.py` held only the machinery behind `from_file`, so it
merges into `references.py` — one module owns the reference story, and only
the PyAV import stays lazy.
The docs build the ref2va example with `from_file`, and a new section feeds a
`t2va` generation straight back as a `ref2va` reference: one pipeline, a
`load_components(workflow=...)` call per task, and an in-memory
`MiniMaxH3VideoReference` built from the generated frames and soundtrack —
already at the model's own rates, so no rate arguments and no lossy container
round-trip.
Parity is untouched: the decode code path is the same, only its packaging
moved.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
part of #14355