Blame for code/model.py

# 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]