"""Packed-sequence geometry for MiniMax-H3. One transformer forward runs over a single packed 1-D sequence: [ text (L) | keyframe conditions (C) | target audio (A) | target video (V) ] This module owns everything needed to place a row in that sequence and give it its (t, h, w) rotary coordinate, plus the sigma-shift math that couples the video (shift 12) and audio (shift 3) flow schedules. Rotary coordinates are built in float64 and with numpy's ``linspace`` because video and audio share one 40-units-per-second rotary clock (video advances 5/3 units per pixel frame at 24 fps, audio one unit per latent at 40/s) and that shared clock is the released checkpoint's audio/video alignment. """ import math from dataclasses import dataclass from typing import List, Optional, Tuple import numpy as np import torch from PIL import Image # Per-row modality tags — these index the transformer's AdaLN table, so the # values are a checkpoint contract. VIDEO_TAG = 0 TEXT_TAG = 1 AUDIO_TAG = 2 PAD_TAG = -1 FPS = 24 SHORT_EDGE = 768 MAX_PIXELS = 768 * 1344 CANVAS_MULTIPLE = 32 # video VAE: 17 pixel frames per chunk -> 5 latent frames, 3 trailing latents # dropped overall, so 17n+5 pixel frames <-> 5n+2 latent frames FRAMES_PER_CHUNK = 17 LATENTS_PER_CHUNK = 5 AUDIO_LATENTS_PER_SECOND = 40 AUDIO_CHANNELS = 2 AUDIO_SAMPLE_RATE = 32000 # released flow shifts (exponential): video 12, audio 3 VIDEO_SIGMA_SHIFT = 12.0 AUDIO_SIGMA_SHIFT = 3.0 # keyframe conditioning rows are noised to t = 0.999 and pinned there; the # posterior sample of the keyframe VAE encode uses a fixed seed of 42 KEYFRAME_NOISE_AUG_T = 0.999 KEYFRAME_ENCODE_SEED = 42 # rotary-time constants: one latent frame spans 5/3 * frames_per_latent units, # the (1, 4, 4, 4, 4) pattern mirroring the VAE's 17 -> 5 frame grouping _ROPE_FRAME_RESCALE = 5.0 / 3.0 _ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) _ROPE_SPATIAL_SCALE = 32 # --------------------------------------------------------------------------- # Frame / canvas arithmetic # --------------------------------------------------------------------------- def align_num_frames(num_frames: int) -> int: """Snap a frame count UP to the next 17n+5 the video VAE can encode.""" if num_frames < 1: raise ValueError(f"num_frames must be positive, got {num_frames}") while num_frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK: num_frames += 1 return num_frames def align_num_frames_down(num_frames: int) -> int: """Snap a frame count DOWN to the previous 17n+5 (minimum 5).""" num_frames = max(num_frames, 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: """17n+5 pixel frames -> 5n+2 latent frames.""" if num_frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK: raise ValueError(f"num_frames must be of the form 17n+5, 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: """Audio latents covering `num_frames` video frames at 24 fps / 40 Hz.""" return int(round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND)) def resolve_canvas_size(aspect_width: float, aspect_height: float) -> Tuple[int, int]: """Aspect ratio -> (height, width): short edge 768, area capped at 768*1344, both axes rounded to the nearest multiple of 32.""" ratio = aspect_width / aspect_height if ratio >= 1.0: width, height = SHORT_EDGE * ratio, float(SHORT_EDGE) else: width, height = float(SHORT_EDGE), SHORT_EDGE / ratio area = width * height if area > MAX_PIXELS: scale = (MAX_PIXELS / area) ** 0.5 width, height = width * scale, height * scale m = CANVAS_MULTIPLE 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 IMAGE sizing (ComfyUI 'match'): aspect-preserving scale DOWN ONLY to the target's pixel area — never upscaled; both axes snap to the canvas multiple. Returns (height, width).""" scale = min( 1.0, 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 reference_video_pixel_size( ref_width: int, ref_height: int, target_height: int, target_width: int ) -> Tuple[int, int]: """Reference VIDEO sizing: match the TARGET's pixel area with the ref's own aspect kept (same aspect -> exactly the target size; other aspects -> the same pixel budget on the /32 grid). 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 ): """Put a keyframe onto the target canvas: the geometry anchor is stretched, a follower keyframe is cover-cropped.""" 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)) # --------------------------------------------------------------------------- # Row packing # --------------------------------------------------------------------------- def patchify_video_latents(latents: torch.Tensor, patch_size=(1, 2, 2)) -> torch.Tensor: """(B, C, T, H, W) -> (B, N, C * prod(patch)) rows, frame-major then row-major, feature order [c, pt, ph, pw].""" pt, ph, pw = patch_size b, c, t, h, w = latents.shape latents = latents.reshape(b, c, t // pt, pt, h // ph, ph, w // pw, pw) latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7) return latents.reshape(b, -1, c * pt * ph * pw).contiguous() def unpatchify_video_tokens( rows: torch.Tensor, num_latent_frames: int, latent_height: int, latent_width: int, channels: int = 24, patch_size=(1, 2, 2), ) -> torch.Tensor: """(B, N, C * prod(patch)) -> (B, C, T, H, W). Inverse of patchify.""" pt, ph, pw = patch_size b = rows.shape[0] rows = rows.reshape( b, num_latent_frames // pt, latent_height // ph, latent_width // pw, channels, pt, ph, pw, ) rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7) return rows.reshape( b, channels, num_latent_frames, latent_height, latent_width ).contiguous() def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor: """(B, 2, C, T) stereo audio latents -> (B, 2*T, C) channel-major rows (all T frames of channel 0, then channel 1).""" return ( latents.permute(0, 1, 3, 2) .reshape(latents.shape[0], -1, latents.shape[2]) .contiguous() ) def unpack_audio_tokens(rows: torch.Tensor, num_audio_latents: int) -> torch.Tensor: """(B, 2*T, C) channel-major rows -> (B, 2, C, T).""" b, _, c = rows.shape rows = rows.reshape(b, AUDIO_CHANNELS, num_audio_latents, c) return rows.permute(0, 1, 3, 2).contiguous() # --------------------------------------------------------------------------- # Rotary grids (float64, numpy linspace — the released grid must reproduce) # --------------------------------------------------------------------------- def _spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: ratio = dim / sqrt_area left = (1.0 - ratio) / 2.0 # numpy linspace(endpoint=False) is start + arange(n) * (stop-start)/n, # which is not bit-identical to torch.linspace grid = ( np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE ) return torch.from_numpy(grid).to(torch.float64) def _temporal_position_grid(num_latent_frames: int, origin: float) -> torch.Tensor: spans = torch.tensor( [ _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] for i in range(num_latent_frames) ], dtype=torch.float64, ) return origin + torch.cat( [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)] ) def _temporal_position_span(num_latent_frames: int) -> float: # numpy pairwise sum on purpose: the reference computes the "last" keyframe # anchor this way and the summation orders differ in the last ulp spans = np.ones(num_latent_frames, dtype=np.float64) * _ROPE_FRAME_RESCALE for i in range(len(_ROPE_FRAMES_PER_LATENT)): spans[i :: len(_ROPE_FRAMES_PER_LATENT)] *= _ROPE_FRAMES_PER_LATENT[i] return float(spans.sum()) # --------------------------------------------------------------------------- # Sequence layout # --------------------------------------------------------------------------- @dataclass class PackedLayout: """Structural description of one packed sequence (one batch item).""" sequence_length: int position_ids: torch.Tensor # (S, 3) float64 token_tags: torch.Tensor # (S,) long video_indices: torch.Tensor # condition rows first, then target rows audio_indices: torch.Tensor text_indices: torch.Tensor num_condition_video_rows: int num_condition_audio_rows: int = 0 def build_packed_sequence( text_token_tags: torch.Tensor, # (L,) long: 1 text, 0 for vision-block rows num_latent_frames: int, latent_height: int, latent_width: int, num_audio_latents: int, patch_size=(1, 2, 2), keyframe_anchors: Tuple[str, ...] = (), 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 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]) # 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 cond_start = num_text 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, 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) sqrt_area = math.sqrt(latent_height * latent_width) height_grid = _spatial_position_grid(latent_height, ph, sqrt_area) width_grid = _spatial_position_grid(latent_width, pw, sqrt_area) frame_grid = torch.stack( [g.reshape(-1) for g in torch.meshgrid(height_grid, width_grid, indexing="ij")], dim=-1, ) for i, anchor in enumerate(keyframe_anchors): if anchor == "first": anchor_time = float(num_text) elif anchor == "last": anchor_time = ( float(num_text) + _temporal_position_span(num_latent_frames) - _ROPE_FRAME_RESCALE ) else: raise ValueError( f"keyframe anchor must be 'first' or 'last', got {anchor!r}" ) rows = slice( cond_start + i * rows_per_frame, cond_start + (i + 1) * rows_per_frame ) position_ids[rows, 0] = anchor_time position_ids[rows, 1:] = frame_grid ref_cursor = cond_start 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), w_grid, indexing="ij", ) ], dim=-1, ) 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 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( [ torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), torch.full( (num_audio_latents,), float(width_grid[-1]), dtype=torch.float64 ), ] ) video_pos = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) video_pos[:, :, 0] = _temporal_position_grid(num_latent_frames, media_origin)[ :, None ] video_pos[:, :, 1:] = frame_grid[None] position_ids[video_start:] = video_pos.reshape(-1, 3) 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) token_tags[text_indices] = text_token_tags.to(torch.long) token_tags[audio_indices] = AUDIO_TAG token_tags[video_indices] = VIDEO_TAG return PackedLayout( sequence_length=seq_len, 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_cond_video, num_condition_audio_rows=num_cond_audio, ) def build_row_timesteps( layout: PackedLayout, video_timestep: float, audio_timestep: float, condition_video_timestep: Optional[float] = None, ) -> torch.Tensor: """Per-row timestep values (S,) float32. Text rows inherit the video timestep; condition video rows stay pinned at their noise-aug level.""" if condition_video_timestep is None: condition_video_timestep = max(video_timestep, KEYFRAME_NOISE_AUG_T) row_t = torch.full( (layout.sequence_length,), float(video_timestep), dtype=torch.float32 ) row_t[layout.video_indices[: layout.num_condition_video_rows]] = float( 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 def pad_layouts_to_batch(layouts: List[PackedLayout]): """Stack per-item layouts that share the same media geometry but may have different text lengths into batched transformer inputs. Items are right-padded in the TEXT segment to the batch's max text length (pad rows tagged -1, masked out of attention as keys, positions zero). Returns (position_ids (B, S, 3) f64, token_tags (B, S) long, video_indices, audio_indices, text_indices, pad_counts) where the index tensors describe the shared padded layout: [text_max | cond | audio | video]. """ max_text = max(int(l.text_indices.shape[0]) for l in layouts) ref = layouts[0] media_len = ref.sequence_length - int(ref.text_indices.shape[0]) for l in layouts: if l.sequence_length - int(l.text_indices.shape[0]) != media_len: raise ValueError( "all layouts in a batch must share the same media geometry" ) seq_len = max_text + media_len b = len(layouts) position_ids = torch.zeros(b, seq_len, 3, dtype=torch.float64) token_tags = torch.full((b, seq_len), PAD_TAG, dtype=torch.long) pad_counts = [] for i, l in enumerate(layouts): lt = int(l.text_indices.shape[0]) position_ids[i, :lt] = l.position_ids[:lt] position_ids[i, max_text:] = l.position_ids[lt:] token_tags[i, :lt] = l.token_tags[:lt] token_tags[i, max_text:] = l.token_tags[lt:] pad_counts.append(max_text - lt) offset = max_text - int(ref.text_indices.shape[0]) video_indices = ref.video_indices + offset audio_indices = ref.audio_indices + offset text_indices = torch.arange(max_text) return ( position_ids, token_tags, video_indices, audio_indices, text_indices, pad_counts, ) # --------------------------------------------------------------------------- # Sigma-shift math (video shift 12, audio shift 3, exponential) # --------------------------------------------------------------------------- def shift_sigma(sigma, shift: float): """Exponential timeshift: shift * sigma / (1 + (shift - 1) * sigma).""" return shift * sigma / (1.0 + (shift - 1.0) * sigma) def pdd_span(sigma_v: float, n: int, video_shift: float, sigma_v_next=None): """Head interval [start, stop) on the unshifted n-interval grid. A sampling step spans sigma -> sigma_next (ComfyUI's rounding); a single training sigma uses the one interval containing it.""" base_t = 1.0 - float(shift_sigma(sigma_v, 1.0 / video_shift)) if sigma_v_next is None: start = min(max(int(math.floor(base_t * n)), 0), n - 1) return start, start + 1 base_t_next = 1.0 - float(shift_sigma(sigma_v_next, 1.0 / video_shift)) start = min(round(base_t * n), n - 1) stop = max(round(base_t_next * n), start + 1) return start, stop def pdd_head_weights(start: int, stop: int, n: int, flow_shift: float): """(n,) weights over the stacked heads: head 0 (the full head) at 1, the offset heads in [max(start, 1), stop) at their normalized dt share.""" grid = torch.linspace(1.0, 0.0, n + 1, dtype=torch.float64) dt = (1.0 - shift_sigma(grid, flow_shift)).diff()[start:stop] w = torch.zeros(n, dtype=torch.float64) w[0] = 1.0 first = max(start, 1) w[first:stop] = (dt / dt.sum())[first - start :] return w def build_pdd_weights( sigma_v, # (B,) video sigmas n: int, video_shift: float, audio_shift: float = AUDIO_SIGMA_SHIFT, sigma_v_next=None, # (B,) or None for single-sigma training ): """Per-sample (video, audio) head weights, each (B, n) float32.""" wv, wa = [], [] for i in range(len(sigma_v)): nxt = None if sigma_v_next is None else float(sigma_v_next[i]) start, stop = pdd_span(float(sigma_v[i]), n, video_shift, nxt) wv.append(pdd_head_weights(start, stop, n, video_shift)) wa.append(pdd_head_weights(start, stop, n, audio_shift)) return torch.stack(wv).float(), torch.stack(wa).float() def remap_sigma( sigma, from_shift: float = VIDEO_SIGMA_SHIFT, to_shift: float = AUDIO_SIGMA_SHIFT ): """Map a sigma on the `from_shift` schedule onto the `to_shift` schedule at the same underlying schedule position (the video/audio coupling).""" base = sigma / (from_shift + sigma * (1.0 - from_shift)) return shift_sigma(base, to_shift) def build_sigma_schedule( num_inference_steps: int, shift: float = VIDEO_SIGMA_SHIFT, t1000_ladder: bool = False, ) -> torch.Tensor: """The released sampling grid: linspace(1, 0, steps + 1) through the exponential shift, consecutive duplicates collapsed — `steps` yields `steps` model evaluations (the released repo counts the terminal 0 in `steps`; we don't, so sample_steps means model evals). ``t1000_ladder`` uses FastH3's trained ladder instead: rounded indices on the shared 1000-step grid ([999, 874, ..., 250, 125] at 8 steps), one forward per entry, each scheduler applying its own shift.""" if t1000_ladder: base = ( torch.linspace(0.0, 999.0, num_inference_steps + 1, dtype=torch.float32) .round() .flip(0) / 1000.0 ) else: base = torch.linspace(1.0, 0.0, num_inference_steps + 1, dtype=torch.float32) sigmas = shift_sigma(base, shift) return torch.unique_consecutive(sigmas)