# Learned spatial upscaler for MiniMax-H3 video latents (24 channels). # # Network, temporal chunking and checkpoint layout from Comfyui_Minimax_h3_latent_Upscaler # (https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler, MIT License), # nodes/minimax_h3_latent_upscaler_3d.py. # # Latent space: ComfyUI's H3 VAE already returns mean/std-normalized latents (as this repo's # does), and the node normalizes them once more with the same per-channel stats before the # network, undoing it after. The network was trained in that twice-normalized space, so # `upscale()` applies the same extra normalization to this repo's latents; calling the network # on them directly gives a magenta dot lattice. import re import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint if __package__: from .stats import _VIDEO_LATENTS_MEAN, _VIDEO_LATENTS_STD else: from stats import _VIDEO_LATENTS_MEAN, _VIDEO_LATENTS_STD VAE_DOWNSAMPLE = 16 FRAME_PER_TOKEN = (1, 4, 4, 4, 4) # pixel frames per latent in each 17-frame chunk of H3's VAE def latent_times(T, fps): """Centre time (s, from the first frame) of each of T latents of a clip at fps.""" start, out = 0, [] for k in range(T): span = FRAME_PER_TOKEN[k % len(FRAME_PER_TOKEN)] out.append((start + (span - 1) / 2) / fps) start += span return torch.tensor(out, dtype=torch.float64) def time_resample_matrix(t_in, fps_in, t_out, fps_out): """[t_out, t_in] linear interpolation weights from latents at fps_in to latents at fps_out over the same start instant, by their centre times (held at the ends).""" m = torch.zeros(t_out, t_in, dtype=torch.float64) if t_in == 1: m[:, 0] = 1 return m src = latent_times(t_in, fps_in) dst = latent_times(t_out, fps_out).clamp(src[0].item(), src[-1].item()) hi = torch.searchsorted(src, dst, right=True).clamp(1, t_in - 1) lo = hi - 1 w = ((dst - src[lo]) / (src[hi] - src[lo])).clamp(0, 1) rows = torch.arange(t_out) m[rows, lo] = 1 - w m[rows, hi] += w return m def normalization(channels): return nn.GroupNorm(32, channels) def zero_module(module): for p in module.parameters(): p.detach().zero_() return module class ResBlockEmb3D(nn.Module): def __init__(self, channels, emb_channels, dropout=0, out_channels=None): super().__init__() self.out_channels = out_channels or channels self.in_layers = nn.Sequential( normalization(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 = normalization(self.out_channels) self.out_layers = nn.Sequential( nn.SiLU(), nn.Dropout(p=dropout), zero_module(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 len(emb_out.shape) < len(h.shape): emb_out = emb_out[..., None] scale, shift = torch.chunk(emb_out, 2, dim=1) h = self.out_norm(h) * (1 + scale) + shift h = self.out_layers(h) return self.skip(x) + h class TemporalConv(nn.Module): def __init__(self, channels, kernel_size=5): super().__init__() self.norm = normalization(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) nn.init.zeros_(self.pwconv.weight) nn.init.zeros_(self.pwconv.bias) def forward(self, x): return x + self.pwconv(self.dwconv(F.silu(self.norm(x)))) class MiniMaxH3LatentUpscaler(nn.Module): """3D-conv latent resizer: conv_in -> in_blocks -> trilinear resize -> out_blocks -> conv_out. The upscale factor is fed to every ResBlock through a small embedding, so one checkpoint covers 1x-4x. Only H and W change; T is preserved. temporal=True adds frame-rate upscaling (e.g. 8 / 12 / 16 -> 24 fps) to the same network: a second embedding of the temporal factor (fps_out / fps_in), whose last layer starts at zero, is added to the spatial one (so a spatial checkpoint loads unchanged and behaves the same at equal frame rates), and the middle resize first interpolates the features in time by each latent's centre time (H3 latents are not evenly spaced: the first of every 17-frame chunk covers 1 frame, the others 4), then in H and W. """ def __init__(self, in_channels=24, in_blocks=12, out_blocks=12, channels=512, dropout=0.1, temporal_every=2, temporal_kernel=5, temporal=False): super().__init__() self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1) embed_dim = 64 self.embed = nn.Sequential(nn.Linear(1, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim)) self.t_embed = None if temporal: self.t_embed = nn.Sequential(nn.Linear(1, embed_dim), nn.SiLU(), zero_module(nn.Linear(embed_dim, embed_dim))) self.in_blocks = self._stage(in_blocks, channels, embed_dim, dropout, temporal_every, temporal_kernel) self.out_blocks = self._stage(out_blocks, channels, embed_dim, dropout, temporal_every, temporal_kernel) self.norm_out = normalization(channels) self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1) self.temporal_kernel = temporal_kernel if temporal_every > 0 else 0 self.gradient_checkpointing = False @staticmethod def _stage(num_blocks, channels, embed_dim, dropout, temporal_every, temporal_kernel): blocks = nn.ModuleList() for b in range(num_blocks): blocks.append(ResBlockEmb3D(channels, embed_dim, dropout)) if temporal_every > 0 and b % temporal_every == 0: blocks.append(TemporalConv(channels, temporal_kernel)) return blocks def _run_blocks(self, blocks, x, emb): for block in blocks: if isinstance(block, ResBlockEmb3D): args = (x, emb.expand(x.shape[0], -1)) else: args = (x,) x = checkpoint(block, *args, use_reentrant=False) if self.gradient_checkpointing and self.training else block(*args) return x def _forward_segment(self, x, scale, size, fps=None): emb = self.embed(torch.tensor([[scale - 1]], dtype=x.dtype, device=x.device)) if self.t_embed is not None: t_scale = fps[1] / fps[0] if fps else 1.0 emb = emb + self.t_embed(torch.tensor([[t_scale - 1]], dtype=x.dtype, device=x.device)) x = self._run_blocks(self.in_blocks, self.conv_in(x), emb) if fps and fps[0] != fps[1]: # in time by centre times; trilinear below then only resizes H and W m = time_resample_matrix(x.shape[2], fps[0], size[0], fps[1]).to(device=x.device, dtype=x.dtype) x = torch.einsum("oi,bcihw->bcohw", m, x) x = F.interpolate(x, size=size, mode="trilinear", align_corners=False) x = self._run_blocks(self.out_blocks, x, emb) return self.conv_out(F.silu(self.norm_out(x))) def forward(self, x, scale, target_size, chunk=32, fps=None): """x: [B, 24, T, H, W] normalized latents; target_size: (T, H_out, W_out); fps: (in, out) frame rates of a frame-rate upscale (both clips start at the same instant), None or equal rates for a spatial one. Long clips run in temporal chunks of `chunk` latent frames with replicate padding and linearly weighted overlaps (overlap = the temporal kernel), so VRAM stays bounded and the last frames do not flicker. chunk=None runs the whole clip at once. """ frame_rate = fps is not None and fps[0] != fps[1] if tuple(target_size) == tuple(x.shape[-3:]) and not frame_rate: return x B, C, T, _, _ = x.shape if frame_rate: if chunk is not None and T > chunk: raise ValueError("temporal chunking is not implemented for frame-rate upscaling; use chunk=None") return self._forward_segment(x, scale, target_size, fps) if chunk is None or T <= chunk: return self._forward_segment(x, scale, target_size) overlap = self.temporal_kernel x_padded = F.pad(x, (0, 0, 0, 0, overlap, overlap), mode="replicate") out = torch.zeros(B, C, T, target_size[-2], target_size[-1], device=x.device, dtype=x.dtype) weight_sum = torch.zeros(1, 1, T, 1, 1, device=x.device, dtype=x.dtype) for seg_start in range(0, T, chunk): seg_end = min(T, seg_start + chunk) out_start, out_end = max(0, seg_start - overlap), min(T, seg_end + overlap) lo, hi = max(0, out_start - overlap), min(T + 2 * overlap, out_end + overlap) seg_out = self._forward_segment(x_padded[:, :, lo:hi].contiguous(), scale, (hi - lo, target_size[-2], target_size[-1])) s0 = out_start + overlap - lo valid = seg_out[:, :, s0:s0 + out_end - out_start] weight = torch.ones(out_end - out_start, device=x.device, dtype=x.dtype) if seg_start > out_start: n = seg_start - out_start weight[:n] = torch.arange(1, n + 1, device=x.device, dtype=x.dtype) / (n + 1) if out_end > seg_end: n = out_end - seg_end weight[-n:] = torch.arange(n, 0, -1, device=x.device, dtype=x.dtype) / (n + 1) out[:, :, out_start:out_end] += valid * weight.view(1, 1, -1, 1, 1) weight_sum[:, :, out_start:out_end] += weight.view(1, 1, -1, 1, 1) return out / weight_sum.clamp(min=1e-8) def upscale(self, latents, scale, target_size, chunk=32, fps=None): """Upscale normalized H3 latents (as the VAE/pipeline produce them) to target_size (T, H, W).""" dtype = next(self.parameters()).dtype mean = torch.tensor(_VIDEO_LATENTS_MEAN, device=latents.device, dtype=dtype).view(1, -1, 1, 1, 1) std = torch.tensor(_VIDEO_LATENTS_STD, device=latents.device, dtype=dtype).view(1, -1, 1, 1, 1) out = self((latents.to(dtype) - mean) / std, scale, target_size, chunk=chunk, fps=fps) return (out * std + mean).to(latents.dtype) @staticmethod def config_from_state_dict(sd): """Architecture is read off the checkpoint (block counts, channels, temporal kernel).""" cfg = {"in_channels": sd["conv_in.weight"].shape[1], "channels": sd["conv_in.weight"].shape[0]} for stage in ("in_blocks", "out_blocks"): cfg[stage] = len({m.group(1) for k in sd if (m := re.match(rf"{stage}\.(\d+)\.in_layers\.", k))}) dwconv = [v for k, v in sd.items() if k.endswith("dwconv.weight")] cfg["temporal_every"] = 2 if dwconv else 0 cfg["temporal_kernel"] = dwconv[0].shape[2] if dwconv else 5 cfg["temporal"] = any(k.startswith("t_embed.") for k in sd) return cfg @classmethod def from_file(cls, path, device="cuda", dtype=torch.float16): from safetensors.torch import load_file sd = load_file(path, device="cpu") if path.endswith(".safetensors") else torch.load(path, map_location="cpu", weights_only=False) sd = sd.get("model", sd) if any(k.startswith("upscaler.") for k in sd): sd = {k[len("upscaler."):]: v for k, v in sd.items() if k.startswith("upscaler.")} sd = {k: v.to(torch.float16) if v.dtype == torch.float8_e4m3fn else v for k, v in sd.items()} model = cls(**cls.config_from_state_dict(sd)) model.load_state_dict(sd, strict=True) return model.to(device=device, dtype=dtype).eval().requires_grad_(False) class MiniMaxH3TemporalUpscaler(MiniMaxH3LatentUpscaler): """H3 backbone with temporal feature interpolation instead of spatial resizing. The parameter layout is identical to MiniMaxH3LatentUpscaler, so its spatial checkpoint can initialize this model. The learned task still changes and needs temporal fine tuning. """ fixed_temporal_scale = 2 def _forward_segment(self, x, scale, size): emb = self.embed(torch.tensor([[scale - 1]], dtype=x.dtype, device=x.device)) x = self._run_blocks(self.in_blocks, self.conv_in(x), emb) x = F.interpolate(x, size=size, mode="trilinear", align_corners=True) x = self._run_blocks(self.out_blocks, x, emb) return self.conv_out(F.silu(self.norm_out(x))) def forward(self, x, scale=2.0, target_size=None, chunk=None): if x.ndim != 5 or x.shape[1] != self.conv_in.in_channels or x.shape[2] < 2: raise ValueError(f"expected [B,{self.conv_in.in_channels},T>=2,H,W], got {tuple(x.shape)}") expected = (2 * x.shape[2] - 1, x.shape[3], x.shape[4]) if scale != 2.0: raise ValueError(f"temporal interpolation has fixed x2 scale, got {scale}") if target_size is not None and tuple(target_size) != expected: raise ValueError(f"temporal interpolation gives {expected}, got target_size={tuple(target_size)}") if chunk is not None: raise ValueError("temporal chunking is not supported; use chunk=None") return self._forward_segment(x, scale, expected) def upscale(self, latents, scale=2.0, target_size=None, chunk=None): return super().upscale(latents, scale, target_size or (2 * latents.shape[2] - 1, latents.shape[3], latents.shape[4]), chunk=chunk) def target_latent_size(latent_h, latent_w, megapixels=None, scale=None, width=None, height=None, align=32): """Output latent (H, W) and the effective scale, as the ComfyUI node computes them. Exactly one of `megapixels`, `scale` or (`width`, `height`) in pixels. The pixel target is rounded to `align` and then to the VAE's 16px grid. """ if megapixels is not None: aspect = latent_w / latent_h h_px = (megapixels * 1024 * 1024 / aspect) ** 0.5 w_px = h_px * aspect elif scale is not None: w_px, h_px = latent_w * VAE_DOWNSAMPLE * scale, latent_h * VAE_DOWNSAMPLE * scale else: w_px, h_px = float(width), float(height) effective_scale = (w_px / (latent_w * VAE_DOWNSAMPLE) + h_px / (latent_h * VAE_DOWNSAMPLE)) / 2.0 w_px = round(round(w_px / align) * align / VAE_DOWNSAMPLE) * VAE_DOWNSAMPLE h_px = round(round(h_px / align) * align / VAE_DOWNSAMPLE) * VAE_DOWNSAMPLE h_out, w_out = max(1, h_px // VAE_DOWNSAMPLE), max(1, w_px // VAE_DOWNSAMPLE) if effective_scale < 1.0 and (w_out < latent_w or h_out < latent_h): raise ValueError("the latent upscaler only upscales (effective scale >= 1)") return (h_out, w_out), effective_scale if __name__ == "__main__": model = MiniMaxH3LatentUpscaler( channels=512, in_blocks=12, out_blocks=12 ).cuda().eval() low = torch.randn(1, 24, 5, 12, 20, device="cuda") # [B, C, T, H, W] with torch.no_grad(): out = model(low, scale=2.0, target_size=(5, 24, 40), chunk=None) print(out.shape) # [1, 24, 5, 24, 40]