digital-garbageandClaude Opus 5.5
fa04d6c200
- region lock: layout patch matched the old PackedLayout signature and crashed every H3 run; interior pins now use upstream's native any-frame placement - REINS/Q-steer/AV decouple/probes: masks cover the target video/audio rows only, not pins, refs or vision pads - video detail: works with the PDD FinalLayer signature; now in Engine Settings - Q-steer: steers queries before RoPE, not after (old banks are not reused) - ALG anchor: blurs the frame-0 pin on H3 - temporal style: hidden on H3 - shadow negative: main-window chip names what it switches off without compose Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1308 lines
57 KiB
Python
1308 lines
57 KiB
Python
"""MiniMax-H3 adapter — the one place that knows how H3 differs from LTXAV.
|
|
|
|
FunPack was built against LTX2.3/LTXAV, and supports BOTH families side by side:
|
|
nothing here changes an LTX run, and every entry point answers "not H3" when the
|
|
model in hand is an LTX one. H3 (merged as ComfyUI 57500fc5, PR #15224, shipped
|
|
in ComfyUI v0.30.0) is close enough that most of the stack works untouched, and
|
|
different in four ways that break silently if nobody accounts for them:
|
|
|
|
1. **Audio latent axis.** LTXAV's audio latent is ``[B, C, T, freq]`` — time on
|
|
dim 2, the same axis as the video latent's frames, which is why the Chain
|
|
Sampler could slice every stream at ``[:, :, a:b]``. H3's audio latent is
|
|
``[B, 32, stereo=2, T]`` — dim 2 is the STEREO CHANNEL. Slicing it as time
|
|
silently crops one speaker instead of trimming the tail. Both are 4-D, so the
|
|
tensor alone cannot tell you which; it has to come from the model.
|
|
|
|
2. **Frame grid.** LTXAV is ``8k+1`` pixel frames -> ``k+1`` latent frames. H3 is
|
|
``17k+5`` -> ``5k+2`` (and the VAE reports ``downscale_index_formula = (4,16,16)``,
|
|
which is the *index* map, not the count map — reading a frame count off it is
|
|
off by ~20%). The count always lives in ``vae.downscale_ratio[0]``.
|
|
|
|
3. **No cross-attention.** H3 is a single packed self-attention stream
|
|
``[text | cond/ref | audio | video]``. There is no ``attn2``, no ``audio_attn2``,
|
|
and no audio text-channel half — so the LTXAV conditioning channel split is a
|
|
no-op here (correctly: the whole tensor conditions both streams).
|
|
|
|
4. **Batch size 1.** ``MiniMaxH3Model._forward`` raises on ``x.shape[0] != 1``,
|
|
and comfy batches cond+uncond into one forward whenever ``cfg != 1.0``. Left
|
|
alone that makes every CFG generation crash; :func:`install_batch_split`
|
|
splits the batch back apart around the model call.
|
|
|
|
Nothing here imports the H3 core modules at module scope: an older ComfyUI (or an
|
|
LTX-only install) has none of them, and every entry point has to degrade to "not
|
|
H3" rather than raise on import.
|
|
|
|
WEIGHTS COME IN TWO CHECKPOINTS, and they are not interchangeable:
|
|
``minimax_h3_fl2va_*`` is the t2v / first-last-frame model (keyframe pins) and
|
|
``minimax_h3_ref2va_*`` is the reference model (reference blocks). Both load
|
|
through the same detection and produce a structurally valid graph either way, so
|
|
nothing crashes when the wrong one is loaded — it just conditions badly. The DiT
|
|
state dict carries no variant marker, so :func:`checkpoint_mode_note` reports
|
|
which mode a run is actually using and lets the user match the file.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import torch
|
|
|
|
# ── the model's own constants (comfy/ldm/minimax/model.py, nodes_minimax_h3.py) ──
|
|
IMAGE_MODEL = "minimax_h3"
|
|
FPS = 24 # H3 generates at a fixed 24 fps
|
|
FRAME_GRID = 17 # pixel frames per grid step ...
|
|
FRAME_BASE = 5 # ... on top of this base (frame_count % 17 == 5)
|
|
LATENT_GRID = 5 # latent frames per grid step ...
|
|
LATENT_BASE = 2 # ... on top of this base
|
|
AUDIO_LATENT_FPS = 40 # audio latent frames per second
|
|
SPATIAL_DOWNSCALE = 16
|
|
VIDEO_LATENT_CHANNELS = 24
|
|
AUDIO_LATENT_CHANNELS = 32
|
|
CANVAS_MULTIPLE = 32 # width/height must be a multiple of this
|
|
REF_IMAGE_SHORT_EDGE = 2048 # the reference pipeline's own "max" sizing
|
|
DEFAULT_SHIFT_VIDEO = 12.0
|
|
DEFAULT_SHIFT_AUDIO = 3.0
|
|
|
|
# transformer_options keys MiniMaxH3SigmaShift sets. The DiT inverts the VIDEO shift to
|
|
# its base grid to derive the audio timestep AND the audio velocity slope, so a schedule
|
|
# built on a different shift has to say so here or the audio stream is integrated against
|
|
# the wrong clock. Absent, the DiT falls back to its own 12.0 / 3.0.
|
|
SHIFT_VIDEO_OPTION = "minimax_h3_sigma_shift_video"
|
|
SHIFT_AUDIO_OPTION = "minimax_h3_sigma_shift_audio"
|
|
|
|
# Audio rows are [B, 32, stereo, T]: time is the LAST axis, not dim 2.
|
|
AUDIO_TIME_DIM = -1
|
|
VIDEO_TIME_DIM = 2
|
|
|
|
|
|
# ── the audio clock ──────────────────────────────────────────────────────────
|
|
# H3 denoises video and audio on two different flow schedules (shift 12 / shift 3).
|
|
# Only ONE sigma grid reaches a sampler, so comfy reconciles them inside the DiT: it
|
|
# returns the audio velocity pre-multiplied by d(sigma_a)/d(sigma_v), which makes the
|
|
# flat ODE a sampler integrates on the video grid equal to the audio stream's true ODE
|
|
# on its own grid (comfy/ldm/minimax/model.py, final lines of _forward).
|
|
#
|
|
# That identity holds in the CONTINUOUS limit. The derivative is the tangent at the
|
|
# start of the step, and the two schedules diverge hard, so on a big step the tangent
|
|
# badly overshoots the chord it is standing in for. On a 4-step shift-12 schedule the
|
|
# video grid ~[1, .973, .923, .8, 0] maps to audio ~[1, .9, .75, .5, 0]; the final step's
|
|
# tangent is 1.56 where the true chord is 0.63 — the audio is driven ~2.5x past its
|
|
# target and comes out distorted. At 20+ steps the two agree to within a percent, which
|
|
# is why this only shows up on distilled / turbo few-step schedules.
|
|
#
|
|
# The fix is to integrate the audio on its own grid: swap the tangent for the chord that
|
|
# actually spans the step. One scalar per step, no extra model evaluation.
|
|
|
|
_AUDIO_CLOCK_MAX_FACTOR = 16.0
|
|
|
|
|
|
def time_shift_sigma(sigma, from_shift, to_shift):
|
|
"""Re-express a sigma from one flow schedule on another. Mirrors the DiT's own."""
|
|
sigma = float(sigma)
|
|
from_shift = float(from_shift)
|
|
to_shift = float(to_shift)
|
|
denom = from_shift + sigma * (1.0 - from_shift)
|
|
if abs(denom) < 1e-12:
|
|
return sigma
|
|
base = sigma / denom
|
|
return to_shift * base / (1.0 + (to_shift - 1.0) * base)
|
|
|
|
|
|
def time_shift_slope(sigma, from_shift, to_shift):
|
|
"""d(sigma_to)/d(sigma_from) at the same base-grid point. Mirrors the DiT's own."""
|
|
sigma = float(sigma)
|
|
from_shift = float(from_shift)
|
|
to_shift = float(to_shift)
|
|
denom = from_shift + sigma * (1.0 - from_shift)
|
|
if abs(denom) < 1e-12 or abs(from_shift) < 1e-12:
|
|
return 1.0
|
|
base = sigma / denom
|
|
lower = from_shift * (1.0 + (to_shift - 1.0) * base) ** 2
|
|
if abs(lower) < 1e-12:
|
|
return 1.0
|
|
return (to_shift * (1.0 + (from_shift - 1.0) * base) ** 2) / lower
|
|
|
|
|
|
|
|
|
|
def resolve_sigma_shifts(model_options):
|
|
"""The (video, audio) flow shifts in force, from a model_options dict.
|
|
|
|
MiniMaxH3SigmaShift writes both into transformer_options; without that node the DiT
|
|
falls back to its own defaults, and so do we.
|
|
"""
|
|
to = {}
|
|
if isinstance(model_options, dict):
|
|
candidate = model_options.get("transformer_options")
|
|
if isinstance(candidate, dict):
|
|
to = candidate
|
|
try:
|
|
shift_v = float(to.get(SHIFT_VIDEO_OPTION, DEFAULT_SHIFT_VIDEO))
|
|
except (TypeError, ValueError):
|
|
shift_v = DEFAULT_SHIFT_VIDEO
|
|
try:
|
|
shift_a = float(to.get(SHIFT_AUDIO_OPTION, DEFAULT_SHIFT_AUDIO))
|
|
except (TypeError, ValueError):
|
|
shift_a = DEFAULT_SHIFT_AUDIO
|
|
return shift_v, shift_a
|
|
|
|
|
|
# ── detection ────────────────────────────────────────────────────────────────
|
|
# All three of these have to work when the H3 PR is not installed at all, so they
|
|
# only ever read attributes that any comfy object has.
|
|
|
|
def _unet_config(model):
|
|
inner = getattr(model, "model", None)
|
|
cfg = getattr(inner, "model_config", None)
|
|
return getattr(cfg, "unet_config", {}) or {}
|
|
|
|
|
|
def is_h3_model(model) -> bool:
|
|
"""True for a ModelPatcher (or BaseModel) wrapping the MiniMax H3 DiT."""
|
|
if model is None:
|
|
return False
|
|
try:
|
|
inner = getattr(model, "model", None)
|
|
if type(inner).__name__ == "MiniMaxH3":
|
|
return True
|
|
if type(model).__name__ == "MiniMaxH3":
|
|
return True
|
|
if str(_unet_config(model).get("image_model", "")).lower() == IMAGE_MODEL:
|
|
return True
|
|
diff = getattr(inner, "diffusion_model", None)
|
|
if type(diff).__name__ == "MiniMaxH3Model":
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
|
|
def is_h3_clip(clip) -> bool:
|
|
"""True for the Qwen3-VL-32B conditioning encoder H3 ships with."""
|
|
if clip is None:
|
|
return False
|
|
try:
|
|
tok = getattr(clip, "tokenizer", None)
|
|
if type(tok).__name__ == "MiniMaxH3Tokenizer":
|
|
return True
|
|
# SD1Tokenizer names its sub-tokenizer after the TE; H3's is unique.
|
|
if tok is not None and hasattr(tok, "qwen3vl_32b"):
|
|
return True
|
|
cond_stage = getattr(clip, "cond_stage_model", None)
|
|
if type(cond_stage).__name__ in ("MiniMaxH3TEModel", "MiniMaxH3TEModel_"):
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
|
|
def is_h3_video_vae(vae) -> bool:
|
|
if vae is None:
|
|
return False
|
|
try:
|
|
return type(getattr(vae, "first_stage_model", None)).__name__ == "MiniMaxH3VideoVAE"
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def is_h3_audio_vae(vae) -> bool:
|
|
if vae is None:
|
|
return False
|
|
try:
|
|
return type(getattr(vae, "first_stage_model", None)).__name__ == "MiniMaxH3AudioVAE"
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
# ── frame-grid geometry ──────────────────────────────────────────────────────
|
|
|
|
def align_frame_count(length: int) -> int:
|
|
"""Snap a pixel frame count UP to the model's 17k+5 grid (the node's rule)."""
|
|
n = max(FRAME_BASE, int(length))
|
|
while n % FRAME_GRID != FRAME_BASE:
|
|
n += 1
|
|
return n
|
|
|
|
|
|
def video_latent_frames(frame_count: int) -> int:
|
|
"""Pixel frames -> video latent frames, for an already-aligned count."""
|
|
n = int(frame_count)
|
|
if n <= FRAME_BASE:
|
|
return LATENT_BASE
|
|
return ((n - FRAME_BASE) // FRAME_GRID) * LATENT_GRID + LATENT_BASE
|
|
|
|
|
|
def audio_latent_frames(frame_count: int) -> int:
|
|
"""Pixel frames -> audio latent frames (40 Hz against a 24 fps video clock)."""
|
|
return int(round(int(frame_count) / float(FPS) * AUDIO_LATENT_FPS))
|
|
|
|
|
|
def temporal_shape(length: int):
|
|
"""(aligned pixel frames, video latent frames, audio latent frames)."""
|
|
frames = align_frame_count(length)
|
|
return frames, video_latent_frames(frames), audio_latent_frames(frames)
|
|
|
|
|
|
def latent_frames_from_vae(vae, pixel_frames: int):
|
|
"""Video latent frame count for `pixel_frames`, asked of the VAE itself.
|
|
|
|
``vae.downscale_ratio[0]`` is the count map every video VAE in comfy carries
|
|
(LTX: ``ceil(n/8)``; H3: ``(n-5)//17*5+2``). Reading it beats hardcoding a
|
|
per-family formula and beats ``downscale_index_formula``, which is the index
|
|
map and gives the wrong count on H3. Returns None when the VAE doesn't
|
|
expose one, so callers can keep their old fallback.
|
|
"""
|
|
try:
|
|
ratio = getattr(vae, "downscale_ratio", None)
|
|
if isinstance(ratio, (list, tuple)) and ratio and callable(ratio[0]):
|
|
return max(1, int(ratio[0](max(1, int(pixel_frames)))))
|
|
if isinstance(ratio, (int, float)) and ratio:
|
|
return ((max(1, int(pixel_frames)) - 1) // max(1, int(ratio))) + 1
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def adapt_canvas(width: int, height: int, base_short_edge=768, max_pixels=768 * 1344):
|
|
"""H3's reference canvas rule: 768 short edge, area-capped, rounded to 32."""
|
|
width = max(1, int(width))
|
|
height = max(1, int(height))
|
|
ratio = width / height
|
|
if ratio >= 1.0:
|
|
nom_w, nom_h = base_short_edge * ratio, float(base_short_edge)
|
|
else:
|
|
nom_w, nom_h = float(base_short_edge), base_short_edge / ratio
|
|
if nom_w * nom_h > max_pixels:
|
|
s = math.sqrt(max_pixels / (nom_w * nom_h))
|
|
nom_w, nom_h = nom_w * s, nom_h * s
|
|
snap = lambda v: max(CANVAS_MULTIPLE, int(round(v / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE)
|
|
return snap(nom_w), snap(nom_h)
|
|
|
|
|
|
# ── per-stream time axes ─────────────────────────────────────────────────────
|
|
|
|
def stream_time_dims(stream_count: int, h3: bool):
|
|
"""Time axis per latent stream, in comfy's pack order (video first).
|
|
|
|
LTXAV: every stream carries time on dim 2. H3: video on dim 2, audio on the
|
|
last axis (dim 2 there is the stereo channel). Returned as a plain list so
|
|
callers can index it by stream position without re-deciding the family.
|
|
"""
|
|
dims = [VIDEO_TIME_DIM] * max(1, int(stream_count))
|
|
if h3:
|
|
for i in range(1, len(dims)):
|
|
dims[i] = AUDIO_TIME_DIM
|
|
return dims
|
|
|
|
|
|
def time_len(tensor, dim: int) -> int:
|
|
return int(tensor.shape[dim])
|
|
|
|
|
|
def time_slice(tensor, start, end, dim: int):
|
|
idx = [slice(None)] * tensor.dim()
|
|
idx[dim] = slice(start, end)
|
|
return tensor[tuple(idx)]
|
|
|
|
|
|
def set_time_slice(tensor, start, end, value, dim: int):
|
|
idx = [slice(None)] * tensor.dim()
|
|
idx[dim] = slice(start, end)
|
|
tensor[tuple(idx)] = value
|
|
return tensor
|
|
|
|
|
|
def time_cat(tensors, dim: int):
|
|
return torch.cat(tensors, dim=dim if dim >= 0 else tensors[0].dim() + dim)
|
|
|
|
|
|
# ── attention bridge ─────────────────────────────────────────────────────────
|
|
# LTX calls ``comfy.ldm.modules.attention.optimized_attention(...)`` through the
|
|
# module, so FunPack's temperature / capture / K-V-lock patches take effect by
|
|
# rebinding that one symbol. H3 does ``from ... import optimized_attention`` at
|
|
# module load, which snapshots the function — rebinding the source module does
|
|
# nothing. Any patch that wants to reach H3 has to rebind H3's own copy too.
|
|
|
|
def attention_patch_targets():
|
|
"""Modules whose ``optimized_attention`` symbol must be rebound together.
|
|
|
|
Always includes comfy's attention module; includes the H3 DiT module only
|
|
when the PR is installed. Callers save/restore whatever they find.
|
|
"""
|
|
targets = []
|
|
try:
|
|
import comfy.ldm.modules.attention as attn_mod
|
|
targets.append(attn_mod)
|
|
except Exception:
|
|
return targets
|
|
try:
|
|
import comfy.ldm.minimax.model as h3_mod
|
|
if hasattr(h3_mod, "optimized_attention"):
|
|
targets.append(h3_mod)
|
|
except Exception:
|
|
pass
|
|
return targets
|
|
|
|
|
|
# H3 calls ``optimized_attention`` with ``skip_reshape=True`` and q/k/v already at
|
|
# ``[1, heads, S, head_dim]``. A patch that reads ``v.shape[1]`` as a token count (the
|
|
# LTX layout) is reading the HEAD count instead, so anything indexed by token position
|
|
# has to check that kwarg before assuming a layout. Temperature is layout agnostic.
|
|
|
|
|
|
# ── batch-size-1 constraint ──────────────────────────────────────────────────
|
|
|
|
def _split_cond_batch(value, index, batch):
|
|
"""Slice one entry of the ``c`` dict down to sample `index` of `batch`."""
|
|
if isinstance(value, torch.Tensor) and value.dim() >= 1 and value.shape[0] == batch:
|
|
return value[index:index + 1]
|
|
if getattr(value, "is_nested", False):
|
|
try:
|
|
import comfy.nested_tensor
|
|
parts = [t[index:index + 1] if t.shape[0] == batch else t for t in value.unbind()]
|
|
return comfy.nested_tensor.NestedTensor(parts)
|
|
except Exception:
|
|
return value
|
|
return value
|
|
|
|
|
|
def install_batch_split(model):
|
|
"""Run the H3 DiT one sample at a time when comfy hands it a batched forward.
|
|
|
|
``MiniMaxH3Model._forward`` raises on ``x.shape[0] != 1``. In practice comfy does
|
|
NOT batch the positive and negative conds together on H3 — each carries its own
|
|
``minimax_payload`` (different token tags, different packed layout), and an unequal
|
|
``CONDConstant`` blocks the concat — so the cfg>1 path survives on its own. What
|
|
does reach the DiT batched is a latent whose batch size is >1 (several videos in one
|
|
request). Splitting here keeps that working at exactly the cost the batch already
|
|
implies, and stays correct if comfy's batching rules change.
|
|
|
|
Returns the previous ``model_function_wrapper`` so the caller can restore it.
|
|
"""
|
|
old_wrapper = model.model_options.get("model_function_wrapper")
|
|
|
|
def _call(apply_fn, args):
|
|
if old_wrapper is not None:
|
|
return old_wrapper(apply_fn, args)
|
|
return apply_fn(args["input"], args["timestep"], **args.get("c", {}))
|
|
|
|
def _batch_split_wrapper(apply_fn, args, _old=old_wrapper):
|
|
x = args.get("input")
|
|
try:
|
|
batch = int(x.shape[0]) if isinstance(x, torch.Tensor) else 1
|
|
except Exception:
|
|
batch = 1
|
|
if batch <= 1:
|
|
return _call(apply_fn, args)
|
|
outs = []
|
|
ts = args.get("timestep")
|
|
c = args.get("c") or {}
|
|
for i in range(batch):
|
|
sub = dict(args)
|
|
sub["input"] = x[i:i + 1]
|
|
if isinstance(ts, torch.Tensor) and ts.shape[0] == batch:
|
|
sub["timestep"] = ts[i:i + 1]
|
|
sub["c"] = {k: _split_cond_batch(v, i, batch) for k, v in c.items()}
|
|
outs.append(_call(apply_fn, sub))
|
|
return torch.cat(outs, dim=0)
|
|
|
|
model.model_options["model_function_wrapper"] = _batch_split_wrapper
|
|
return old_wrapper, _batch_split_wrapper
|
|
|
|
|
|
# ── conditioning payload ─────────────────────────────────────────────────────
|
|
# H3 carries its visual conditioning on the CONDITIONING dict, not in the latent:
|
|
# minimax_keyframes — fl2va first/last frame pins (never denoised)
|
|
# minimax_frame_count— pixel length, so the layout can place the last-frame pin
|
|
# minimax_refs — ref2va identity/style/audio reference blocks
|
|
# model_base.MiniMaxH3.extra_conds turns these into the DiT payload.
|
|
|
|
def _set_values(conditioning, values):
|
|
out = []
|
|
for entry in conditioning or []:
|
|
cond, meta = entry[0], (entry[1] if len(entry) > 1 else {})
|
|
new_meta = dict(meta) if isinstance(meta, dict) else {}
|
|
for key, value in values.items():
|
|
if value is None:
|
|
new_meta.pop(key, None)
|
|
else:
|
|
new_meta[key] = value
|
|
out.append([cond, new_meta])
|
|
return out
|
|
|
|
|
|
def set_keyframes(conditioning, keyframes, frame_count):
|
|
"""Attach fl2va keyframe pins. `keyframes` = [{resolved_frame_index, latent}]."""
|
|
if not keyframes:
|
|
return conditioning
|
|
return _set_values(conditioning, {
|
|
"minimax_keyframes": list(keyframes),
|
|
"minimax_frame_count": int(frame_count),
|
|
})
|
|
|
|
|
|
def set_refs(conditioning, ref_blocks):
|
|
"""Attach ref2va reference blocks (images / videos / audio)."""
|
|
if not ref_blocks:
|
|
return conditioning
|
|
return _set_values(conditioning, {"minimax_refs": list(ref_blocks)})
|
|
|
|
|
|
def checkpoint_mode_note(has_keyframes: bool, has_refs: bool):
|
|
"""One line naming the checkpoint this run's conditioning actually needs, or None.
|
|
|
|
H3 ships as two separate DiTs — ``minimax_h3_fl2va_*`` (t2v / first-last frame) and
|
|
``minimax_h3_ref2va_*`` (references) — and community checkpoints merge both. Nothing in
|
|
the state dict identifies the variant, so THIS CANNOT AND DOES NOT CHECK THE LOADED FILE.
|
|
It names the conditioning the run is using and stops there.
|
|
|
|
It used to add "fl2va weights will read it, but were not trained on it", which reads as a
|
|
finding and is not one: it fires on every run with references, including on the ref2va and
|
|
hybrid checkpoints where nothing is wrong. Asserting a mismatch nothing detected is worse
|
|
than silence — the reader has to learn to ignore the line, and then ignores it when it
|
|
matters.
|
|
"""
|
|
if has_keyframes and has_refs:
|
|
# One sentence, because the log prints the first one: the finding has to be IN it.
|
|
return ("BOTH keyframe pins and references, which live in different DiTs unless "
|
|
"yours is a hybrid — one of the two is then untrained conditioning. Drop the "
|
|
"anchor image or the references, or run them as separate scenes.")
|
|
if has_keyframes:
|
|
return ("keyframe pins (fl2va conditioning). Needs fl2va or hybrid weights — which "
|
|
"variant is loaded cannot be read from the file.")
|
|
if has_refs:
|
|
return ("references (ref2va conditioning). Needs ref2va or hybrid weights — which "
|
|
"variant is loaded cannot be read from the file.")
|
|
return None
|
|
|
|
|
|
# ── keyframe pins: any frame, and region locks ──────────────────────────────
|
|
#
|
|
# Upstream places a pin at ANY pixel frame itself (``cond_t = cursor + FRAME_RESCALE *
|
|
# resolved_frame_index``, since #15439). FunPack used to patch that in; the patch passed a
|
|
# ``frame_count`` argument upstream then removed, which broke EVERY later H3 layout in the
|
|
# process. Interior pins now need nothing from FunPack.
|
|
#
|
|
# fl2va was trained on first/last pins only, so a mid-clip pin is representable rather than
|
|
# learned — the sampler still says so when it places one.
|
|
_REGION_LOCKS = {"state": None} # None = not tried, True/False = patch verdict
|
|
_LAYOUT_PATCH = {"state": None} # is PackedLayout.__init__ wrapped at all
|
|
|
|
#: key a keyframe pin carries its keep-mask under (bool tensor over the frame's patch rows)
|
|
REGION_META = "funpack_region"
|
|
|
|
|
|
def _video_pins(keyframes):
|
|
"""The pins that own condition VIDEO rows, in row order.
|
|
|
|
Upstream emits one ``cond`` segment and one ``cond_video_latents`` entry per pin WITH a
|
|
video latent; an audio-only pin has neither. Refs follow the pins in both.
|
|
"""
|
|
return [kf for kf in keyframes or () if kf.get("latent") is not None]
|
|
|
|
|
|
def _region_keep(region, n_rows):
|
|
"""The keep-mask to actually apply to a pin's ``n_rows`` patch rows, or None.
|
|
|
|
None means "keep every row", and it is returned for every case that is not an
|
|
unambiguous partial selection: no region, a mask whose length disagrees with the rows it
|
|
would filter, an empty selection, or a full one. Both sides of the patch — the layout and
|
|
the condition rows — go through this one function with the same mask and the same row
|
|
count, so they cannot disagree about which rows survive.
|
|
"""
|
|
if region is None:
|
|
return None
|
|
try:
|
|
mask = region.reshape(-1).to(torch.bool)
|
|
except Exception:
|
|
return None
|
|
if mask.numel() != int(n_rows) or not bool(mask.any()) or bool(mask.all()):
|
|
return None
|
|
return mask
|
|
|
|
|
|
def region_rows_from_alpha(image, grid_h, grid_w, threshold=0.5):
|
|
"""An RGBA guide image's alpha channel as a keep-mask over one frame's patch rows.
|
|
|
|
Transparent means "I have nothing to say about this part of the frame". Alpha is used
|
|
rather than a separate MASK input because it needs no new wire: the Editor's fixed graph
|
|
has no mask source, so a MASK socket would be a sampler-only feature, and the standing
|
|
rule is that a new sampler capability reaches the Editor in the same change.
|
|
|
|
Returns None unless the alpha is a genuine PARTIAL selection — a fully opaque image (the
|
|
overwhelmingly common case) behaves exactly as it did before this existed.
|
|
"""
|
|
try:
|
|
if getattr(image, "ndim", 0) != 4 or image.shape[-1] < 4:
|
|
return None
|
|
alpha = image[:1, ..., 3].float() # [1, H, W]
|
|
pooled = torch.nn.functional.adaptive_avg_pool2d(
|
|
alpha.unsqueeze(0), (int(grid_h), int(grid_w))) # [1, 1, gh, gw]
|
|
except Exception:
|
|
return None
|
|
mask = (pooled.reshape(-1) >= float(threshold))
|
|
if not bool(mask.any()) or bool(mask.all()):
|
|
return None
|
|
return mask
|
|
|
|
|
|
def _apply_region_locks(layout, spans, regions):
|
|
"""Drop the condition rows a pin's region excludes, and re-derive the layout around them.
|
|
|
|
A region lock is a row DELETION, not a blanked or noised row. A blanked row still says
|
|
something ("this part is flat grey") and a noised one still occupies a position the
|
|
picture attends to; a row that does not exist says nothing at all, which is what "I only
|
|
care about this part of the frame" means. Attention is dense, so the removed rows also
|
|
stop costing anything.
|
|
|
|
Everything here is re-derived from upstream's own outputs rather than rebuilt, so a
|
|
change to how upstream lays the sequence out is inherited. ``img_pos``/``audio_pos`` are
|
|
renumbered even though the current forward never reads them — they are part of the
|
|
layout's contract, and a stale one would be a trap for whatever reads them next.
|
|
"""
|
|
keep = torch.ones(layout.seq_len, dtype=torch.bool)
|
|
for (start, stop), region in zip(spans, regions):
|
|
mask = _region_keep(region, stop - start)
|
|
if mask is not None:
|
|
keep[start:stop] = mask
|
|
if bool(keep.all()):
|
|
return False
|
|
new_index = keep.cumsum(0) - 1
|
|
img_keep = keep[layout.img_pos]
|
|
audio_keep = keep[layout.audio_pos]
|
|
layout.position_ids = layout.position_ids[keep]
|
|
layout.img_update = layout.img_update[img_keep]
|
|
layout.audio_update = layout.audio_update[audio_keep]
|
|
layout.img_pos = new_index[layout.img_pos[img_keep]]
|
|
layout.audio_pos = new_index[layout.audio_pos[audio_keep]]
|
|
segments, offset = [], 0
|
|
for a, b, kind in layout.segments:
|
|
n = int(keep[a:b].sum())
|
|
if n:
|
|
segments.append((offset, offset + n, kind))
|
|
offset += n
|
|
layout.segments = segments
|
|
layout.seq_len = offset
|
|
return True
|
|
|
|
|
|
def _install_layout_patch():
|
|
"""Wrap ``PackedLayout.__init__`` once, for region locks.
|
|
|
|
Deliberately NOT a re-implementation of it: upstream builds the layout exactly as it
|
|
always has, and this removes the rows a region excludes afterwards. Arguments pass
|
|
through untouched (``*args, **kwargs``), so an upstream signature change is inherited
|
|
instead of breaking every layout — which is what the old interior-pin wrapper did.
|
|
|
|
A run with no region never enters the rewrite, so its layout is upstream's, bit for bit.
|
|
"""
|
|
if _LAYOUT_PATCH["state"] is not None:
|
|
return _LAYOUT_PATCH["state"]
|
|
try:
|
|
from comfy.ldm.minimax.model import PackedLayout
|
|
except Exception as error: # no H3 in this ComfyUI
|
|
_LAYOUT_PATCH["state"] = False
|
|
print(f"[FunPack H3] packed-layout extensions unavailable ({error}).")
|
|
return False
|
|
if getattr(PackedLayout.__init__, "_funpack_layout", False):
|
|
_LAYOUT_PATCH["state"] = True
|
|
return True
|
|
|
|
original = PackedLayout.__init__
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
original(self, *args, **kwargs)
|
|
if not _REGION_LOCKS["state"]:
|
|
return
|
|
pins = _video_pins(kwargs.get("keyframes"))
|
|
regions = [kf.get(REGION_META) for kf in pins]
|
|
if not any(r is not None for r in regions):
|
|
return
|
|
# Condition segments are emitted in pin order, immediately after text.
|
|
spans = [(a, b) for a, b, kind in self.segments if kind == "cond"]
|
|
if len(spans) != len(pins):
|
|
return # our accounting disagrees with upstream's: change nothing
|
|
_apply_region_locks(self, spans, regions)
|
|
|
|
__init__._funpack_layout = True
|
|
PackedLayout.__init__ = __init__
|
|
_LAYOUT_PATCH["state"] = True
|
|
return True
|
|
|
|
|
|
def install_region_locks():
|
|
"""Let a keyframe pin cover PART of a frame instead of all of it.
|
|
|
|
Needs two patches that must agree exactly on which rows survive: the layout drops the
|
|
excluded rows from the packed sequence, and ``_cond_video_rows`` drops the same rows from
|
|
the condition latents. Both go through :func:`_region_keep` with the same mask, so they
|
|
cannot diverge — and if either patch cannot be installed, neither is used.
|
|
|
|
EXPERIMENTAL: fl2va saw whole condition frames, so a partial one is representable
|
|
(nothing in the layout or the attention requires a cond segment to be a complete frame)
|
|
without being something the weights have seen.
|
|
"""
|
|
if _REGION_LOCKS["state"] is not None:
|
|
return _REGION_LOCKS["state"]
|
|
if not _install_layout_patch():
|
|
_REGION_LOCKS["state"] = False
|
|
return False
|
|
try:
|
|
from comfy.ldm.minimax.model import MiniMaxH3Model
|
|
except Exception as error:
|
|
_REGION_LOCKS["state"] = False
|
|
print(f"[FunPack H3] region locks unavailable ({error}).")
|
|
return False
|
|
if getattr(MiniMaxH3Model._cond_video_rows, "_funpack_regions", False):
|
|
_REGION_LOCKS["state"] = True
|
|
return True
|
|
|
|
original = MiniMaxH3Model._cond_video_rows
|
|
|
|
def _cond_video_rows(self, payload, device):
|
|
rows = original(self, payload, device)
|
|
try:
|
|
if rows is None:
|
|
return rows
|
|
pins = _video_pins(payload.get("keyframes"))
|
|
if not any(pin.get(REGION_META) is not None for pin in pins):
|
|
return rows
|
|
# Pins' latents lead the list; reference latents (if any) follow and are kept whole.
|
|
latents = payload.get("cond_video_latents") or []
|
|
if len(latents) < len(pins):
|
|
return rows
|
|
kept, offset = [], 0
|
|
for index, latent in enumerate(latents):
|
|
b, _c, t, h, w = latent.shape
|
|
n = b * t * (h // 2) * (w // 2)
|
|
span = rows[offset:offset + n]
|
|
offset += n
|
|
mask = _region_keep(pins[index].get(REGION_META), n) if index < len(pins) else None
|
|
kept.append(span if mask is None else span[mask.to(span.device)])
|
|
if offset != rows.shape[0]:
|
|
# Our row accounting disagrees with upstream's: change nothing rather than
|
|
# hand the layout a set of rows it did not plan for.
|
|
return rows
|
|
return torch.cat(kept, dim=0)
|
|
except Exception as error: # noqa: BLE001
|
|
print(f"[FunPack H3] region lock skipped on the condition rows ({error}); "
|
|
f"the pin covers the whole frame.")
|
|
return rows
|
|
|
|
_cond_video_rows._funpack_regions = True
|
|
MiniMaxH3Model._cond_video_rows = _cond_video_rows
|
|
_REGION_LOCKS["state"] = True
|
|
return True
|
|
|
|
|
|
def keyframe_is_endpoint(frame_index, frame_count):
|
|
"""True for the first or last pixel frame — the two positions fl2va was trained on."""
|
|
if frame_index is None:
|
|
return False
|
|
idx = int(frame_index)
|
|
if idx == 0:
|
|
return True
|
|
return frame_count is not None and idx == int(frame_count) - 1
|
|
|
|
|
|
def keyframe_indices_supported(frame_index, frame_count):
|
|
"""Can a pin be placed at this pixel frame? Any frame inside the clip: upstream places it."""
|
|
return 0 <= int(frame_index) <= int(frame_count) - 1
|
|
|
|
|
|
def encode_keyframe(vae, image, width, height, frame_index, crop="disabled"):
|
|
"""VAE-encode one pixel frame into an H3 keyframe pin entry.
|
|
|
|
`crop` mirrors the reference node: the frame-0 anchor is stretched to the
|
|
canvas (it defines the geometry), a last-frame follower is cover-cropped.
|
|
"""
|
|
import comfy.utils
|
|
# one frame only: a pin is a single condition frame, and an IMAGE batch here would
|
|
# otherwise encode into a latent with a batch dimension the packed layout cannot place.
|
|
# Alpha rides through the SAME resize as the colour so a cover-cropped pin's region ends
|
|
# up over the part of the frame it was painted on, not over the pre-crop coordinates.
|
|
channels = 4 if image.shape[-1] >= 4 else 3
|
|
samples = image[:1, ..., :channels].movedim(-1, 1)
|
|
samples = comfy.utils.common_upscale(samples, int(width), int(height), "lanczos", crop)
|
|
resized = samples.movedim(1, -1)
|
|
pixels = resized[..., :3]
|
|
region = None
|
|
if channels == 4 and install_region_locks():
|
|
grid_h = int(height) // SPATIAL_DOWNSCALE // 2
|
|
grid_w = int(width) // SPATIAL_DOWNSCALE // 2
|
|
region = region_rows_from_alpha(resized, grid_h, grid_w)
|
|
if region is not None:
|
|
# The excluded ROWS are dropped, but the VAE is convolutional: whatever colour sits
|
|
# under the transparent area still bleeds into the latent cells that are kept. An
|
|
# unpainted PNG leaves black there, which would print a hard edge into the kept
|
|
# boundary. Filling with the kept area's own mean colour makes that edge as quiet as
|
|
# the encoder can make it.
|
|
alpha = resized[..., 3:4]
|
|
keep = (alpha >= 0.5).float()
|
|
weight = keep.sum().clamp_min(1.0)
|
|
fill = (pixels * keep).sum(dim=(0, 1, 2), keepdim=True) / weight
|
|
pixels = pixels * keep + fill * (1.0 - keep)
|
|
pin = {"resolved_frame_index": int(frame_index), "latent": vae.encode(pixels)}
|
|
if region is not None:
|
|
pin[REGION_META] = region
|
|
return pin
|
|
|
|
|
|
def image_ref_block(vae, image, width, height, size_mode="match"):
|
|
"""Encode a pixel image into a ref2va IMAGE block (identity / style reference).
|
|
|
|
`size_mode` mirrors the reference node's own control: "match" scales the reference
|
|
(down only, keeping aspect) to the generation's pixel area; "max" uses the reference
|
|
pipeline's 2048px short edge for the best identity fidelity. Reference rows ride
|
|
through every sampling step, so "max" costs real time on every step of every scene.
|
|
"""
|
|
import comfy.utils
|
|
h, w = int(image.shape[1]), int(image.shape[2])
|
|
if str(size_mode).lower() == "max":
|
|
scale = min(1.0, REF_IMAGE_SHORT_EDGE / max(1, min(w, h)))
|
|
else:
|
|
scale = min(1.0, math.sqrt((int(width) * int(height)) / max(1, w * h)))
|
|
snap = lambda v: max(CANVAS_MULTIPLE, int(round(v / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE)
|
|
tw, th = snap(w * scale), snap(h * scale)
|
|
samples = image[:1, ..., :3].movedim(-1, 1)
|
|
samples = comfy.utils.common_upscale(samples, tw, th, "lanczos", "disabled")
|
|
resized = samples.movedim(1, -1)
|
|
return ({"type": "image", "data": resized},
|
|
{"kind": "image", "latent_h": th // SPATIAL_DOWNSCALE,
|
|
"latent_w": tw // SPATIAL_DOWNSCALE, "latent": vae.encode(resized)})
|
|
|
|
|
|
def audio_ref_block(audio_vae, audio):
|
|
"""Encode an AUDIO dict into a ref2va AUDIO block."""
|
|
import torchaudio
|
|
waveform = audio["waveform"]
|
|
sr = int(audio.get("sample_rate", 32000))
|
|
vae_sr = int(getattr(audio_vae, "audio_sample_rate", 32000))
|
|
if sr != vae_sr:
|
|
waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
|
|
z = audio_vae.encode(waveform[:1].movedim(1, -1))
|
|
return ({"type": "audio"},
|
|
{"kind": "audio", "ref_audio_t": int(z.shape[-1]), "audio_latent": z})
|
|
|
|
|
|
# ── ref2va: one ordered spec, two consumers ──────────────────────────────────
|
|
# A reference has to reach the model twice and in the SAME order both times:
|
|
# * as PRESENTATION — "<Picture 1>: <vision block>" spliced into the Qwen prompt, so
|
|
# the text encoder has seen it and the prompt can refer to "<Picture 1>";
|
|
# * as a LATENT BLOCK — packed into the DiT sequence, never denoised.
|
|
# Those two encodes happen in different nodes (Studio owns the CLIP, the Chain Sampler
|
|
# owns the VAE), so the ordered spec travels between them on the conditioning and both
|
|
# sides walk it with the helpers below. Ordinals are 1-based per type, exactly as the
|
|
# reference node numbers them, so a drift in order silently re-points "<Picture 2>".
|
|
|
|
REF_KINDS = ("image", "audio", "video")
|
|
|
|
|
|
def normalize_ref_spec(raw):
|
|
"""Coerce a user/JSON reference list into an ordered list of {kind, filename, ...}.
|
|
|
|
Unknown kinds and entries without a filename are dropped rather than guessed at.
|
|
Returns [] for anything unusable, so callers can treat "no references" and "bad
|
|
references" the same way.
|
|
"""
|
|
if isinstance(raw, str):
|
|
import json
|
|
try:
|
|
raw = json.loads(raw or "[]")
|
|
except Exception:
|
|
return []
|
|
if isinstance(raw, dict):
|
|
raw = raw.get("references") or raw.get("refs") or []
|
|
if not isinstance(raw, (list, tuple)):
|
|
return []
|
|
out = []
|
|
for item in raw:
|
|
if isinstance(item, str):
|
|
item = {"kind": "image", "filename": item}
|
|
if not isinstance(item, dict):
|
|
continue
|
|
kind = str(item.get("kind") or item.get("type") or "image").strip().lower()
|
|
filename = str(item.get("filename") or item.get("file") or "").strip()
|
|
if kind not in REF_KINDS or not filename:
|
|
continue
|
|
entry = {"kind": kind, "filename": filename}
|
|
if kind == "video" and item.get("audio"):
|
|
entry["audio"] = str(item["audio"]).strip()
|
|
if kind == "image":
|
|
size = str(item.get("size") or "match").strip().lower()
|
|
entry["size"] = size if size in ("match", "max") else "match"
|
|
out.append(entry)
|
|
return out
|
|
|
|
|
|
def load_input_image(filename):
|
|
"""Load an image from ComfyUI's input directory as [1, H, W, 3] in 0..1."""
|
|
import os
|
|
try:
|
|
import folder_paths
|
|
import numpy as np
|
|
from PIL import Image
|
|
except ImportError:
|
|
return None
|
|
path = os.path.join(folder_paths.get_input_directory(), filename)
|
|
if not os.path.isfile(path):
|
|
return None
|
|
try:
|
|
img = Image.open(path).convert("RGB")
|
|
arr = np.array(img).astype("float32") / 255.0
|
|
return torch.from_numpy(arr)[None,]
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def load_input_audio(filename):
|
|
"""Load an audio file from ComfyUI's input directory as comfy's AUDIO dict."""
|
|
import os
|
|
try:
|
|
import folder_paths
|
|
import torchaudio
|
|
except ImportError:
|
|
return None
|
|
path = os.path.join(folder_paths.get_input_directory(), filename)
|
|
if not os.path.isfile(path):
|
|
return None
|
|
try:
|
|
waveform, sample_rate = torchaudio.load(path)
|
|
return {"waveform": waveform.unsqueeze(0), "sample_rate": int(sample_rate)}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def load_input_video(filename, fps=FPS, max_seconds=15.0):
|
|
"""Decode a video from ComfyUI's input directory to frames [T, H, W, 3] in 0..1.
|
|
|
|
ffmpeg, which FunPack already requires for preview/export, rather than a node
|
|
dependency — the reference pipeline wants frames at the model's own 24 fps and the
|
|
tokenizer wants every second frame of that, so decoding once here keeps both views
|
|
of the same clip consistent. Capped at `max_seconds` because reference tokens ride
|
|
through every sampling step: a long clip is a silent, permanent slowdown.
|
|
"""
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
try:
|
|
import folder_paths
|
|
import numpy as np
|
|
except ImportError:
|
|
return None
|
|
path = os.path.join(folder_paths.get_input_directory(), filename)
|
|
if not os.path.isfile(path):
|
|
return None
|
|
ff = shutil.which("ffmpeg")
|
|
if ff is None:
|
|
return None
|
|
# rgb24 rawvideo on stdout: no temp files, no container parsing on our side. The
|
|
# scale filter pins the frame size so the byte stream is a flat T x H x W x 3 block.
|
|
probe = shutil.which("ffprobe")
|
|
width = height = None
|
|
if probe is not None:
|
|
try:
|
|
out = subprocess.run(
|
|
[probe, "-v", "error", "-select_streams", "v:0", "-show_entries",
|
|
"stream=width,height", "-of", "csv=p=0:s=x", path],
|
|
capture_output=True, text=True, timeout=30)
|
|
width, height = (int(v) for v in out.stdout.strip().split("x")[:2])
|
|
except Exception:
|
|
width = height = None
|
|
if not width or not height:
|
|
return None
|
|
# even dimensions keep the decode simple and match the model's 32-multiple canvas rule
|
|
width -= width % 2
|
|
height -= height % 2
|
|
try:
|
|
proc = subprocess.run(
|
|
[ff, "-v", "error", "-i", path, "-t", str(float(max_seconds)),
|
|
"-vf", f"fps={int(fps)},scale={width}:{height}",
|
|
"-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
|
|
capture_output=True, timeout=300)
|
|
except Exception:
|
|
return None
|
|
if proc.returncode != 0 or not proc.stdout:
|
|
return None
|
|
frame_bytes = width * height * 3
|
|
count = len(proc.stdout) // frame_bytes
|
|
if count < 5: # the model's floor: reference videos need at least 5 frames
|
|
return None
|
|
arr = np.frombuffer(proc.stdout[:count * frame_bytes], dtype=np.uint8)
|
|
arr = arr.reshape(count, height, width, 3).astype("float32") / 255.0
|
|
return torch.from_numpy(arr.copy())
|
|
|
|
|
|
def trim_video_to_grid(frames):
|
|
"""Trim reference video frames DOWN to the model's 17k+5 grid (never up).
|
|
|
|
Trims rather than pads because padding a reference invents motion that was never in
|
|
the source. Returns None when there is not a single whole grid step left, which is
|
|
the same floor the model enforces. Deliberately independent of the generation length:
|
|
Studio and the Chain Sampler both trim this way, so the clip the text encoder saw and
|
|
the clip that was encoded to latents are always the same frames.
|
|
"""
|
|
n = int(frames.shape[0])
|
|
while n > 0 and n % FRAME_GRID != FRAME_BASE:
|
|
n -= 1
|
|
if n < FRAME_BASE:
|
|
return None
|
|
return frames[:n]
|
|
|
|
|
|
def video_presentation_frames(frames):
|
|
"""The 2 fps view Qwen sees, with its timestamps — the tokenizer's own sampling rate.
|
|
|
|
Returns (frames_at_2fps, timestamps_in_seconds). The tokenizer repeat-pads an odd
|
|
count into its 2-frame temporal patch, so an odd number here is fine.
|
|
"""
|
|
step = max(1, FPS // 2)
|
|
idx = list(range(0, int(frames.shape[0]), step))
|
|
return frames[idx], [i / 2.0 for i in range(len(idx))]
|
|
|
|
|
|
def video_ref_block(vae, frames, audio_vae=None, audio=None):
|
|
"""Encode reference video frames (plus an optional soundtrack) into a ref2va block.
|
|
|
|
The block's audio rows pack immediately before its video rows and share a cursor
|
|
origin, which is what ties a soundtrack to its own clip rather than to the target.
|
|
"""
|
|
h, w = int(frames.shape[1]), int(frames.shape[2])
|
|
cw, ch = adapt_canvas(w, h)
|
|
if w * h < cw * ch:
|
|
# never upscale a reference: keep its own size, snapped to the canvas multiple
|
|
snap = lambda v: max(CANVAS_MULTIPLE, int(round(v / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE)
|
|
cw, ch = snap(w), snap(h)
|
|
import comfy.utils
|
|
samples = frames[..., :3].movedim(-1, 1)
|
|
samples = comfy.utils.common_upscale(samples, cw, ch, "lanczos", "disabled")
|
|
resized = samples.movedim(1, -1)
|
|
|
|
z = vae.encode(resized)
|
|
audio_latent, ref_audio_t = None, 0
|
|
if audio is not None and audio_vae is not None:
|
|
_item, audio_block = audio_ref_block(audio_vae, audio)
|
|
audio_latent = audio_block["audio_latent"]
|
|
ref_audio_t = audio_block["ref_audio_t"]
|
|
|
|
qwen_frames, timestamps = video_presentation_frames(resized)
|
|
item = {"type": "video", "data": qwen_frames, "timestamps": timestamps}
|
|
block = {"kind": "video_audio" if ref_audio_t else "video",
|
|
"latent_t": int(z.shape[2]), "latent_h": ch // SPATIAL_DOWNSCALE,
|
|
"latent_w": cw // SPATIAL_DOWNSCALE, "ref_audio_t": ref_audio_t,
|
|
"latent": z, "audio_latent": audio_latent}
|
|
return item, block
|
|
|
|
|
|
def ref_items_from_spec(spec):
|
|
"""Spec -> the tokenizer's `minimax_ref_items` (presentation only, no VAE needed).
|
|
|
|
Only images carry pixels into Qwen; audio contributes a bare "<Audio j>: " label
|
|
(the reference node does the same — audio never enters the text encoder). Entries
|
|
whose file will not load are dropped from BOTH lists by the caller pairing this
|
|
with :func:`ref_blocks_from_spec` over the same resolved spec.
|
|
"""
|
|
items = []
|
|
for entry in spec:
|
|
kind = entry["kind"]
|
|
if kind == "image":
|
|
image = entry.get("image")
|
|
if image is None:
|
|
continue
|
|
items.append({"type": "image", "data": image})
|
|
elif kind == "audio":
|
|
items.append({"type": "audio"})
|
|
elif kind == "video":
|
|
frames = entry.get("frames")
|
|
if frames is None:
|
|
continue
|
|
# a soundtrack gets its own "<Audio j>" label, emitted BEFORE its "<Video k>",
|
|
# which is how the model learns the two belong to the same clip
|
|
if entry.get("audio_data") is not None:
|
|
items.append({"type": "audio"})
|
|
qwen_frames, timestamps = video_presentation_frames(frames)
|
|
items.append({"type": "video", "data": qwen_frames, "timestamps": timestamps})
|
|
return items
|
|
|
|
|
|
def resolve_ref_spec(spec, load_image=None, load_audio=None, load_video=None):
|
|
"""Load every reference's media once, dropping the ones that fail.
|
|
|
|
Returns (resolved, skipped) where `resolved` entries carry the loaded media under
|
|
"image" / "audio_data" / "frames". Both the presentation and the latent blocks are
|
|
built from this single resolved list, so the two can never fall out of order.
|
|
"""
|
|
load_image = load_image or load_input_image
|
|
load_audio = load_audio or load_input_audio
|
|
load_video = load_video or load_input_video
|
|
resolved, skipped = [], []
|
|
for entry in spec:
|
|
kind = entry["kind"]
|
|
if kind == "image":
|
|
image = load_image(entry["filename"])
|
|
if image is None:
|
|
skipped.append((entry["filename"], "image could not be loaded"))
|
|
continue
|
|
resolved.append({**entry, "image": image})
|
|
elif kind == "audio":
|
|
audio = load_audio(entry["filename"])
|
|
if audio is None:
|
|
skipped.append((entry["filename"], "audio could not be loaded"))
|
|
continue
|
|
resolved.append({**entry, "audio_data": audio})
|
|
elif kind == "video":
|
|
frames = load_video(entry["filename"])
|
|
if frames is None:
|
|
skipped.append((entry["filename"],
|
|
"video could not be decoded, or is shorter than 5 frames "
|
|
"(~0.2s) — ffmpeg must be on PATH"))
|
|
continue
|
|
trimmed = trim_video_to_grid(frames)
|
|
if trimmed is None:
|
|
skipped.append((entry["filename"],
|
|
"video is shorter than one 17-frame grid step (~0.9s at 24 fps)"))
|
|
continue
|
|
out = {**entry, "frames": trimmed}
|
|
# an index-paired soundtrack: the clip's own audio, not a standalone reference
|
|
if entry.get("audio"):
|
|
track = load_audio(entry["audio"])
|
|
if track is None:
|
|
skipped.append((entry["audio"],
|
|
"soundtrack could not be loaded — the video reference "
|
|
"is still used, silently"))
|
|
else:
|
|
out["audio_data"] = track
|
|
resolved.append(out)
|
|
return resolved, skipped
|
|
|
|
|
|
def ref_blocks_from_spec(resolved, vae, width, height, audio_vae=None):
|
|
"""Resolved spec -> the DiT's `minimax_refs` blocks, in the same order.
|
|
|
|
`vae` encodes image references; `audio_vae` encodes audio ones. An audio reference
|
|
with no audio VAE connected cannot produce a block, so it is dropped AND returned in
|
|
`skipped` — the caller must say so out loud, because the presentation was already
|
|
baked at encode time and every later "<Audio j>" in the prompt now points one
|
|
reference earlier than the user wrote it.
|
|
|
|
Reference videos are capped by :func:`load_input_video`'s own `max_seconds`, NOT by
|
|
the generation's length (which the reference node uses): the presentation is baked by
|
|
Studio, which never sees a frame count, so a sampler-side cap would silently show Qwen
|
|
more of the clip than the packed rows contain.
|
|
"""
|
|
blocks, skipped = [], []
|
|
for entry in resolved:
|
|
kind = entry["kind"]
|
|
if kind == "image":
|
|
_item, block = image_ref_block(vae, entry["image"], width, height,
|
|
size_mode=entry.get("size", "match"))
|
|
blocks.append(block)
|
|
elif kind == "audio":
|
|
if audio_vae is None:
|
|
skipped.append((entry["filename"],
|
|
"audio reference needs an audio_vae connected"))
|
|
continue
|
|
_item, block = audio_ref_block(audio_vae, entry["audio_data"])
|
|
blocks.append(block)
|
|
elif kind == "video":
|
|
track = entry.get("audio_data")
|
|
if track is not None and audio_vae is None:
|
|
# the clip still goes in, just mute — better than dropping the whole
|
|
# reference over its soundtrack, and the caller reports it either way
|
|
skipped.append((entry["filename"],
|
|
"video reference used WITHOUT its soundtrack — encoding "
|
|
"the audio needs an audio_vae connected"))
|
|
track = None
|
|
_item, block = video_ref_block(vae, entry["frames"],
|
|
audio_vae=audio_vae, audio=track)
|
|
blocks.append(block)
|
|
return blocks, skipped
|
|
|
|
|
|
def token_tags_length(meta):
|
|
"""Length of the ``minimax_token_tags`` run this conditioning carries, or None.
|
|
|
|
The DiT walks the tag vector position-by-position across the text span, so a
|
|
conditioning whose token count no longer matches its tags is a hard failure
|
|
(IndexError mid-forward), not a soft one. Anything that changes the token
|
|
count — appending identity tokens, truncating, concatenating prompts — must
|
|
either update the tags or drop the whole feature for H3.
|
|
"""
|
|
if not isinstance(meta, dict):
|
|
return None
|
|
tags = meta.get("minimax_token_tags")
|
|
if tags is None:
|
|
return None
|
|
try:
|
|
return int(tags.reshape(-1).shape[0])
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def tags_match(cond, meta) -> bool:
|
|
"""True when `meta`'s token tags still line up with `cond`'s token count."""
|
|
n = token_tags_length(meta)
|
|
if n is None:
|
|
return True # no tags to invalidate
|
|
try:
|
|
return int(cond.shape[1]) == n
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def trim_token_tags(meta, want):
|
|
"""Cut the tag vector down to `want` positions, keeping the head aligned.
|
|
|
|
Tags are built from the tokenizer's INPUT sequence (`embeds.shape[1]` in comfy's
|
|
MiniMaxQwen3VL.forward) while the conditioning is what the encoder returned, so an
|
|
image prompt normally comes back one tag long — the vision block's expansion does not
|
|
survive the encode 1:1. Position i still means position i in both, so trimming the tail
|
|
is exact; the alternative (dropping the tags) would throw away the vision span's
|
|
modality-0 marks and hand the image tokens to the DiT's text modulation.
|
|
"""
|
|
n = token_tags_length(meta)
|
|
if n is None or want < 0 or want >= n:
|
|
return meta
|
|
out = dict(meta)
|
|
out["minimax_token_tags"] = meta["minimax_token_tags"].reshape(-1)[:int(want)]
|
|
return out
|
|
|
|
|
|
def extend_token_tags(meta, added_tokens, tag=1):
|
|
"""Grow the tag vector by `added_tokens` positions so appended tokens stay legal.
|
|
|
|
Tag 1 is the text modality (0 = video, 2 = audio) — appended context tokens
|
|
are text-like, so they take the text adaLN row.
|
|
"""
|
|
n = token_tags_length(meta)
|
|
if n is None or added_tokens <= 0:
|
|
return meta
|
|
out = dict(meta)
|
|
tags = meta["minimax_token_tags"].reshape(-1)
|
|
pad = torch.full((int(added_tokens),), int(tag), dtype=tags.dtype, device=tags.device)
|
|
out["minimax_token_tags"] = torch.cat([tags, pad], dim=0)
|
|
return out
|
|
|
|
|
|
# ── patching H3 in place ────────────────────────────────────────────────────
|
|
#
|
|
# Patching a module's FORWARD rather than replacing the module is what keeps the state
|
|
# dict intact: a wrapper module renames every weight beneath it, and ComfyUI records
|
|
# weight names off the LIVE tree, so an unwrap by any other patcher sharing this model
|
|
# would leave a restore walking a path that no longer exists.
|
|
|
|
|
|
def unwrapped_forward(module):
|
|
"""The module's own ``forward``, never a FunPack wrapper already sitting on it.
|
|
|
|
ComfyUI restores an object patch by writing the previous value back, and the previous
|
|
value of ``<module>.forward`` may itself be one of our wrappers from a run that is still
|
|
winding down. Capturing that would nest wrappers on every scene and eventually stack a
|
|
scale per scene, so the chain is followed back to the real bound method every time.
|
|
"""
|
|
fn = getattr(module, "forward", module)
|
|
seen = 0
|
|
while getattr(fn, "_funpack_wrapper", False) and seen < 16:
|
|
fn = getattr(fn, "original", fn)
|
|
seen += 1
|
|
return fn
|
|
|
|
|
|
def patch_module_forward(patcher, key, wrapper):
|
|
"""Install `wrapper` as `<key>.forward` through ComfyUI's own object-patch mechanism.
|
|
|
|
Patching the FORWARD rather than replacing the module is what keeps the state dict
|
|
intact: a wrapper module would rename every weight under `key` to `<key>.inner.*`, and
|
|
ComfyUI captures weight names off the live tree, so an unwrap by any other patcher
|
|
sharing this model would leave a restore walking a `.inner` that no longer exists.
|
|
"""
|
|
patcher.add_object_patch(f"{key}.forward", wrapper)
|
|
|
|
|
|
# ── final layer: the one video-only edit ─────────────────────────────────────
|
|
#
|
|
# Every block shares one attention pass, so anything written into the video rows is read by
|
|
# the audio rows in the next block — which is what made post-block injection corrupt the
|
|
# soundtrack. The final layer is past the last attention: video and audio are modulated on
|
|
# separate rows and leave through separate heads, so an edit to the video branch here has
|
|
# no path to the audio at all.
|
|
#
|
|
# The arithmetic below MIRRORS ``FinalLayer.forward`` using the module's own submodules
|
|
# rather than wrapping its result, because the value being changed is consumed inside that
|
|
# method. It is checked against the module's attributes before it runs, so a changed upstream
|
|
# shape declines instead of computing something else.
|
|
#
|
|
# WHAT IS SCALED, and why it is not the obvious thing. The layer computes
|
|
# ``norm(x) * (1 + scale[row]) + shift[row]``. Scaling ``scale`` itself reads as the natural
|
|
# edit, and it was the first build — but H3's ``scale`` is NEGATIVE here, so raising the
|
|
# multiplier SHRANK ``1 + scale`` and the dial ran backwards: measured on a rental, 0.8 gave
|
|
# more contrast and 1.2 less, and 1.15 was nearly inert. Scaling the whole ``(1 + scale)``
|
|
# term instead is monotone in magnitude whichever sign the checkpoint carries, so "above 1.0
|
|
# is more" is true by construction rather than by luck. ``shift`` is left alone either way:
|
|
# it is the bias, and moving it walks every channel off its trained centre.
|
|
_FINAL_LAYER_PARTS = ("norm", "adaln_proj", "video_out", "audio_out")
|
|
|
|
|
|
class FinalLayerVideoScale:
|
|
"""Scales how strongly the final layer reads the accumulated video rows.
|
|
|
|
Upstream's own forward does all the work — every argument passes straight through, so
|
|
its signature (it grew ``sigma, sample_sigmas, shifts`` for PDD heads) and its head
|
|
branches are inherited, never copied. The one change: while it runs, the final norm's
|
|
output for the VIDEO slice is multiplied by the scale, which is exactly
|
|
``norm(x) * (1 + scale) * s + shift`` on video rows and nothing on audio rows.
|
|
"""
|
|
|
|
_funpack_wrapper = True
|
|
|
|
def __init__(self, inner, scale):
|
|
self.inner = inner
|
|
self.original = unwrapped_forward(inner)
|
|
self.scale = float(scale)
|
|
|
|
def __call__(self, x, *args, **kwargs):
|
|
norm = getattr(self.inner, "norm", None)
|
|
try:
|
|
va, vb = int(args[1][0]), int(args[1][1]) # video_seg = (start, stop, row)
|
|
except Exception: # noqa: BLE001
|
|
va = vb = None
|
|
if self.scale == 1.0 or norm is None or va is None or vb <= va:
|
|
return self.original(x, *args, **kwargs)
|
|
video_ptr, video_rows = x[va:vb].data_ptr(), vb - va
|
|
gain = self.scale
|
|
had = "forward" in norm.__dict__
|
|
prev = norm.__dict__.get("forward")
|
|
inner_forward = norm.forward
|
|
|
|
def forward(inp, *a, **k):
|
|
out = inner_forward(inp, *a, **k)
|
|
if inp.data_ptr() == video_ptr and inp.shape[0] == video_rows:
|
|
out = out * gain
|
|
return out
|
|
|
|
norm.forward = forward
|
|
try:
|
|
return self.original(x, *args, **kwargs)
|
|
finally:
|
|
if had:
|
|
norm.forward = prev
|
|
else:
|
|
del norm.forward
|
|
|
|
|
|
def apply_final_video_scale(model, scale):
|
|
"""Attach the video-only final-layer scale. Returns (model, note)."""
|
|
if float(scale) == 1.0:
|
|
return model, None
|
|
key = "diffusion_model.final_layer"
|
|
try:
|
|
inner = model.get_model_object(key)
|
|
except Exception as error:
|
|
return model, f"Final-layer video scale skipped — no final layer to patch ({error})."
|
|
if not all(hasattr(inner, part) for part in _FINAL_LAYER_PARTS):
|
|
return model, ("Final-layer video scale skipped — this final layer is not the shape "
|
|
"the edit was written against.")
|
|
patched = model.clone()
|
|
patch_module_forward(patched, key, FinalLayerVideoScale(inner, scale))
|
|
return patched, (f"Final-layer video scale {float(scale):g} — past the last attention, "
|
|
f"so the audio stream cannot see it.")
|