From 97bf49edadf4644fba6dacfaf5b2037281316d16 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 15 Aug 2026 06:18:09 -0600 Subject: [PATCH] Add support for video references in MiniMax H3 ref2va --- .../diffusion_models/minimax_h3/minimax_h3.py | 211 ++++++++++++++++-- .../minimax_h3/src/packing.py | 90 ++++++-- .../minimax_h3/src/pipeline.py | 55 ++++- .../minimax_h3/src/ref_video_cache.py | 159 +++++++++++++ .../minimax_h3/src/text_encoder.py | 123 +++++++++- extensions_built_in/sd_trainer/SDTrainer.py | 29 ++- toolkit/data_transfer_object/data_loader.py | 8 + toolkit/dataloader_mixins.py | 19 ++ toolkit/models/base_model.py | 27 ++- ui/src/components/SampleControlImage.tsx | 21 +- 10 files changed, 668 insertions(+), 74 deletions(-) create mode 100644 extensions_built_in/diffusion_models/minimax_h3/src/ref_video_cache.py diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 7cac9cd8..0cff7d76 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -62,6 +62,8 @@ from toolkit.util.quantize import get_qtype, quantize, quantize_model from optimum.quanto import freeze from .src import packing + +packing_video_exts = [".mp4", ".avi", ".mov", ".webm", ".mkv", ".wmv", ".m4v", ".flv"] from .src.audio_vae import MiniMaxH3AudioVAE, fold_audio_vae_weight_norm from .src.packing import ( KEYFRAME_ENCODE_SEED, @@ -75,7 +77,13 @@ from .src.packing import ( unpatchify_video_tokens, ) from .src.pipeline import MiniMaxH3Pipeline -from .src.text_encoder import TEXT_ENCODER_LAYER, encode_minimax_h3_prompt +from .src.ref_video_cache import load_ref_video_latent +from .src.text_encoder import ( + TEXT_ENCODER_LAYER, + VideoRef, + encode_minimax_h3_prompt, + load_video_ref, +) from .src.transformer import MiniMaxH3Transformer, MiniMaxH3TransformerParams from .src.vae import MiniMaxH3VideoVAE @@ -664,6 +672,9 @@ class MinimaxH3Model(BaseModel): pil_images.append( Image.fromarray(arr.permute(1, 2, 0).cpu().numpy()) ) + elif isinstance(img, str): + # a control VIDEO path: 2 fps timestamped presentation + pil_images.append(load_video_ref(img)) else: pil_images.append(img) if len(pil_images) == 1: @@ -824,7 +835,7 @@ class MinimaxH3Model(BaseModel): """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`` + ``(cond_rows, cond_audio_rows, keyframe_anchors, ref_blocks)`` 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 @@ -835,7 +846,7 @@ class MinimaxH3Model(BaseModel): and getattr(batch, "num_frames", 1) > 1 ) if not do_i2v: - return None, (), () + return None, None, (), () if batch.first_frame_latents is not None: first_latents = batch.first_frame_latents.to(device, torch.float32) @@ -857,7 +868,7 @@ class MinimaxH3Model(BaseModel): KEYFRAME_NOISE_AUG_T * first_latents + (1.0 - KEYFRAME_NOISE_AUG_T) * cond_noise ) - return patchify_video_latents(first_latents).to(dtype), ("first",), () + return patchify_video_latents(first_latents).to(dtype), None, ("first",), () def get_noise_prediction( self, @@ -896,9 +907,12 @@ class MinimaxH3Model(BaseModel): t_a = 1.0 - sigma_a # --- 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, + cond_audio_rows, + keyframe_anchors, + ref_blocks, + ) = self._build_condition(batch, (t_lat, h_lat, w_lat), device, dtype) # --- audio rows ------------------------------------------------- if batch is not None and getattr(batch, "num_frames", None): @@ -983,7 +997,7 @@ class MinimaxH3Model(BaseModel): latent_width=w_lat, num_audio_latents=a_lat, keyframe_anchors=keyframe_anchors, - image_ref_shapes=image_ref_shapes, + ref_blocks=ref_blocks, ) ) ( @@ -995,14 +1009,20 @@ class MinimaxH3Model(BaseModel): _, ) = pad_layouts_to_batch(layouts) num_cond = layouts[0].num_condition_video_rows + num_cond_audio = layouts[0].num_condition_audio_rows # per-row timesteps: text/video rows at t_v, audio rows at t_a, - # condition rows pinned at max(t_v, 0.999) + # condition rows pinned at max(t_v, 0.999); ref soundtracks clean row_t = t_v.view(-1, 1).expand(-1, token_tags.shape[1]).clone() row_t[:, audio_indices] = t_a.view(-1, 1) if num_cond > 0: cond_t = torch.maximum(t_v, torch.full_like(t_v, KEYFRAME_NOISE_AUG_T)) row_t[:, video_indices[:num_cond]] = cond_t.view(-1, 1) + if num_cond_audio > 0: + row_t[:, audio_indices[:num_cond_audio]] = 1.0 + audio_rows = torch.cat( + [cond_audio_rows.to(audio_rows.dtype), audio_rows], dim=1 + ) # pad text embeds to the batch max length max_text = int(text_indices.shape[0]) @@ -1034,6 +1054,9 @@ class MinimaxH3Model(BaseModel): text_indices=text_indices.to(device), ) + if num_cond_audio > 0: + # reference soundtrack rows are conditioning, not targets + audio_pred = audio_pred[:, num_cond_audio:] if batch is not None and batch.audio_target is not None: # flip to ai-toolkit's noise - clean convention if is_primary_pred: @@ -1199,6 +1222,9 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): # their ORIGINAL size — this model does its own area-matched resize self.has_multiple_control_images = True self.use_raw_control_images = True + # control VIDEOS are cached like dataset items and consumed as + # multi-frame reference blocks + self.supports_video_control_images = True def _dit_component(self) -> str: partition = str( @@ -1220,7 +1246,7 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): # 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, (), () + return None, None, (), () controls_per_item = None if batch.control_tensor_list is not None: controls_per_item = batch.control_tensor_list @@ -1229,7 +1255,11 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): [batch.control_tensor[b]] for b in range(batch.control_tensor.shape[0]) ] if not controls_per_item: - return None, (), () + controls_per_item = ( + [[] for _ in batch.file_items] if batch.file_items else [] + ) + if not controls_per_item: + return None, None, (), () ref_count = len(controls_per_item[0]) if any(len(c) != ref_count for c in controls_per_item): raise ValueError( @@ -1278,7 +1308,151 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): ) 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) + + blocks = [(1, h, w, 0) for h, w in ref_shapes] + audio_rows = [] + self._append_video_ref_blocks( + batch, all_rows, audio_rows, blocks, device, dtype + ) + if not all_rows: + return None, None, (), () + cond_audio = torch.cat(audio_rows, dim=1) if audio_rows else None + return torch.cat(all_rows, dim=1), cond_audio, (), tuple(blocks) + + @torch.no_grad() + def _encode_ref_video_for_sampling(self, path: str, gen_config) -> torch.Tensor: + """Decode a reference video, sample it evenly onto the 17n+5 grid + (capped at the sample's frame count), area-match it to the target with + its own aspect, and encode with the released keyframe recipe. Returns + normalized latents (C, T, h, w).""" + import cv2 + import numpy as np + + cap = cv2.VideoCapture(path) + if not cap.isOpened(): + raise ValueError(f"Could not open reference video {path}") + total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + n = packing.align_num_frames_down(min(total, max(gen_config.num_frames, 5))) + indices = [round(i * (total - 1) / max(n - 1, 1)) for i in range(n)] + frames = [] + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ok, frame = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + cap.release() + if len(frames) != n: + n = packing.align_num_frames_down(max(len(frames), 5)) + frames = frames[:n] + h0, w0 = frames[0].shape[:2] + ph, pw = packing.reference_pixel_size( + w0, h0, gen_config.height, gen_config.width + ) + pixels = torch.from_numpy(np.stack(frames)).float() / 255.0 * 2.0 - 1.0 + pixels = pixels.permute(3, 0, 1, 2)[None] # (1, 3, T, H, W) + pixels = ( + torch.nn.functional.interpolate( + pixels.reshape(1 * 3, len(frames), h0, w0).transpose(0, 1), + size=(ph, pw), + mode="bilinear", + antialias=True, + ) + .transpose(0, 1) + .reshape(1, 3, len(frames), ph, pw) + ) + generator = torch.Generator(device="cpu").manual_seed(KEYFRAME_ENCODE_SEED) + latents = self.video_vae.encode( + pixels.to(self.vae.device, self.video_vae.dtype), + sample=True, + generator=generator, + fp16_round=True, + ) + # soundtrack rides clean when the clip has one (best effort) + audio_rows = None + try: + import torchaudio + + waveform, sample_rate = torchaudio.load(path) + rows = self.encode_audio( + [{"waveform": waveform, "sample_rate": sample_rate}] + )[0] + a_lat = packing.audio_latent_num_frames(n) + audio_rows = self._fit_audio_rows(rows.float(), a_lat) + except Exception: + pass + return {"latent": latents[0].float(), "audio_rows": audio_rows} + + def _append_video_ref_blocks( + self, batch, all_rows, audio_rows, blocks, device, dtype + ): + # control videos get dataset-identical treatment (frame count, fps, + # bucket) and one VAE encode, disk-cached next to the video; the + # resulting latents become multi-frame reference blocks packed after + # the image references + paths_per_item = getattr(batch, "control_video_paths_list", None) + if not paths_per_item: + return + vid_count = len(paths_per_item[0]) + if any(len(v) != vid_count for v in paths_per_item): + raise ValueError( + "ref2va: every item in a batch must have the same number of " + "reference videos" + ) + for ref_idx in range(vid_count): + lats = [] + auds = [] + for per_item in paths_per_item: + entry = load_ref_video_latent( + self, per_item[ref_idx], batch.dataset_config + ) + lats.append(entry["latent"].to(device, torch.float32)) + auds.append(entry.get("audio_rows")) + shapes = {tuple(l.shape) for l in lats} + if len(shapes) > 1: + raise ValueError( + "ref2va: reference video latent shapes must match across a " + f"batch (got {sorted(shapes)}); use batch_size 1" + ) + ref_latents = torch.stack(lats) # (B, C, T, h, w) + ref_noise = torch.randn_like(ref_latents) + ref_latents = ( + KEYFRAME_NOISE_AUG_T * ref_latents + + (1.0 - KEYFRAME_NOISE_AUG_T) * ref_noise + ) + # soundtrack rows ride clean when every item in the batch has them + a_lat = 0 + if all(a is not None for a in auds): + a_lat = packing.audio_latent_num_frames(entry["num_frames"]) + trimmed = [ + self._fit_audio_rows(a.to(device, torch.float32), a_lat) + for a in auds + ] + if len({t.shape for t in trimmed}) == 1: + audio_rows.append(torch.stack(trimmed).to(dtype)) + else: + a_lat = 0 + blocks.append( + ( + ref_latents.shape[2], + ref_latents.shape[3], + ref_latents.shape[4], + a_lat, + ) + ) + all_rows.append(patchify_video_latents(ref_latents).to(dtype)) + + @staticmethod + def _fit_audio_rows(rows: torch.Tensor, a_lat: int) -> torch.Tensor: + """Trim/pad channel-major packed rows (2*T, C) to 2*a_lat rows, + per stereo channel.""" + t = rows.shape[0] // 2 + per_ch = rows.reshape(2, t, rows.shape[-1]) + if t > a_lat: + per_ch = per_ch[:, :a_lat] + elif t < a_lat: + per_ch = torch.nn.functional.pad(per_ch, (0, 0, 0, a_lat - t)) + return per_ch.reshape(2 * a_lat, rows.shape[-1]) def generate_single_image( self, @@ -1304,8 +1478,9 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): 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 + # every ctrl file is a REFERENCE (never a first frame): images are + # resized to the target's pixel area with their own aspect kept; + # videos are dataset-style encoded into multi-frame latent blocks ref_images = [] for path in ( gen_config.ctrl_img, @@ -1313,7 +1488,13 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): gen_config.ctrl_img_2, gen_config.ctrl_img_3, ): - if path is not None: + if path is None: + continue + if os.path.splitext(str(path))[1].lower() in packing_video_exts: + ref_images.append( + self._encode_ref_video_for_sampling(str(path), gen_config) + ) + else: img = Image.open(path).convert("RGB") ref_images.append( packing.prepare_reference_image( diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/packing.py b/extensions_built_in/diffusion_models/minimax_h3/src/packing.py index a0b960f6..48b5d5d0 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/packing.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/packing.py @@ -269,6 +269,7 @@ class PackedLayout: audio_indices: torch.Tensor text_indices: torch.Tensor num_condition_video_rows: int + num_condition_audio_rows: int = 0 def build_packed_sequence( @@ -279,24 +280,32 @@ def build_packed_sequence( num_audio_latents: int, patch_size=(1, 2, 2), keyframe_anchors: Tuple[str, ...] = (), - image_ref_shapes: Tuple[Tuple[int, int], ...] = (), + ref_blocks: Tuple[Tuple[int, int, int], ...] = (), ) -> PackedLayout: """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") + spatial grid) or ref2va references (``ref_blocks``: one + ``(t_lat, latent_height, latent_width)`` per reference — ``t_lat == 1`` + for images. References keep their OWN aspect on their own + aspect-normalized grid; an image block advances the shared media clock by + 1.0, a video block by its temporal span, and the target streams start + after the cumulative advance).""" + if keyframe_anchors and ref_blocks: + raise ValueError("keyframe_anchors and ref_blocks are mutually exclusive") _, ph, pw = patch_size rows_per_frame = (latent_height // ph) * (latent_width // pw) num_text = int(text_token_tags.shape[0]) - ref_rows = [(h // ph) * (w // pw) for h, w in image_ref_shapes] - num_cond = len(keyframe_anchors) * rows_per_frame + sum(ref_rows) + # a ref block is (t_lat, h, w) or (t_lat, h, w, audio_latents): a video + # reference's soundtrack packs as clean audio rows immediately BEFORE its + # own video rows + ref_blocks = tuple(tuple(b) + (0,) * (4 - len(b)) for b in ref_blocks) + ref_vid_rows = [t * (h // ph) * (w // pw) for t, h, w, _ in ref_blocks] + ref_aud_rows = [a * AUDIO_CHANNELS for _, _, _, a in ref_blocks] + num_cond = ( + len(keyframe_anchors) * rows_per_frame + sum(ref_vid_rows) + sum(ref_aud_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 @@ -305,10 +314,15 @@ def build_packed_sequence( audio_start = cond_start + num_cond video_start = audio_start + num_audio_rows + def _block_advance(t, a): + span = 1.0 if t == 1 else _temporal_position_span(t) + return max(span, float(a)) if a else span + # text rows sit on the time axis at their row index; the 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)) + # continues from there past the reference blocks, so prompt length (and + # reference count/length) shifts the whole media clock + media_advance = sum(_block_advance(t, a) for t, _, _, a in ref_blocks) + media_origin = float(num_text) + media_advance position_ids = torch.zeros(seq_len, 3, dtype=torch.float64) position_ids[:num_text, 0] = torch.arange(num_text, dtype=torch.float64) @@ -340,25 +354,50 @@ def build_packed_sequence( position_ids[rows, 1:] = frame_grid ref_cursor = cond_start - for i, (ref_h, ref_w) in enumerate(image_ref_shapes): + ref_clock = float(num_text) + cond_audio_idx = [] + cond_video_idx = [] + if keyframe_anchors: + cond_video_idx.append(torch.arange(cond_start, audio_start)) + ref_cursor = audio_start + for i, (ref_t, ref_h, ref_w, ref_a) in enumerate(ref_blocks): # 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) + w_grid = _spatial_position_grid(ref_w, pw, ref_sqrt_area) 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), + w_grid, 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] + if ref_a: + # soundtrack rows first: channel-major, shared 40/s clock from the + # block origin, width pinned to the ref grid's extremes + a_time = ref_clock + torch.arange(ref_a, dtype=torch.float64) + rows = slice(ref_cursor, ref_cursor + ref_aud_rows[i]) + position_ids[rows, 0] = a_time.repeat(AUDIO_CHANNELS) + position_ids[rows, 2] = torch.cat( + [ + torch.full((ref_a,), float(w_grid[0]), dtype=torch.float64), + torch.full((ref_a,), float(w_grid[-1]), dtype=torch.float64), + ] + ) + cond_audio_idx.append(torch.arange(rows.start, rows.stop)) + ref_cursor += ref_aud_rows[i] + rows_per_ref_frame = ref_grid.shape[0] + block = torch.empty(ref_t, rows_per_ref_frame, 3, dtype=torch.float64) + block[:, :, 0] = _temporal_position_grid(ref_t, ref_clock)[:, None] + block[:, :, 1:] = ref_grid[None] + position_ids[ref_cursor : ref_cursor + ref_vid_rows[i]] = block.reshape(-1, 3) + cond_video_idx.append(torch.arange(ref_cursor, ref_cursor + ref_vid_rows[i])) + ref_cursor += ref_vid_rows[i] + ref_clock += _block_advance(ref_t, ref_a) # 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 @@ -380,10 +419,10 @@ def build_packed_sequence( video_pos[:, :, 1:] = frame_grid[None] position_ids[video_start:] = video_pos.reshape(-1, 3) - video_indices = torch.cat( - [torch.arange(cond_start, audio_start), torch.arange(video_start, seq_len)] - ) - audio_indices = torch.arange(audio_start, video_start) + num_cond_video = sum(int(x.shape[0]) for x in cond_video_idx) + num_cond_audio = sum(int(x.shape[0]) for x in cond_audio_idx) + video_indices = torch.cat(cond_video_idx + [torch.arange(video_start, seq_len)]) + audio_indices = torch.cat(cond_audio_idx + [torch.arange(audio_start, video_start)]) text_indices = torch.arange(num_text) token_tags = torch.empty(seq_len, dtype=torch.long) @@ -398,7 +437,8 @@ def build_packed_sequence( video_indices=video_indices, audio_indices=audio_indices, text_indices=text_indices, - num_condition_video_rows=num_cond, + num_condition_video_rows=num_cond_video, + num_condition_audio_rows=num_cond_audio, ) @@ -419,6 +459,8 @@ def build_row_timesteps( condition_video_timestep ) row_t[layout.audio_indices] = float(audio_timestep) + # reference soundtracks ride clean + row_t[layout.audio_indices[: layout.num_condition_audio_rows]] = 1.0 return row_t diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py index 4aa86ee5..f1bec546 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py @@ -108,9 +108,20 @@ class MiniMaxH3Pipeline: 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 []) - ) + # (PIL images are single-frame blocks; video refs arrive as latent + # tensors (C, T, h, w) already encoded by the caller) + ref_blocks = [] + for r in ref_images or []: + if isinstance(r, dict): + lat = r["latent"] + a = r.get("audio_rows") + a_lat = int(a.shape[0]) // 2 if a is not None else 0 + ref_blocks.append((lat.shape[1], lat.shape[2], lat.shape[3], a_lat)) + elif isinstance(r, torch.Tensor): + ref_blocks.append((r.shape[1], r.shape[2], r.shape[3])) + else: + ref_blocks.append((1, r.size[1] // 16, r.size[0] // 16)) + ref_blocks = tuple(ref_blocks) layout = build_packed_sequence( text_token_tags=token_tags, num_latent_frames=t_lat, @@ -118,7 +129,7 @@ class MiniMaxH3Pipeline: latent_width=w_lat, num_audio_latents=a_lat, keyframe_anchors=anchors, - image_ref_shapes=ref_shapes, + ref_blocks=ref_blocks, ) num_cond = layout.num_condition_video_rows @@ -140,13 +151,36 @@ class MiniMaxH3Pipeline: ) return patchify_video_latents(cond_latents) # (1, rows, 96) + def noise_aug_rows(latents: torch.Tensor) -> torch.Tensor: + cond_noise = randn_tensor( + latents.shape, generator=generator, dtype=torch.float32 + ).to(device) + mixed = ( + KEYFRAME_NOISE_AUG_T * latents.to(device, torch.float32) + + (1.0 - KEYFRAME_NOISE_AUG_T) * cond_noise + ) + return patchify_video_latents(mixed) + cond_rows = None + cond_audio_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 - ) + parts = [] + audio_parts = [] + for r in ref_images: + if isinstance(r, dict): + # pre-encoded video reference (+ optional clean soundtrack) + parts.append(noise_aug_rows(r["latent"][None])) + if r.get("audio_rows") is not None: + audio_parts.append(r["audio_rows"][None].to(device)) + elif isinstance(r, torch.Tensor): + parts.append(noise_aug_rows(r[None])) + else: + parts.append(encode_condition_image(r)) + cond_rows = torch.cat(parts, dim=1) + if audio_parts: + cond_audio_rows = torch.cat(audio_parts, dim=1).float() # --- initial noise ------------------------------------------------- if latents is None: @@ -186,10 +220,13 @@ class MiniMaxH3Pipeline: video_in = video_rows if cond_rows is not None: video_in = torch.cat([cond_rows, video_rows], dim=1) + audio_in = audio_rows + if cond_audio_rows is not None: + audio_in = torch.cat([cond_audio_rows, audio_rows], dim=1) video_pred, audio_pred = transformer( hidden_states=video_in.to(dtype), - audio_hidden_states=audio_rows.to(dtype), + audio_hidden_states=audio_in.to(dtype), encoder_hidden_states=text_embeds[None], row_timesteps=row_t, token_tags=tags, @@ -199,7 +236,7 @@ class MiniMaxH3Pipeline: text_indices=text_indices, ) v_video = video_pred[:, num_cond:].float() - v_audio = audio_pred.float() + v_audio = audio_pred[:, layout.num_condition_audio_rows :].float() denoised_v = video_rows + sv * v_video ratio_v = sv_next / sv diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/ref_video_cache.py b/extensions_built_in/diffusion_models/minimax_h3/src/ref_video_cache.py new file mode 100644 index 00000000..a0e8f614 --- /dev/null +++ b/extensions_built_in/diffusion_models/minimax_h3/src/ref_video_cache.py @@ -0,0 +1,159 @@ +"""Reference-video latents for ref2va, without dataloader machinery. + +A control VIDEO gets the dataset's treatment — num_frames / auto_frame_count, +fps, resolution bucket with center crop — then a single VAE encode whose +result is cached next to the video in ``_latent_cache/`` (keyed like normal +latent caches: file signature + the config values that shape the latent). +Everything is deterministic (even frame spread, no random start) so the cache +is stable; the audio track is encoded alongside when possible (cached for +later use, unused in conditioning for now). +""" + +import base64 +import hashlib +import json +import os + +import cv2 +import numpy as np +import torch +from safetensors.torch import load_file, save_file + +from toolkit.basic import get_quick_signature_string +from toolkit.buckets import get_bucket_for_image_size + + +def _cache_path(path: str, hash_dict: dict) -> str: + latent_dir = os.path.join(os.path.dirname(path), "_latent_cache") + name = os.path.splitext(os.path.basename(path))[0] + hash_input = json.dumps(hash_dict, sort_keys=True).encode("utf-8") + hash_str = ( + base64.urlsafe_b64encode(hashlib.md5(hash_input).digest()) + .decode("ascii") + .replace("=", "") + ) + return os.path.join(latent_dir, f"{name}_{hash_str}.safetensors") + + +@torch.no_grad() +def load_ref_video_latent(model, path: str, dataset_config) -> dict: + """Returns {"latent": (C, T, h, w) cpu tensor, "num_frames": int}, + encoding + disk-caching on first use. ``model`` is the MinimaxH3 model + (used for the VAE, audio encode and the frame-count snapper).""" + mem_cache = getattr(model, "_ref_video_cache", None) + if mem_cache is None: + mem_cache = {} + model._ref_video_cache = mem_cache + if path in mem_cache: + return mem_cache[path] + + cap = cv2.VideoCapture(path) + if not cap.isOpened(): + raise ValueError(f"Could not open reference video {path}") + src_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + src_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + src_fps = cap.get(cv2.CAP_PROP_FPS) or dataset_config.fps + + # dataset-identical frame count + if dataset_config.auto_frame_count: + num_frames = int(total / src_fps * dataset_config.fps) + snapper = model.get_frame_count_snapper() + if snapper is not None: + num_frames = snapper(num_frames) + else: + num_frames = dataset_config.num_frames + + trim_tail = bool( + dataset_config.auto_frame_count and dataset_config.trim_auto_frame_count_tail + ) + hash_dict = { + "signature": get_quick_signature_string(path), + "resolution": dataset_config.resolution, + "num_frames": num_frames, + "fps": dataset_config.fps, + "trim_tail": trim_tail, + "latent_space_version": model.latent_space_version, + "is_ref_video": True, + } + cache_file = _cache_path(path, hash_dict) + if os.path.exists(cache_file): + cap.release() + sd = load_file(cache_file, device="cpu") + entry = { + "latent": sd["latent"], + "num_frames": int(sd["num_frames"].item()), + "audio_rows": sd.get("audio_latent"), + } + mem_cache[path] = entry + return entry + + # dataset-identical bucket sizing (center crop, no random) + bucket = get_bucket_for_image_size( + src_w, + src_h, + resolution=dataset_config.resolution, + divisibility=dataset_config.bucket_tolerance, + ) + scale = max(bucket["width"] / src_w, bucket["height"] / src_h) + scale_w, scale_h = int(np.ceil(src_w * scale)), int(np.ceil(src_h * scale)) + crop_x = (scale_w - bucket["width"]) // 2 + crop_y = (scale_h - bucket["height"]) // 2 + + if trim_tail: + # dataset trim mode: real-time pacing from the start, tail trimmed — + # keeps motion speed honest and the soundtrack in sync + fps_ratio = src_fps / dataset_config.fps if src_fps > 0 else 1.0 + indices = [min(round(i * fps_ratio), total - 1) for i in range(num_frames)] + else: + # deterministic even frame spread across the clip + indices = [ + min(round(i * (total - 1) / max(num_frames - 1, 1)), total - 1) + for i in range(num_frames) + ] + frames = [] + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ok, frame = cap.read() + if not ok: + raise ValueError(f"Could not read frame {idx} of {path}") + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = cv2.resize(frame, (scale_w, scale_h), interpolation=cv2.INTER_AREA) + frame = frame[ + crop_y : crop_y + bucket["height"], crop_x : crop_x + bucket["width"] + ] + frames.append(frame) + cap.release() + + pixels = torch.from_numpy(np.stack(frames)).float() / 255.0 * 2.0 - 1.0 + pixels = pixels.permute(0, 3, 1, 2) # (T, C, H, W), [-1, 1] + latent = model.encode_images([pixels])[0].to("cpu", torch.float16) + + state_dict = { + "latent": latent, + "num_frames": torch.tensor(num_frames, dtype=torch.int64), + } + # the soundtrack rides as clean condition rows; best effort (no track = None) + audio_rows = None + try: + import torchaudio + + waveform, sample_rate = torchaudio.load(path) + if trim_tail: + # frames cover [0, num_frames / fps) seconds; trim the soundtrack + # to the same window so it stays in sync with the sampled frames + keep = int(round(num_frames / dataset_config.fps * sample_rate)) + waveform = waveform[:, :keep] + audio_latent = model.encode_audio( + [{"waveform": waveform, "sample_rate": sample_rate}] + )[0] + audio_rows = audio_latent.to("cpu", torch.float16) + state_dict["audio_latent"] = audio_rows + except Exception: + pass + + os.makedirs(os.path.dirname(cache_file), exist_ok=True) + save_file(state_dict, cache_file) + entry = {"latent": latent, "num_frames": num_frames, "audio_rows": audio_rows} + mem_cache[path] = entry + return entry diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py b/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py index 58613c45..6eaef6b5 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py @@ -19,11 +19,60 @@ from typing import List, Optional import torch +from dataclasses import dataclass, field + from .packing import TEXT_TAG, VIDEO_TAG TEXT_ENCODER_LAYER = 50 +@dataclass +class VideoRef: + """A reference VIDEO for the Qwen3-VL presentation: frames sampled at + 2 fps with their timestamps (seconds). Presented ComfyUI-style as + ``