File size: 8,739 Bytes
22a73d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | #!/usr/bin/env python3
"""Repository-local MiniMax-H3 video resampling and RoPE time helpers.
The video normalizer is a drop-in alternative to
``MiniMaxH3Ref2VASetupStep._normalize_video_condition``. It retains the
upstream input conversion, output-length rounding, truncation and LANCZOS
resize policy. Its default ``zero`` mode anchors the target frame grid at
media time zero; ``legacy_round`` reproduces Diffusers' current slot-rounding
behavior for compatibility.
The upstream implementation is Apache-2.0 licensed and lives in
``diffusers/modular_pipelines/minimax_h3/before_encoder.py``.
"""
from __future__ import annotations
import math
from typing import Literal
import numpy as np
import torch
from PIL import Image
from diffusers.modular_pipelines.minimax_h3.modular_pipeline import resolve_canvas_size
ResampleStartMode = Literal["zero", "legacy_round"]
ROPE_TIME_UNITS_PER_SECOND = 40.0
def resample_frame_indices(
num_source_frames: int,
source_fps: float,
target_fps: float,
start_mode: ResampleStartMode = "zero",
) -> np.ndarray:
"""Return source indices on either a zero-origin or legacy rounded grid."""
if num_source_frames < 0:
raise ValueError(f"num_source_frames must be non-negative, got {num_source_frames}.")
if source_fps <= 0 or target_fps <= 0:
raise ValueError(
f"source_fps and target_fps must be positive, got {source_fps} and {target_fps}."
)
if start_mode not in ("zero", "legacy_round"):
raise ValueError(f"Unknown resampling start mode {start_mode!r}.")
count = math.floor(num_source_frames * target_fps / source_fps + 0.5)
if count == 0:
return np.empty(0, dtype=np.int64)
if start_mode == "legacy_round":
# This is Diffusers' current implementation expressed as indices:
# every source frame is repeated by the change in its rounded output
# slot. For 24 -> 8 it retains 1,4,7,...; for 24 -> 12 it retains
# 0,2,4,... due to the half-up boundary.
scale = target_fps / source_fps
slots = np.floor(np.arange(num_source_frames) * scale + 0.5).astype(np.int64)
repeats = np.diff(slots, append=count)
return np.repeat(np.arange(num_source_frames, dtype=np.int64), repeats)
indices = np.floor(
np.arange(count, dtype=np.float64) * source_fps / target_fps + 1e-9
).astype(np.int64)
return np.minimum(indices, num_source_frames - 1)
def zero_origin_frame_indices(
num_source_frames: int,
source_fps: float,
target_fps: float,
) -> np.ndarray:
"""Return source indices for a target grid beginning at exactly ``t=0``.
The output count deliberately matches Diffusers' current stream-end
rounding. Each target timestamp ``k / target_fps`` selects the source
frame whose interval contains it, namely
``floor(k * source_fps / target_fps)``.
"""
return resample_frame_indices(num_source_frames, source_fps, target_fps, "zero")
def sampling_time_origin_seconds(indices: np.ndarray, source_fps: float) -> float:
"""Return the media timestamp represented by the first sampled frame."""
indices = np.asarray(indices)
if indices.ndim != 1 or len(indices) == 0:
raise ValueError("indices must be a non-empty one-dimensional array.")
if source_fps <= 0:
raise ValueError(f"source_fps must be positive, got {source_fps}.")
return float(indices[0]) / float(source_fps)
def normalize_video_condition(
frames,
fps: float,
num_frames: int,
canvas_multiple: int,
canvas_short_edge: int,
canvas_max_pixels: int,
target_fps: float,
start_mode: ResampleStartMode = "zero",
) -> np.ndarray:
"""Normalize a MiniMax-H3 video reference on a selectable time grid.
The signature intentionally matches Diffusers' private normalization
helper for its first seven arguments. The repository-local eighth argument
selects zero-origin behavior by default or exact upstream compatibility.
"""
if isinstance(frames, list):
frames = np.stack([np.asarray(frame.convert("RGB")) for frame in frames])
if isinstance(frames, torch.Tensor):
frames = frames.movedim(-3, -1).cpu().numpy()
frames = np.asarray(frames)
if frames.dtype != np.uint8:
frames = (frames * 255.0).round().clip(0, 255).astype(np.uint8)
if frames.ndim != 4 or frames.shape[3] != 3:
raise ValueError(
"A reference video must be `(num_frames, height, width, 3)` RGB frames, "
f"got {tuple(frames.shape)}."
)
if fps <= 0 or target_fps <= 0:
raise ValueError(f"Video frame rates must be positive, got {fps} and {target_fps}.")
if num_frames <= 0:
raise ValueError(f"num_frames must be positive, got {num_frames}.")
if fps != target_fps:
indices = resample_frame_indices(frames.shape[0], fps, target_fps, start_mode)
frames = frames[indices]
frames = frames[:num_frames]
if len(frames) == 0:
raise ValueError("Video resampling produced no frames.")
height, width = resolve_canvas_size(
frames.shape[2],
frames.shape[1],
canvas_multiple,
canvas_short_edge,
canvas_max_pixels,
)
if frames.shape[1:3] == (height, width):
return frames
return np.stack(
[np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS)) for frame in frames]
)
def normalize_video_condition_zero_origin(
frames,
fps: float,
num_frames: int,
canvas_multiple: int,
canvas_short_edge: int,
canvas_max_pixels: int,
target_fps: float,
) -> np.ndarray:
"""Backward-compatible convenience wrapper fixed to zero-origin mode."""
return normalize_video_condition(
frames,
fps,
num_frames,
canvas_multiple,
canvas_short_edge,
canvas_max_pixels,
target_fps,
start_mode="zero",
)
def transform_reference_video_rope_time(
position_ids: torch.Tensor,
condition_video_indices: torch.Tensor,
*,
video_span_rope_units: float,
time_scale: float = 1.0,
time_origin_seconds: float = 0.0,
audio_span_rope_units: float = 0.0,
following_start_index: int | None = None,
rope_units_per_second: float = ROPE_TIME_UNITS_PER_SECOND,
) -> torch.Tensor:
"""Apply explicit scale and media-time origin to one reference video.
Audio coordinates are intentionally untouched. If rows follow the
reference block, ``following_start_index`` shifts their time coordinates
by the change in ``max(audio_end, video_end)`` so subsequent blocks begin
after the corrected reference end.
"""
if position_ids.ndim != 2 or position_ids.shape[1] != 3:
raise ValueError(f"position_ids must have shape (rows, 3), got {tuple(position_ids.shape)}.")
if condition_video_indices.ndim != 1 or condition_video_indices.numel() == 0:
raise ValueError("condition_video_indices must be a non-empty one-dimensional tensor.")
values = (video_span_rope_units, time_scale, time_origin_seconds, audio_span_rope_units, rope_units_per_second)
if not all(math.isfinite(float(value)) for value in values):
raise ValueError("RoPE time parameters must be finite.")
if video_span_rope_units < 0 or audio_span_rope_units < 0:
raise ValueError("Video and audio spans must be non-negative.")
if time_scale <= 0 or rope_units_per_second <= 0:
raise ValueError("time_scale and rope_units_per_second must be positive.")
if following_start_index is not None and not 0 <= following_start_index <= len(position_ids):
raise ValueError(f"following_start_index is out of range: {following_start_index}.")
result = position_ids.clone()
origin = result[condition_video_indices[0], 0].clone()
phase = float(time_origin_seconds) * float(rope_units_per_second)
old_time = result[condition_video_indices, 0]
result[condition_video_indices, 0] = origin + (old_time - origin) * time_scale + phase
old_block_span = max(float(audio_span_rope_units), float(video_span_rope_units))
new_video_span = phase + float(video_span_rope_units) * float(time_scale)
new_block_span = max(float(audio_span_rope_units), new_video_span)
if following_start_index is not None:
result[following_start_index:, 0] += new_block_span - old_block_span
return result
__all__ = [
"ROPE_TIME_UNITS_PER_SECOND",
"ResampleStartMode",
"normalize_video_condition",
"normalize_video_condition_zero_origin",
"resample_frame_indices",
"sampling_time_origin_seconds",
"transform_reference_video_rope_time",
"zero_origin_frame_indices",
]
|