File size: 11,887 Bytes
bc50ea0 | 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 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 | """Pure tensor/media helpers for general MiniMax H3 masked targets.
The 32-pixel grid mirrors H3's 16x video VAE scale followed by the DiT's 2x2
latent patching. A mask value of one means generate and zero means preserve.
"""
from __future__ import annotations
from typing import Any
import torch
import torch.nn.functional as F
H3_FRAME_OFFSET = 5
H3_FRAME_STRIDE = 17
H3_VIDEO_FPS = 24
H3_VAE_SCALE = 16
H3_PIXEL_CELL = 32
H3_FRAMES_PER_VIDEO_LATENT = (1, 4, 4, 4, 4)
MASK_CONVERSION_MODES = (
"H3 exact (causal/token max)",
"legacy trilinear",
)
GRID_MODES = (
"runtime exact (latent max)",
"any pixel coverage",
"50% pixel coverage",
"full pixel coverage",
)
def floor_h3_frame_count(frame_count: int) -> int:
"""Return the largest valid 17k+5 frame count not above ``frame_count``."""
count = int(frame_count)
if count < H3_FRAME_OFFSET:
raise ValueError(
"MiniMax H3 needs at least %d video frames; received %d." %
(H3_FRAME_OFFSET, count)
)
return H3_FRAME_OFFSET + (
(count - H3_FRAME_OFFSET) // H3_FRAME_STRIDE
) * H3_FRAME_STRIDE
def trim_audio_to_frames(
audio: Any,
frame_count: int,
fps: int = H3_VIDEO_FPS,
):
"""Trim ComfyUI AUDIO to the exact video duration without padding."""
if audio is None:
return None, None
if not isinstance(audio, dict):
raise ValueError(
"H3 source AV trim needs a ComfyUI AUDIO value when connected."
)
waveform = audio.get("waveform")
sample_rate = audio.get("sample_rate")
if waveform is None or not hasattr(waveform, "shape") or waveform.ndim < 2:
raise ValueError("Connected AUDIO has no valid waveform tensor.")
try:
resolved_rate = int(sample_rate)
except (TypeError, ValueError) as exc:
raise ValueError("Connected AUDIO has no valid sample_rate.") from exc
if resolved_rate < 1:
raise ValueError("Connected AUDIO sample_rate must be positive.")
target_samples = round(int(frame_count) * resolved_rate / int(fps))
trimmed_samples = min(int(waveform.shape[-1]), target_samples)
output = audio.copy()
output["waveform"] = waveform[..., :trimmed_samples]
return output, target_samples
def normalize_comfy_mask(mask: torch.Tensor) -> torch.Tensor:
"""Normalize common Comfy MASK spellings to float ``[frames,H,W]``."""
if not isinstance(mask, torch.Tensor):
raise ValueError("H3 masking expects a ComfyUI MASK tensor.")
if mask.ndim == 2:
mask = mask.unsqueeze(0)
elif mask.ndim == 4 and int(mask.shape[1]) == 1:
mask = mask[:, 0]
elif mask.ndim == 4 and int(mask.shape[-1]) == 1:
mask = mask[..., 0]
if mask.ndim != 3:
raise ValueError(
"Expected a ComfyUI MASK [frames,H,W], got %s." %
(list(mask.shape),)
)
if not int(mask.shape[0]) or not int(mask.shape[1]) or not int(mask.shape[2]):
raise ValueError("H3 masking received an empty mask tensor.")
return mask.to(torch.float32).clamp(0.0, 1.0)
def resize_mask_to_canvas(
mask: torch.Tensor,
frames: int,
height: int,
width: int,
) -> torch.Tensor:
"""Resize a Comfy mask to ``[frames,1,height,width]``."""
if frames < 1 or height < 1 or width < 1:
raise ValueError("Mask canvas dimensions must be positive.")
work = normalize_comfy_mask(mask)[None, None]
work = F.interpolate(
work,
size=(int(frames), int(height), int(width)),
mode="trilinear",
align_corners=False,
)
return work[0].movedim(0, 1).contiguous()
def resize_mask_to_video_latent(
mask: torch.Tensor,
video: torch.Tensor,
) -> torch.Tensor:
"""Legacy trilinear conversion to the target H3 video latent dimensions."""
if video.ndim != 5:
raise ValueError(
"H3 target video latent must be [B,C,T,H,W], got %s." %
(list(video.shape),)
)
if int(video.shape[0]) != 1:
raise ValueError("H3 masked targets currently support batch size 1.")
work = normalize_comfy_mask(mask).to(device=video.device)[None, None]
work = F.interpolate(
work,
size=(int(video.shape[2]), int(video.shape[3]), int(video.shape[4])),
mode="trilinear",
align_corners=False,
)
return work.to(torch.float32).contiguous()
def h3_pixel_frames_for_video_latents(latent_steps: int) -> int:
"""Return the pixel-frame span represented by H3 video latent steps."""
steps = int(latent_steps)
if steps < 1:
raise ValueError("H3 video latent length must be positive.")
pattern = H3_FRAMES_PER_VIDEO_LATENT
complete, trailing = divmod(steps, len(pattern))
return complete * sum(pattern) + sum(pattern[:trailing])
def h3_video_latents_for_pixel_frames(pixel_frames: int) -> int:
"""Return H3 video steps for an exact 1-frame or 17k+5 pixel run."""
frames = int(pixel_frames)
if frames == 1:
return 1
if frames < H3_FRAME_OFFSET or (
frames - H3_FRAME_OFFSET) % H3_FRAME_STRIDE:
raise ValueError(
"H3 exact mask conversion needs a 1-frame image or a 17k+5 "
"video run; received %d frames." % frames)
return 2 + 5 * ((frames - H3_FRAME_OFFSET) // H3_FRAME_STRIDE)
def h3_video_latent_frame_groups(
latent_steps: int,
) -> tuple[tuple[int, int], ...]:
"""Return the exact causal pixel-frame range feeding each H3 latent step."""
groups = []
start = 0
pattern = H3_FRAMES_PER_VIDEO_LATENT
for index in range(int(latent_steps)):
end = start + pattern[index % len(pattern)]
groups.append((start, end))
start = end
if not groups:
raise ValueError("H3 video latent length must be positive.")
return tuple(groups)
def _h3_token_snap_max(mask: torch.Tensor) -> torch.Tensor:
"""Unify each 2x2 H3 latent-pixel token using conservative max coverage."""
if mask.ndim != 3:
raise ValueError(
"H3 token snapping expects [frames,H,W], got %s." %
(list(mask.shape),)
)
height, width = int(mask.shape[-2]), int(mask.shape[-1])
work = F.pad(
mask[:, None],
(0, -width % 2, 0, -height % 2),
mode="replicate",
)
work = F.max_pool2d(work, kernel_size=2, stride=2)
work = work.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
return work[:, 0, :height, :width].contiguous()
def reduce_h3_mask_to_video_latent(
mask: torch.Tensor,
video: torch.Tensor,
) -> torch.Tensor:
"""Map a pixel mask to H3's causal VAE groups and 2x2 token grid.
A tracked mask must contain exactly the pixel-frame span represented by
the target latent. A single mask is broadcast over that span. Spatial,
temporal, and token reduction all use max coverage so a moving or thin
selected region cannot disappear between causal frames or token rows.
"""
if video.ndim != 5:
raise ValueError(
"H3 target video latent must be [B,C,T,H,W], got %s." %
(list(video.shape),)
)
if int(video.shape[0]) != 1:
raise ValueError("H3 masked targets currently support batch size 1.")
latent_steps = int(video.shape[2])
expected_frames = h3_pixel_frames_for_video_latents(latent_steps)
source = normalize_comfy_mask(mask).to(device=video.device)
source_frames = int(source.shape[0])
static = source_frames == 1
if not static and source_frames != expected_frames:
raise ValueError(
"H3 exact mask conversion needs one static mask or exactly %d "
"tracked masks for this %d-step target; received %d. In a loop, "
"connect Loop Mask Slice. Choose legacy trilinear only to retain "
"the older interpolated behavior." %
(expected_frames, latent_steps, source_frames)
)
reduced = F.adaptive_max_pool2d(
source[:, None],
(int(video.shape[3]), int(video.shape[4])),
)[:, 0]
reduced = _h3_token_snap_max(reduced)
if static:
return reduced[None, None].expand(
1, 1, latent_steps,
int(reduced.shape[-2]), int(reduced.shape[-1]),
).to(torch.float32).contiguous()
groups = h3_video_latent_frame_groups(latent_steps)
reduced = torch.stack(
[reduced[start:end].amax(dim=0) for start, end in groups], dim=0)
return reduced[None, None].to(torch.float32).contiguous()
def temporal_audio_mask(
mask: torch.Tensor,
audio: torch.Tensor,
) -> torch.Tensor:
"""Reduce a video/custom mask to an H3 audio-latent timeline."""
if audio.ndim != 4:
raise ValueError(
"H3 target audio latent must be [B,C,channels,T], got %s." %
(list(audio.shape),)
)
if int(audio.shape[0]) != 1:
raise ValueError("H3 masked targets currently support batch size 1.")
source = normalize_comfy_mask(mask).to(device=audio.device)
timeline = source.amax(dim=(-2, -1))[None, None]
timeline = F.interpolate(
timeline,
size=int(audio.shape[-1]),
mode="linear",
align_corners=False,
)
return timeline[:, :, None].expand(
int(audio.shape[0]), 1, int(audio.shape[2]), int(audio.shape[3])
).to(torch.float32).contiguous()
def _grow_or_shrink_cells(cells: torch.Tensor, amount: int) -> torch.Tensor:
if amount == 0:
return cells
radius = abs(int(amount))
kernel = radius * 2 + 1
if amount > 0:
return F.max_pool2d(
cells, kernel_size=kernel, stride=1, padding=radius)
inverse = 1.0 - cells
inverse = F.pad(
inverse,
(radius, radius, radius, radius),
mode="constant",
value=1.0,
)
return 1.0 - F.max_pool2d(inverse, kernel_size=kernel, stride=1)
def quantize_h3_pixel_mask(
mask: torch.Tensor,
frames: int,
height: int,
width: int,
mode: str = GRID_MODES[0],
cell_adjust: int = 0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Return resized raw mask, snapped pixel mask, and 32px cell map."""
if mode not in GRID_MODES:
raise ValueError("Unknown H3 mask grid mode: %s" % mode)
if height % H3_PIXEL_CELL or width % H3_PIXEL_CELL:
raise ValueError(
"Accurate H3 grid preview requires width and height divisible by "
"%d; received %dx%d." % (H3_PIXEL_CELL, width, height)
)
raw = resize_mask_to_canvas(mask, frames, height, width)
grid_h = height // H3_PIXEL_CELL
grid_w = width // H3_PIXEL_CELL
if mode == "runtime exact (latent max)":
latent = F.adaptive_max_pool2d(
raw,
(height // H3_VAE_SCALE, width // H3_VAE_SCALE),
)[:, 0]
latent = _h3_token_snap_max(latent)
latent_steps = h3_video_latents_for_pixel_frames(frames)
groups = h3_video_latent_frame_groups(latent_steps)
latent = torch.cat([
latent[start:end].amax(dim=0, keepdim=True).expand(
end - start, int(latent.shape[1]), int(latent.shape[2]))
for start, end in groups
], dim=0)
cells = latent[:, None, ::2, ::2]
else:
pixels = raw.reshape(
frames, 1, grid_h, H3_PIXEL_CELL, grid_w, H3_PIXEL_CELL)
if mode == "any pixel coverage":
cells = (pixels.amax(dim=(-3, -1)) >= 0.5).to(raw.dtype)
elif mode == "50% pixel coverage":
cells = (pixels.mean(dim=(-3, -1)) >= 0.5).to(raw.dtype)
else:
cells = (pixels.amin(dim=(-3, -1)) >= 0.5).to(raw.dtype)
cells = _grow_or_shrink_cells(cells, int(cell_adjust))
snapped = cells.repeat_interleave(
H3_PIXEL_CELL, dim=-2).repeat_interleave(H3_PIXEL_CELL, dim=-1)
return raw, snapped[:, 0].contiguous(), cells
|