Handle images with MiniMax H3.

This commit is contained in:
Jaret Burkett
2026-08-03 11:58:44 -06:00
parent bf739ff966
commit 88ac27fc8f
4 changed files with 18 additions and 8 deletions

View File

@@ -1,7 +1,9 @@
"""MiniMax-H3 (33B joint video+audio DiT) for ai-toolkit.
Supports t2v (t2va) and first-frame i2v (fl2va) training and sampling, with
joint audio when the dataset provides it. The architecture lives in ./src/:
joint audio when the dataset provides it. Image datasets train as single
latent frames (keyframe-row geometry) and sampling with num_frames 1 renders
a single image. The architecture lives in ./src/:
- transformer.py: packed-sequence DiT, weight-compatible with the original
``MiniMaxAI/MiniMax-H3`` checkpoint keys

View File

@@ -86,12 +86,13 @@ class MiniMaxH3Pipeline:
is_video = num_frames > 1
if is_video:
num_frames = packing.align_num_frames_down(num_frames)
t_lat = packing.video_latent_num_frames(num_frames)
else:
# a single image still runs through the video path with the
# minimum 5-frame clip; the first frame is returned
num_frames = packing.LATENTS_PER_CHUNK
t_lat = packing.video_latent_num_frames(num_frames)
# true single-frame generation (image mode, LTX-2.3 style): one
# latent frame with keyframe-row geometry — the same layout image
# datasets train with, so image LoRAs sample in-distribution
num_frames = 1
t_lat = 1
h_lat = height // 16
w_lat = width // 16
a_lat = packing.audio_latent_num_frames(num_frames)

View File

@@ -777,7 +777,13 @@ class MiniMaxH3VideoVAE(nn.Module):
)
with ctx:
if z.shape[2] == 1:
dec = self._decode_clip(z)[:, :, -1:]
# a lone temporal token is OOD for the chunk-trained ViT
# decoder (visibly patchy output, ~16 dB vs ~35 dB roundtrip).
# Decode it as the first latent of a 2-latent clip instead:
# the first latent of a chunk covers exactly pixel frame 0 in
# the (1, 4, 4, 4, 4) grouping, so that frame is a clean,
# in-distribution single-frame decode.
dec = self._decode_video(torch.cat([z, z], dim=2))[:, :, :1]
else:
dec = self._decode_video(z)

View File

@@ -764,7 +764,8 @@ export const modelArchs: ModelArch[] = [
<p>
Supports t2v and first-frame i2v (ctrl img / i2v datasets) with joint audio. The model is guidance-distilled —
keep guidance scale at 1. Video is fixed 24 fps and frame counts snap down to the 17n+5 grid (5, 22, 39, 56,
..., 107, 124 ≈ 5s).
..., 107, 124 ≈ 5s). Image datasets (num_frames 1) train as single frames, and a sample with num_frames 1
renders a single image.
</p>
</div>
),