github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add support for MiniMax H3 Ref2Vid training

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-13 12:36:03 -06:00
parent ab18528fdb
commit 6ea281973d
6 changed files with 574 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
extensions_built_in/diffusion_models/__init__.py
View File

@@ -19,7 +19,7 @@ from .prx_pixel_t2i import PRXPixelT2IModel
from .krea2 import Krea2Model
from .boogu_image import BooguImageModel, BooguImageEditModel
from .mageflow import MageFlowModel, MageFlowEditModel
from .minimax_h3 import MinimaxH3Model
from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel
AI_TOOLKIT_MODELS = [
# put a list of models here
@@ -57,4 +57,5 @@ AI_TOOLKIT_MODELS = [
MageFlowModel,
MageFlowEditModel,
MinimaxH3Model,
MinimaxH3Ref2VAModel,
]

2
extensions_built_in/diffusion_models/minimax_h3/__init__.py
View File

@@ -1 +1 @@
from .minimax_h3 import MinimaxH3Model
from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel

241
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
View File

@@ -818,6 +818,47 @@ class MinimaxH3Model(BaseModel):
# ------------------------------------------------------------------
# Training forward
# ------------------------------------------------------------------
def _build_condition(
self, batch: "DataLoaderBatchDTO", latent_shape, device, dtype
):
"""Build the packed sequence's condition segment for one train step.
``latent_shape`` is the target's ``(t_lat, h_lat, w_lat)``. Returns
``(cond_rows, keyframe_anchors, image_ref_shapes)`` where ``cond_rows``
is ``(B, num_condition_rows, 96)`` or None. The base model implements
fl2va: the clip's first frame as a keyframe when the dataset asks for
i2v. MinimaxH3Ref2VAModel overrides this with image references from
the control images."""
do_i2v = (
batch is not None
and batch.dataset_config.do_i2v
and getattr(batch, "num_frames", 1) > 1
)
if not do_i2v:
return None, (), ()
if batch.first_frame_latents is not None:
first_latents = batch.first_frame_latents.to(device, torch.float32)
else:
frames = batch.tensor
if frames is None:
raise ValueError(
"do_i2v needs the first frame; no cached "
"first_frame_latents or raw tensors in batch"
)
first_frames = frames[:, 0] if frames.ndim == 5 else frames
first_latents = self.encode_keyframe_latents(
first_frames.unsqueeze(2).to(device)
)
if first_latents.ndim == 4:
first_latents = first_latents.unsqueeze(2)
cond_noise = torch.randn_like(first_latents)
first_latents = (
KEYFRAME_NOISE_AUG_T * first_latents
+ (1.0 - KEYFRAME_NOISE_AUG_T) * cond_noise
)
return patchify_video_latents(first_latents).to(dtype), ("first",), ()
def get_noise_prediction(
self,
latent_model_input: torch.Tensor, # (B, 24, t, h, w) noisy latents
@@ -854,35 +895,10 @@ class MinimaxH3Model(BaseModel):
t_v = 1.0 - sigma_v
t_a = 1.0 - sigma_a
# --- i2v first-frame conditioning rows -------------------------
do_i2v = (
batch is not None
and batch.dataset_config.do_i2v
and getattr(batch, "num_frames", 1) > 1
# --- conditioning rows (fl2va keyframe / ref2va references) ----
cond_rows, keyframe_anchors, image_ref_shapes = self._build_condition(
batch, (t_lat, h_lat, w_lat), device, dtype
)
cond_rows = None
if do_i2v:
if batch.first_frame_latents is not None:
first_latents = batch.first_frame_latents.to(device, torch.float32)
else:
frames = batch.tensor
if frames is None:
raise ValueError(
"do_i2v needs the first frame; no cached "
"first_frame_latents or raw tensors in batch"
)
first_frames = frames[:, 0] if frames.ndim == 5 else frames
first_latents = self.encode_keyframe_latents(
first_frames.unsqueeze(2).to(device)
)
if first_latents.ndim == 4:
first_latents = first_latents.unsqueeze(2)
cond_noise = torch.randn_like(first_latents)
first_latents = (
KEYFRAME_NOISE_AUG_T * first_latents
+ (1.0 - KEYFRAME_NOISE_AUG_T) * cond_noise
)
cond_rows = patchify_video_latents(first_latents).to(dtype)
# --- audio rows -------------------------------------------------
if batch is not None and getattr(batch, "num_frames", None):
@@ -958,7 +974,6 @@ class MinimaxH3Model(BaseModel):
# --- packed layout (per item: text lengths differ) --------------
layouts = []
anchors = ("first",) if cond_rows is not None else ()
for i in range(batch_size):
layouts.append(
build_packed_sequence(
@@ -967,7 +982,8 @@ class MinimaxH3Model(BaseModel):
latent_height=h_lat,
latent_width=w_lat,
num_audio_latents=a_lat,
keyframe_anchors=anchors,
keyframe_anchors=keyframe_anchors,
image_ref_shapes=image_ref_shapes,
)
)
(
@@ -1158,3 +1174,168 @@ class MinimaxH3Model(BaseModel):
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}
class MinimaxH3Ref2VAModel(MinimaxH3Model):
"""Reference-to-video (ref2va): the control images ride along as reference
blocks in the packed sequence (plus ``<Picture i>: `` vision blocks in the
Qwen3-VL conditioning) instead of anchoring the first frame.
Image references only for now. Every reference keeps its OWN aspect ratio
and is resized to the TARGET's pixel area (axes snapped to /32), then sits
on its own aspect-normalized rotary grid — like the released ref2va's
per-reference grids, but area-matched to the target instead of the 2048px
reference short edge. Each reference block advances the rotary media
clock by 1.0.
At sampling, ctrl images are ALWAYS references, never first frames.
"""
arch = "minimax_h3_ref2va"
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# references arrive as a list per sample (multi_control_paths), at
# their ORIGINAL size — this model does its own area-matched resize
self.has_multiple_control_images = True
self.use_raw_control_images = True
def _dit_component(self) -> str:
partition = str(
self.model_config.model_kwargs.get("partition", "ref2va_pruned")
).lower()
if partition not in ("ref2va", "ref2va_pruned"):
raise ValueError(
f"model_kwargs.partition must be ref2va or ref2va_pruned for "
f"{self.arch}, got {partition}"
)
return f"dit_{partition}"
def get_base_model_version(self):
return "minimax_h3_ref2va"
def _build_condition(
self, batch: "DataLoaderBatchDTO", latent_shape, device, dtype
):
# raw-size control tensors in [0, 1]; each is resized to the TARGET's
# pixel area with its own aspect kept, then encoded on its own grid
if batch is None:
return None, (), ()
controls_per_item = None
if batch.control_tensor_list is not None:
controls_per_item = batch.control_tensor_list
elif batch.control_tensor is not None:
controls_per_item = [
[batch.control_tensor[b]] for b in range(batch.control_tensor.shape[0])
]
if not controls_per_item:
return None, (), ()
ref_count = len(controls_per_item[0])
if any(len(c) != ref_count for c in controls_per_item):
raise ValueError(
"ref2va: every item in a batch must have the same number of "
"reference images"
)
_, h_lat, w_lat = latent_shape
target_h, target_w = h_lat * 16, w_lat * 16
all_rows = []
ref_shapes = []
for ref_idx in range(ref_count):
resized = []
for c in controls_per_item:
img = c[ref_idx]
if img.ndim == 4:
img = img[0]
ph, pw = packing.reference_pixel_size(
img.shape[2], img.shape[1], target_h, target_w
)
resized.append(
torch.nn.functional.interpolate(
img[None].to(device, torch.float32),
size=(ph, pw),
mode="bilinear",
antialias=True,
)[0]
)
shapes = {tuple(r.shape[1:]) for r in resized}
if len(shapes) > 1:
raise ValueError(
"ref2va: reference aspect ratios must match across a batch "
f"(got pixel shapes {sorted(shapes)}); use batch_size 1 for "
"mixed-aspect references"
)
frames = torch.stack(resized)
# [0, 1] control -> [-1, 1] pixels, single-frame keyframe encode
ref_latents = self.encode_keyframe_latents(
(frames * 2.0 - 1.0).unsqueeze(2)
)
ref_noise = torch.randn_like(ref_latents)
ref_latents = (
KEYFRAME_NOISE_AUG_T * ref_latents
+ (1.0 - KEYFRAME_NOISE_AUG_T) * ref_noise
)
ref_shapes.append((ref_latents.shape[3], ref_latents.shape[4]))
all_rows.append(patchify_video_latents(ref_latents).to(dtype))
return torch.cat(all_rows, dim=1), (), tuple(ref_shapes)
def generate_single_image(
self,
pipeline: MiniMaxH3Pipeline,
gen_config: GenerateImageConfig,
conditional_embeds: AdvancedPromptEmbeds,
unconditional_embeds: AdvancedPromptEmbeds,
generator: torch.Generator,
extra: dict,
):
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
sc = self.get_bucket_divisibility()
gen_config.width = max(sc, int(gen_config.width // sc * sc))
gen_config.height = max(sc, int(gen_config.height // sc * sc))
is_video = gen_config.num_frames > 1
if is_video:
gen_config.num_frames = packing.align_num_frames_down(gen_config.num_frames)
gen_config.fps = packing.FPS
gen_config.save_image = partial(new_save_image_function, gen_config)
gen_config.log_image = partial(blank_log_image_function, gen_config)
gen_config.output_ext = "mp4"
# every ctrl image is a REFERENCE (never a first frame): resized to
# the target's pixel area with its own aspect kept
ref_images = []
for path in (
gen_config.ctrl_img,
gen_config.ctrl_img_1,
gen_config.ctrl_img_2,
gen_config.ctrl_img_3,
):
if path is not None:
img = Image.open(path).convert("RGB")
ref_images.append(
packing.prepare_reference_image(
img, gen_config.height, gen_config.width
)
)
with_audio = bool(self.model_config.model_kwargs.get("sample_audio", True))
result = pipeline(
conditional_embeds=conditional_embeds,
unconditional_embeds=unconditional_embeds,
height=gen_config.height,
width=gen_config.width,
num_frames=gen_config.num_frames,
num_inference_steps=gen_config.num_inference_steps,
guidance_scale=gen_config.guidance_scale,
latents=gen_config.latents,
generator=generator,
ref_images=ref_images or None,
with_audio=with_audio and is_video,
)
if is_video:
return result # dict consumed by new_save_image_function
return result[0]

71
extensions_built_in/diffusion_models/minimax_h3/src/packing.py
View File

@@ -109,6 +109,31 @@ def resolve_canvas_size(aspect_width: float, aspect_height: float) -> Tuple[int,
return max(m, round(height / m) * m), max(m, round(width / m) * m)
def reference_pixel_size(
ref_width: int, ref_height: int, target_height: int, target_width: int
) -> Tuple[int, int]:
"""Reference images match the TARGET's pixel area while keeping their own
aspect ratio; both axes snap to the canvas multiple. Returns (height, width)."""
scale = math.sqrt((target_height * target_width) / float(ref_width * ref_height))
m = CANVAS_MULTIPLE
height = max(m, round(ref_height * scale / m) * m)
width = max(m, round(ref_width * scale / m) * m)
return height, width
def prepare_reference_image(
image: Image.Image, target_height: int, target_width: int
) -> Image.Image:
"""Resize a reference onto the target's pixel budget, aspect preserved
(up to the /32 snap)."""
height, width = reference_pixel_size(
image.size[0], image.size[1], target_height, target_width
)
if image.size == (width, height):
return image
return image.resize((width, height), Image.Resampling.LANCZOS)
def prepare_keyframe_image(
image: Image.Image, height: int, width: int, stretch: bool = True
):
@@ -254,13 +279,24 @@ def build_packed_sequence(
num_audio_latents: int,
patch_size=(1, 2, 2),
keyframe_anchors: Tuple[str, ...] = (),
image_ref_shapes: Tuple[Tuple[int, int], ...] = (),
) -> PackedLayout:
"""Build the [text | keyframe conditions | target audio | target video]
layout used by t2va and fl2va."""
"""Build the [text | conditions | target audio | target video] layout.
The condition segment holds either fl2va keyframes (``keyframe_anchors``:
rows pinned at the first/last target frame's rotary time, on the target's
spatial grid) or ref2va image references (``image_ref_shapes``: one
``(latent_height, latent_width)`` per reference — references keep their
OWN aspect on their own aspect-normalized grid; each block sits at its
own rotary time and advances the shared media clock by 1.0, so the target
streams start at ``num_text + len(image_ref_shapes)``)."""
if keyframe_anchors and image_ref_shapes:
raise ValueError("keyframe_anchors and image_ref_shapes are mutually exclusive")
_, ph, pw = patch_size
rows_per_frame = (latent_height // ph) * (latent_width // pw)
num_text = int(text_token_tags.shape[0])
num_cond = len(keyframe_anchors) * rows_per_frame
ref_rows = [(h // ph) * (w // pw) for h, w in image_ref_shapes]
num_cond = len(keyframe_anchors) * rows_per_frame + sum(ref_rows)
num_audio_rows = num_audio_latents * AUDIO_CHANNELS
num_video_rows = num_latent_frames * rows_per_frame
seq_len = num_text + num_cond + num_audio_rows + num_video_rows
@@ -270,7 +306,9 @@ def build_packed_sequence(
video_start = audio_start + num_audio_rows
# text rows sit on the time axis at their row index; the media clock
# continues from there, so prompt length shifts the whole media clock
# continues from there (past the reference blocks, which advance it 1.0
# each), so prompt length shifts the whole media clock
media_origin = float(num_text + len(image_ref_shapes))
position_ids = torch.zeros(seq_len, 3, dtype=torch.float64)
position_ids[:num_text, 0] = torch.arange(num_text, dtype=torch.float64)
@@ -301,9 +339,30 @@ def build_packed_sequence(
position_ids[rows, 0] = anchor_time
position_ids[rows, 1:] = frame_grid
ref_cursor = cond_start
for i, (ref_h, ref_w) in enumerate(image_ref_shapes):
# each reference on its own aspect-normalized grid (area-matched to
# the target, so the grids span comparable ranges)
ref_sqrt_area = math.sqrt(ref_h * ref_w)
ref_grid = torch.stack(
[
g.reshape(-1)
for g in torch.meshgrid(
_spatial_position_grid(ref_h, ph, ref_sqrt_area),
_spatial_position_grid(ref_w, pw, ref_sqrt_area),
indexing="ij",
)
],
dim=-1,
)
rows = slice(ref_cursor, ref_cursor + ref_rows[i])
position_ids[rows, 0] = float(num_text + i)
position_ids[rows, 1:] = ref_grid
ref_cursor += ref_rows[i]
# audio rows: channel-major, one rotary unit per latent (40/s = 24fps*5/3),
# no height coordinate, width pinned to the grid extremes per channel
audio_time = float(num_text) + torch.arange(num_audio_latents, dtype=torch.float64)
audio_time = media_origin + torch.arange(num_audio_latents, dtype=torch.float64)
position_ids[audio_start:video_start, 0] = audio_time.repeat(AUDIO_CHANNELS)
position_ids[audio_start:video_start, 2] = torch.cat(
[
@@ -315,7 +374,7 @@ def build_packed_sequence(
)
video_pos = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64)
video_pos[:, :, 0] = _temporal_position_grid(num_latent_frames, float(num_text))[
video_pos[:, :, 0] = _temporal_position_grid(num_latent_frames, media_origin)[
:, None
]
video_pos[:, :, 1:] = frame_grid[None]

29
extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py
View File

@@ -75,6 +75,9 @@ class MiniMaxH3Pipeline:
ctrl_img: Optional[
Image.Image
] = None, # first-frame keyframe, already canvas-sized
ref_images: Optional[
list
] = None, # ref2va references, already area-matched (own aspect, /32)
with_audio: bool = True,
**kwargs,
):
@@ -101,7 +104,13 @@ class MiniMaxH3Pipeline:
token_tags = conditional_embeds.text_token_tags[0].to("cpu", torch.long)
# --- packed layout -------------------------------------------------
if ctrl_img is not None and ref_images:
raise ValueError("ctrl_img (first frame) and ref_images are exclusive")
anchors = ("first",) if ctrl_img is not None else ()
# references keep their own aspect: latent dims come from each image
ref_shapes = tuple(
(img.size[1] // 16, img.size[0] // 16) for img in (ref_images or [])
)
layout = build_packed_sequence(
text_token_tags=token_tags,
num_latent_frames=t_lat,
@@ -109,16 +118,18 @@ class MiniMaxH3Pipeline:
latent_width=w_lat,
num_audio_latents=a_lat,
keyframe_anchors=anchors,
image_ref_shapes=ref_shapes,
)
num_cond = layout.num_condition_video_rows
# --- conditioning rows (draw order: condition noise, video, audio) --
cond_rows = None
if ctrl_img is not None:
def encode_condition_image(img: Image.Image) -> torch.Tensor:
cond_noise = randn_tensor(
(1, 24, 1, h_lat, w_lat), generator=generator, dtype=torch.float32
(1, 24, 1, img.size[1] // 16, img.size[0] // 16),
generator=generator,
dtype=torch.float32,
).to(device)
frame = torch.from_numpy(np.array(ctrl_img)).float()
frame = torch.from_numpy(np.array(img)).float()
frame = (frame / 255.0) * 2.0 - 1.0 # (H, W, 3) -> [-1, 1]
frame = frame.permute(2, 0, 1)[None, :, None] # (1, 3, 1, H, W)
cond_latents = model.encode_keyframe_latents(frame) # (1, 24, 1, h, w) fp32
@@ -127,7 +138,15 @@ class MiniMaxH3Pipeline:
KEYFRAME_NOISE_AUG_T * cond_latents.to(device)
+ (1.0 - KEYFRAME_NOISE_AUG_T) * cond_noise
)
cond_rows = patchify_video_latents(cond_latents) # (1, rows, 96)
return patchify_video_latents(cond_latents) # (1, rows, 96)
cond_rows = None
if ctrl_img is not None:
cond_rows = encode_condition_image(ctrl_img)
elif ref_images:
cond_rows = torch.cat(
[encode_condition_image(img) for img in ref_images], dim=1
)
# --- initial noise -------------------------------------------------
if latents is None: