digital-garbageandClaude Opus 5.5
5787190b3f
H3's latent upscaler loads through the H3 module (ComfyUI's loader has no branch for it); an upscale resizes H3's anchor pin so pass 2 stays pinned. FunPack Sampler takes an optional wired schedule, run exactly as written. New preset: H3 Reference to Video + second pass. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
237 lines
10 KiB
Python
237 lines
10 KiB
Python
"""MiniMax H3's latent upscaler: the 3D-conv resizer, which ComfyUI cannot load.
|
|
|
|
ComfyUI's Load Latent Upscale Model knows Lightricks' LatentUpsampler and two
|
|
Hunyuan SR variants (checked on upstream master, 2026-09-27); this is a fourth,
|
|
and the only published 24-channel one, so without it H3 has no second-pass
|
|
upscaler at all.
|
|
|
|
The architecture is dictated by the checkpoint: every module name matches a key
|
|
in the published state dict and each shape is read back off the weights
|
|
(`detect_config`). Latent normalisation uses the upscaler's own published
|
|
per-channel statistics -- H3's VAE statistics are a different thing.
|
|
|
|
Reference implementation: LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler.
|
|
Reconstructed rather than vendored -- that repository ships no licence. Ported
|
|
from FunPack v4's h3_latent_upscaler.py.
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# Per-channel statistics the upscaler was trained against, from the model's own card. The
|
|
# net sees normalised latents and returns normalised latents; skipping this does not
|
|
# degrade the output so much as produce a different distribution entirely.
|
|
LATENTS_MEAN = [
|
|
0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075,
|
|
-0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975,
|
|
-0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923,
|
|
-0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543,
|
|
-0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279,
|
|
-0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264,
|
|
]
|
|
LATENTS_STD = [
|
|
1.2223774194717407, 1.2767263650894165, 1.6831774711608887, 1.7549455165863037,
|
|
1.5636216402053833, 2.194143533706665, 0.9653137922286987, 1.0569885969161987,
|
|
0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647,
|
|
0.7996809482574463, 0.44988900423049927, 0.7197399735450745, 0.6936293244361877,
|
|
2.961095094680786, 2.7694199085235596, 3.0496184825897217, 2.1088054180145264,
|
|
3.276226282119751, 3.1627357006073, 2.2816812992095947, 2.6127843856811523,
|
|
]
|
|
|
|
_EMBED_DIM = 64
|
|
_GROUPS = 32
|
|
|
|
|
|
def _norm(channels):
|
|
return nn.GroupNorm(_GROUPS, channels)
|
|
|
|
|
|
class _ResBlockEmb3D(nn.Module):
|
|
"""Residual block conditioned on the scale factor (FiLM-style shift/scale)."""
|
|
|
|
def __init__(self, channels, emb_channels, out_channels=None):
|
|
super().__init__()
|
|
self.out_channels = out_channels or channels
|
|
self.in_layers = nn.Sequential(
|
|
_norm(channels), nn.SiLU(), nn.Conv3d(channels, self.out_channels, 3, padding=1))
|
|
self.emb_layers = nn.Sequential(
|
|
nn.SiLU(), nn.Linear(emb_channels, 2 * self.out_channels))
|
|
self.out_norm = _norm(self.out_channels)
|
|
self.out_layers = nn.Sequential(
|
|
nn.SiLU(), nn.Dropout(p=0.0),
|
|
nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1))
|
|
self.skip = (nn.Conv3d(channels, self.out_channels, 1)
|
|
if self.out_channels != channels else nn.Identity())
|
|
|
|
def forward(self, x, emb):
|
|
h = self.in_layers(x)
|
|
emb_out = self.emb_layers(emb).type(h.dtype)
|
|
while emb_out.dim() < h.dim():
|
|
emb_out = emb_out[..., None]
|
|
scale, shift = torch.chunk(emb_out, 2, dim=1)
|
|
h = self.out_norm(h) * (1 + scale) + shift
|
|
return self.skip(x) + self.out_layers(h)
|
|
|
|
|
|
class _TemporalConv(nn.Module):
|
|
"""Depthwise conv along time only — what makes this a video upscaler rather than a
|
|
per-frame image one, and why an upscaled clip does not shimmer frame to frame."""
|
|
|
|
def __init__(self, channels, kernel_size=5):
|
|
super().__init__()
|
|
self.norm = _norm(channels)
|
|
self.dwconv = nn.Conv3d(channels, channels, kernel_size=(kernel_size, 1, 1),
|
|
padding=(kernel_size // 2, 0, 0), groups=channels)
|
|
self.pwconv = nn.Conv3d(channels, channels, kernel_size=1)
|
|
|
|
def forward(self, x):
|
|
return x + self.pwconv(self.dwconv(F.silu(self.norm(x))))
|
|
|
|
|
|
class H3LatentResizer3D(nn.Module):
|
|
"""Refine, resample, refine. The resample is a plain trilinear interpolation; the
|
|
learned part is what the blocks do on either side of it."""
|
|
|
|
def __init__(self, in_channels=24, in_blocks=12, out_blocks=12, channels=512,
|
|
temporal_every=2, temporal_kernel=5):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1)
|
|
self.embed = nn.Sequential(
|
|
nn.Linear(1, _EMBED_DIM), nn.SiLU(), nn.Linear(_EMBED_DIM, _EMBED_DIM))
|
|
|
|
def stack(count):
|
|
blocks = nn.ModuleList()
|
|
for index in range(count):
|
|
blocks.append(_ResBlockEmb3D(channels, _EMBED_DIM))
|
|
if temporal_every > 0 and index % temporal_every == 0:
|
|
blocks.append(_TemporalConv(channels, temporal_kernel))
|
|
return blocks
|
|
|
|
# Interleaved in one flat list, so the index a key names depends on this exact
|
|
# construction order — it is part of the format, not a style choice.
|
|
self.in_blocks = stack(in_blocks)
|
|
self.out_blocks = stack(out_blocks)
|
|
self.norm_out = _norm(channels)
|
|
self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1)
|
|
|
|
def _run(self, blocks, x, emb):
|
|
for block in blocks:
|
|
if isinstance(block, _ResBlockEmb3D):
|
|
x = block(x, emb.expand(x.shape[0], -1))
|
|
else:
|
|
x = block(x)
|
|
return x
|
|
|
|
def forward(self, x, scale=None, target_size=None):
|
|
if target_size is not None:
|
|
size = tuple(int(v) for v in target_size)
|
|
elif scale is not None:
|
|
size = tuple(int(round(dim * scale)) for dim in x.shape[-3:])
|
|
else:
|
|
return x
|
|
if size == tuple(x.shape[-3:]):
|
|
return x
|
|
|
|
# The block conditioning is scale-1, so an unchanged size embeds as zero.
|
|
embed_in = torch.tensor([[0.0 if scale is None else float(scale) - 1.0]],
|
|
dtype=x.dtype, device=x.device)
|
|
emb = self.embed(embed_in)
|
|
|
|
x = self._run(self.in_blocks, self.conv_in(x), emb)
|
|
x = F.interpolate(x, size=size, mode="trilinear", align_corners=False)
|
|
x = self._run(self.out_blocks, x, emb)
|
|
return self.conv_out(F.silu(self.norm_out(x)))
|
|
|
|
# Marks this as an upscaler that owns its normalisation and takes a scale factor,
|
|
# so the caller does not have to know which architecture it loaded.
|
|
def funpack_latent_upscale(self, x, scale=2.0):
|
|
mean = torch.tensor(LATENTS_MEAN, dtype=x.dtype, device=x.device).view(1, -1, 1, 1, 1)
|
|
std = torch.tensor(LATENTS_STD, dtype=x.dtype, device=x.device).view(1, -1, 1, 1, 1)
|
|
if mean.shape[1] != x.shape[1]:
|
|
raise ValueError(
|
|
f"this upscaler's statistics are {mean.shape[1]}-channel and the latent is "
|
|
f"{x.shape[1]}-channel")
|
|
# Spatial only: H3's video latent is on a 5k+2 time grid and resampling it would
|
|
# leave a frame count the VAE has no defined decode for.
|
|
#
|
|
# Latent H/W snap to even numbers: H3 patchifies 2x2, so an odd latent dimension is
|
|
# not a size the model can take at all. At scale 2.0 this is a no-op; it only bites
|
|
# on the in-between factors this net also supports.
|
|
def side(dim):
|
|
return max(2, int(round(dim * scale / 2.0)) * 2)
|
|
|
|
target = (int(x.shape[2]), side(x.shape[3]), side(x.shape[4]))
|
|
out = self((x - mean) / std, scale=scale, target_size=target)
|
|
return out * std + mean
|
|
|
|
|
|
def is_h3_latent_upscaler(sd):
|
|
"""Whether a state dict is this architecture. Checked before ComfyUI's own loader,
|
|
which has no branch for it and fails with a Python error rather than a diagnosis."""
|
|
sd = strip_prefix(sd)
|
|
return "conv_in.weight" in sd and any(k.startswith("in_blocks.") for k in sd) \
|
|
and "post_upsample_res_blocks.0.conv2.bias" not in sd
|
|
|
|
|
|
def strip_prefix(sd):
|
|
"""Some checkpoints wrap the net under `upscaler.` (and a `model` top-level key)."""
|
|
if isinstance(sd, dict) and "model" in sd and isinstance(sd["model"], dict):
|
|
sd = sd["model"]
|
|
if any(str(k).startswith("upscaler.") for k in sd):
|
|
return {k[len("upscaler."):]: v for k, v in sd.items() if str(k).startswith("upscaler.")}
|
|
return sd
|
|
|
|
|
|
def detect_config(sd):
|
|
"""Read the architecture back off the weights — width, depth and temporal layout are
|
|
all recoverable, so there is no config to carry and no version to get wrong."""
|
|
import re
|
|
|
|
config = {"in_channels": 24, "in_blocks": 12, "out_blocks": 12, "channels": 512,
|
|
"temporal_every": 2, "temporal_kernel": 5}
|
|
conv_in = sd.get("conv_in.weight")
|
|
if conv_in is not None:
|
|
config["channels"] = int(conv_in.shape[0])
|
|
config["in_channels"] = int(conv_in.shape[1])
|
|
|
|
res, temporal = {"in": set(), "out": set()}, False
|
|
for key in sd:
|
|
match = re.match(r"(in|out)_blocks\.(\d+)\.in_layers\.", str(key))
|
|
if match:
|
|
res[match.group(1)].add(int(match.group(2)))
|
|
if str(key).endswith("dwconv.weight"):
|
|
temporal = True
|
|
config["temporal_kernel"] = int(sd[key].shape[2])
|
|
if res["in"]:
|
|
config["in_blocks"] = len(res["in"])
|
|
if res["out"]:
|
|
config["out_blocks"] = len(res["out"])
|
|
if not temporal:
|
|
config["temporal_every"] = 0
|
|
return config
|
|
|
|
|
|
def load_upscaler(sd):
|
|
"""The `latent_upscaler` capability: a built upscaler, or None if this state
|
|
dict is not this architecture. Returns an object with
|
|
`funpack_latent_upscale(x, scale)`, which is all a caller needs."""
|
|
if not is_h3_latent_upscaler(sd):
|
|
return None
|
|
return from_state_dict(sd)
|
|
|
|
|
|
def from_state_dict(sd, dtype=None):
|
|
"""Build the resizer this state dict describes and load it."""
|
|
sd = strip_prefix(sd)
|
|
if any(".q.weight" in str(k) or ".proj_out.weight" in str(k) for k in sd):
|
|
raise ValueError(
|
|
"this H3 latent upscaler has attention blocks, which this loader does not build "
|
|
"— use the companion custom node for that checkpoint")
|
|
model = H3LatentResizer3D(**detect_config(sd))
|
|
model.load_state_dict(sd, strict=True)
|
|
model.eval().requires_grad_(False)
|
|
if dtype is not None:
|
|
model = model.to(dtype=dtype)
|
|
return model
|