diff --git a/CHANGELOG.md b/CHANGELOG.md index 631d9737..cc8638cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,16 @@ backend installed on its own. Kernel and block map vendored from LightX2V (Apache-2.0) via ComfyUI-H3-SLA-Attention (MIT). +- **MiniMax H3's latent upscaler loads without its custom node pack**, so `second_pass_op` + and segmented detailing work on H3 — the operation needs an upsampler whose latents are + the model's width, and the only published 24-channel one shipped behind its own node. + Architecture and normalisation statistics are read off the checkpoint; verified + bit-identical against the reference implementation. +- **H3 scenes continue from the previous one.** The previous scene's last latent frame + becomes the next scene's frame-0 keyframe pin — the only continuity conditioning H3 + accepts. A carried latent tail is not conditioning on this model, so the seam matched + while the rest of the shot knew nothing about the scene before it. + ### Fixed - **Multi-scene chains work with a second pass again.** `second_pass_op="upscale_2x"` hands back a scene at twice the latent size, but every later scene is still built from the latent diff --git a/detailing.py b/detailing.py index 93181569..515662d5 100644 --- a/detailing.py +++ b/detailing.py @@ -204,6 +204,16 @@ def load_latent_upsampler(model_name): model_path = folder_paths.get_full_path_or_raise("latent_upscale_models", model_name) sd, metadata = comfy.utils.load_torch_file(model_path, safe_load=True, return_metadata=True) + try: + from . import h3_latent_upscaler as _h3up + except ImportError: + import h3_latent_upscaler as _h3up + if _h3up.is_h3_latent_upscaler(sd): + # MiniMax H3's own upscaler. ComfyUI has no branch for this architecture, so + # without this the only published 24-channel upsampler needs a custom node pack — + # and second_pass_op on H3 needs a 24-channel one to run at all. + return _h3up.from_state_dict(sd, dtype=comfy.model_management.vae_dtype( + allowed_dtypes=[torch.bfloat16, torch.float32])) if "post_upsample_res_blocks.0.conv2.bias" not in sd: # Not Lightricks' architecture. ComfyUI's own loader knows the others (and will know # the next one), so hand it over rather than declaring the file unusable here. @@ -472,14 +482,18 @@ def _run_upsampler(upsampler, video_crop, vae, debug=False): device = comfy.model_management.get_torch_device() model_dtype = next(module.parameters()).dtype in_dtype, in_device = video_crop.dtype, video_crop.device - stats = _per_channel_stats(vae) + # An upscaler that owns its normalisation (H3's) says so, because its statistics are the + # model's own — running it through the VAE's per-channel stats instead would hand it a + # distribution it never saw. + own_norm = getattr(module, "funpack_latent_upscale", None) + stats = None if callable(own_norm) else _per_channel_stats(vae) try: module.to(device) x = video_crop.to(device=device, dtype=model_dtype) if stats is not None: x = stats.un_normalize(x) with torch.no_grad(): - x = module(x) + x = own_norm(x, scale=2.0) if callable(own_norm) else module(x) if stats is not None: x = stats.normalize(x) finally: diff --git a/h3_latent_upscaler.py b/h3_latent_upscaler.py new file mode 100644 index 00000000..d11ab447 --- /dev/null +++ b/h3_latent_upscaler.py @@ -0,0 +1,222 @@ +"""MiniMax H3's latent upscaler: the 3D-conv resizer, loadable without a node pack. + +ComfyUI's own Load Latent Upscale Model knows three architectures (Lightricks' +LatentUpsampler and two Hunyuan SR variants) and this is a fourth, so an H3 upsampler +file could not be loaded at all — which is what blocked `second_pass_op` on H3: the +operation needs an upsampler whose latents are the same width as the model's, and the only +published 24-channel one lives behind its own custom node. + +The architecture is dictated by the checkpoint, not chosen here: every module name below +exists to match a key in the published state dict, and the shape of each is read back off +the weights (`detect_config`) rather than configured. Latent normalisation uses the model's +own published per-channel statistics — H3's VAE statistics are a different thing and would +be wrong here. + +Reference implementation: LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler (the companion node +for LBH-123-AI/Minimax_h3_latent_Upscaler). Reconstructed rather than vendored — that +repository ships no licence. +""" +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. + target = (int(x.shape[2]), int(round(x.shape[3] * scale)), int(round(x.shape[4] * scale))) + 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 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 diff --git a/tests/test_h3_latent_upscaler.py b/tests/test_h3_latent_upscaler.py new file mode 100644 index 00000000..ac204fdc --- /dev/null +++ b/tests/test_h3_latent_upscaler.py @@ -0,0 +1,93 @@ +"""MiniMax H3's latent upscaler, loaded without its custom node pack. + +The architecture is dictated by the checkpoint, so the tests that matter are the ones that +prove a state dict round-trips: build a small net, save its keys, rebuild from them alone, +and load strictly. If the module names or the block interleaving drift, that fails. +""" +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import h3_latent_upscaler as h3up # noqa: E402 + + +def _small(**kw): + cfg = dict(in_channels=24, in_blocks=2, out_blocks=2, channels=32, + temporal_every=2, temporal_kernel=5) + cfg.update(kw) + return h3up.H3LatentResizer3D(**cfg) + + +def test_the_architecture_is_read_back_off_the_weights(): + """No config travels with these files, so every dimension has to be recoverable.""" + sd = _small(in_blocks=3, out_blocks=2, channels=64, temporal_kernel=3).state_dict() + cfg = h3up.detect_config(sd) + assert cfg == {"in_channels": 24, "in_blocks": 3, "out_blocks": 2, "channels": 64, + "temporal_every": 2, "temporal_kernel": 3} + + +def test_a_checkpoint_without_temporal_convs_is_detected_as_such(): + sd = _small(temporal_every=0).state_dict() + assert h3up.detect_config(sd)["temporal_every"] == 0 + + +def test_a_state_dict_rebuilds_into_the_same_net_strictly(): + """strict=True is the point: a silently-renamed module would load as random weights.""" + original = _small(in_blocks=3, channels=64) + rebuilt = h3up.from_state_dict(original.state_dict()) + assert isinstance(rebuilt, h3up.H3LatentResizer3D) + for key, tensor in original.state_dict().items(): + assert torch.equal(rebuilt.state_dict()[key], tensor) + + +def test_the_wrapper_prefix_some_checkpoints_use_is_stripped(): + sd = {"upscaler." + k: v for k, v in _small().state_dict().items()} + assert h3up.is_h3_latent_upscaler(sd) + assert h3up.from_state_dict(sd) is not None + + +def test_it_is_told_apart_from_the_architectures_comfyui_already_knows(): + assert not h3up.is_h3_latent_upscaler({"post_upsample_res_blocks.0.conv2.bias": None, + "conv_in.weight": None, "in_blocks.0.x": None}) + assert not h3up.is_h3_latent_upscaler({"blocks.0.block.0.conv.weight": None}) + assert h3up.is_h3_latent_upscaler(_small().state_dict()) + + +def test_upscaling_doubles_the_spatial_axes_and_leaves_time_alone(): + """H3's video latent is on a 5k+2 time grid — resampling time would leave a frame count + the VAE has no defined decode for.""" + model = _small() + latent = torch.randn(1, 24, 3, 8, 12) + out = model.funpack_latent_upscale(latent, scale=2.0) + assert out.shape == (1, 24, 3, 16, 24) + + +def test_a_latent_of_the_wrong_width_is_refused_by_name(): + """LTX latents are 128-channel; this model's statistics are 24.""" + model = _small() + with pytest.raises(ValueError) as exc: + model.funpack_latent_upscale(torch.randn(1, 128, 2, 8, 8), scale=2.0) + assert "24-channel" in str(exc.value) and "128-channel" in str(exc.value) + + +def test_the_published_statistics_cover_every_channel(): + assert len(h3up.LATENTS_MEAN) == len(h3up.LATENTS_STD) == 24 + + +def test_the_upscaler_reports_the_latent_width_it_takes(): + """detailing refuses a mismatched upsampler by naming both numbers; that check reads + the width off this model.""" + import detailing + assert detailing.upsampler_in_channels(_small()) == 24 + + +def test_an_attention_checkpoint_says_what_to_do_instead_of_loading_wrong(): + sd = _small().state_dict() + sd["in_blocks.0.q.weight"] = torch.zeros(1) + with pytest.raises(ValueError) as exc: + h3up.from_state_dict(sd) + assert "custom node" in str(exc.value)