From 57500fc5bc92566a63f2046824f522cd55c335ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Mon, 3 Aug 2026 05:28:29 +0300 Subject: [PATCH] feat: Support MiniMax-H3 (CORE-375) (#15224) --- comfy/latent_formats.py | 39 ++ comfy/ldm/minimax/audio_vae.py | 442 ++++++++++++++++++++ comfy/ldm/minimax/model.py | 646 ++++++++++++++++++++++++++++ comfy/ldm/minimax/vae.py | 694 +++++++++++++++++++++++++++++++ comfy/model_base.py | 52 +++ comfy/model_detection.py | 29 ++ comfy/ops.py | 55 ++- comfy/quant_ops.py | 2 +- comfy/sd.py | 55 +++ comfy/supported_models.py | 29 ++ comfy/text_encoders/llama.py | 11 + comfy/text_encoders/minimax.py | 201 +++++++++ comfy/text_encoders/qwen3vl.py | 5 +- comfy_extras/nodes_lt.py | 5 +- comfy_extras/nodes_minimax_h3.py | 337 +++++++++++++++ nodes.py | 3 +- 16 files changed, 2599 insertions(+), 6 deletions(-) create mode 100644 comfy/ldm/minimax/audio_vae.py create mode 100644 comfy/ldm/minimax/model.py create mode 100644 comfy/ldm/minimax/vae.py create mode 100644 comfy/text_encoders/minimax.py create mode 100644 comfy_extras/nodes_minimax_h3.py diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 63eefd420..c4270022b 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -567,6 +567,45 @@ class LTXAV(LTXV): ] self.latent_rgb_factors_bias = [-0.347892, -0.363814, -0.370287] +class MiniMaxH3Video(LatentFormat): + latent_channels = 24 + latent_dimensions = 3 + spacial_downscale_ratio = 16 + temporal_downscale_ratio = 4 + scale_factor = 1.0 + + latent_rgb_factors = [ + [-0.018555, 0.024344, -0.017536], + [ 0.150164, 0.137244, 0.129221], + [ 0.027367, -0.050369, -0.208606], + [-0.000793, -0.164622, -0.323161], + [-0.048556, 0.013970, -0.074286], + [ 0.011740, 0.014172, -0.006906], + [ 0.061517, 0.061212, 0.110025], + [ 0.035321, 0.086879, 0.110059], + [-0.017426, 0.002997, 0.035356], + [ 0.531539, 0.548819, 0.624404], + [-0.024968, -0.040234, -0.034302], + [-0.032549, -0.029096, -0.017221], + [ 0.022609, 0.020286, 0.050661], + [-0.084001, -0.038131, -0.020805], + [-0.018830, 0.010412, 0.061120], + [ 0.020777, 0.011196, -0.030994], + [-0.008390, -0.012201, -0.025687], + [-0.013281, -0.002924, 0.006331], + [ 0.000260, 0.001833, -0.011038], + [ 0.105471, 0.100482, 0.132106], + [ 0.016529, 0.015213, 0.009999], + [-0.014015, -0.017438, -0.019134], + [-0.033787, -0.009984, -0.019725], + [ 0.004224, 0.017284, 0.027196], + ] + latent_rgb_factors_bias = [ 0.057426, -0.022078, -0.071449] + +class MiniMaxH3AV(MiniMaxH3Video): + # max channels across the two streams (video 24, audio 32) so per-stream slices keep both streams whole + latent_channels = 32 + class HunyuanVideo(LatentFormat): latent_channels = 16 latent_dimensions = 3 diff --git a/comfy/ldm/minimax/audio_vae.py b/comfy/ldm/minimax/audio_vae.py new file mode 100644 index 000000000..033ae6966 --- /dev/null +++ b/comfy/ldm/minimax/audio_vae.py @@ -0,0 +1,442 @@ +# MiniMax H3 audio VAE: DAC-lineage waveform encoder + BigVGAN decoder. +# Weight-norm parametrizations are folded into plain conv weights, so this +# module uses ordinary ops.Conv1d / ops.ConvTranspose1d and loads the converted +# checkpoint (plain "*.weight" tensors) with strict=True. +# +# Lineage / licenses of the reference implementation: +# DAC encoder: descript-audio-codec (MIT) +# BigVGAN decoder: NVIDIA BigVGAN (MIT), adapted from hifi-gan (MIT) +# Alias-free ops: junjun3518/alias-free-torch (Apache-2.0), julius (MIT) + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import comfy.ops + +ops = comfy.ops.disable_weight_init + + +# Snake activations + +def snake(x, alpha, beta): + # x + 1/beta * sin^2(alpha * x) + t = torch.sin(alpha * x) + return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x) + + +class Snake1d(nn.Module): + """Snake activation with per-channel alpha (encoder side).""" + + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.empty(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha, self.alpha) + + +class SnakeBeta(nn.Module): + """SnakeBeta := x + 1/beta * sin^2(alpha * x); alpha/beta stored in log scale.""" + + def __init__(self, in_features): + super().__init__() + self.alpha = nn.Parameter(torch.empty(in_features)) + self.beta = nn.Parameter(torch.empty(in_features)) + + def forward(self, x): + alpha = torch.exp(self.alpha).view(1, -1, 1) + beta = torch.exp(self.beta).view(1, -1, 1) + return snake(x, alpha, beta) + + +# Alias-free (anti-aliased) activation: kaiser-windowed sinc resampling + +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): + # returns filter [1, 1, kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # kaiser window design + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + + filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + # Normalize filter to have sum = 1, otherwise there is a small leakage of + # the constant component in the input signal. + filter_ /= filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=12): + super().__init__() + self.ratio = ratio + self.stride = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2 + self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2 + self.register_buffer( + "filter", + kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size), + ) + + def forward(self, x): + _, C, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = F.conv_transpose1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C).mul_(self.ratio) + x = x[..., self.pad_left:-self.pad_right] + return x + + +class LowPassFilter1d(nn.Module): + def __init__(self, cutoff=0.5, half_width=0.6, stride=1, kernel_size=12): + super().__init__() + self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + + def forward(self, x): + _, C, _ = x.shape + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C) + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=12): + super().__init__() + self.ratio = ratio + self.kernel_size = kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x): + return self.lowpass(x) + + +class Activation1d(nn.Module): + """upsample x2 -> pointwise activation -> downsample x2 (anti-aliased).""" + + def __init__(self, activation, up_ratio=2, down_ratio=2, up_kernel_size=12, down_kernel_size=12): + super().__init__() + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + return x + + +# DAC encoder + +class ResidualUnit(nn.Module): + def __init__(self, dim=16, dilation=1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Snake1d(dim), + ops.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + Snake1d(dim), + ops.Conv1d(dim, dim, kernel_size=1), + ) + + def forward(self, x): + y = self.block(x) + pad = (x.shape[-1] - y.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return y.add_(x) + + +class EncoderBlock(nn.Module): + def __init__(self, dim=16, stride=1): + super().__init__() + self.block = nn.Sequential( + ResidualUnit(dim // 2, dilation=1), + ResidualUnit(dim // 2, dilation=3), + ResidualUnit(dim // 2, dilation=9), + Snake1d(dim // 2), + ops.Conv1d( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, x): + return self.block(x) + + +class Encoder(nn.Module): + def __init__(self, d_model=64, strides=(2, 4, 4, 5, 5), d_latent=2048): + super().__init__() + block = [ops.Conv1d(1, d_model, kernel_size=7, padding=3)] + for stride in strides: + d_model *= 2 + block += [EncoderBlock(d_model, stride=stride)] + block += [ + Snake1d(d_model), + ops.Conv1d(d_model, d_latent, kernel_size=3, padding=1), + ] + self.block = nn.Sequential(*block) + + def forward(self, x): + return self.block(x) + + +# Attention projection (encoder posterior head) + +class GeGluMlp(nn.Module): + def __init__(self, in_features, hidden_features): + super().__init__() + self.norm = ops.LayerNorm(in_features) + self.act = nn.GELU(approximate="tanh") + self.w0 = ops.Linear(in_features, hidden_features) + self.w1 = ops.Linear(in_features, hidden_features) + self.w2 = ops.Linear(hidden_features, in_features) + + def forward(self, x): + x = self.norm(x) + return self.w2(self.act(self.w0(x)).mul_(self.w1(x))) + + +class CausalAttention(nn.Module): + def __init__(self, in_dim, out_dim, num_heads): + super().__init__() + self.head_dim = in_dim // num_heads + self.num_heads = num_heads + self.out_dim = out_dim + self.qkv = ops.Linear(in_dim, in_dim * 3, bias=False) + self.q_bias = nn.Parameter(torch.empty(in_dim)) + self.v_bias = nn.Parameter(torch.empty(in_dim)) + self.register_buffer("zero_k_bias", torch.empty(in_dim)) + self.proj = ops.Linear(out_dim, out_dim) + + def forward(self, x): + B, N, C = x.shape + weight, _, offload_stream = comfy.ops.cast_bias_weight(self.qkv, x, offloadable=True) + qkv = F.linear(x, weight=weight, bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias))) + comfy.ops.uncast_bias_weight(self.qkv, weight, None, offload_stream) + q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0) + + # mean over heads then pool down to the latent width (in_dim >> out_dim) + x = comfy.ops.scaled_dot_product_attention(q, k, v, is_causal=True) + x = F.adaptive_avg_pool1d(torch.mean(x, dim=1), self.out_dim) + return self.proj(x) + + +class AttnProjection(nn.Module): + def __init__(self, in_dim, out_dim, num_heads, mlp_ratio=2): + super().__init__() + self.norm1 = ops.LayerNorm(in_dim) + self.attn = CausalAttention(in_dim, out_dim, num_heads) + self.proj = ops.Linear(in_dim, out_dim) + self.norm3 = ops.LayerNorm(in_dim) + + self.norm2 = ops.LayerNorm(out_dim) + hidden_dim = int(out_dim * mlp_ratio) + self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim) + + def forward(self, x): + # x: [B, T, in_dim] + x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x))) + return x.add_(self.mlp(self.norm2(x))) + + +# BigVGAN decoder + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +class AMPBlock1(nn.Module): + def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)): + super().__init__() + self.convs1 = nn.ModuleList( + [ + ops.Conv1d(channels, channels, kernel_size, stride=1, dilation=d, padding=get_padding(kernel_size, d)) + for d in dilation + ] + ) + self.convs2 = nn.ModuleList( + [ + ops.Conv1d(channels, channels, kernel_size, stride=1, dilation=1, padding=get_padding(kernel_size, 1)) + for _ in range(len(dilation)) + ] + ) + self.num_layers = len(self.convs1) + len(self.convs2) + self.activations = nn.ModuleList( + [Activation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)] + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt.add_(x) + return x + + +class BigVGAN(nn.Module): + """BigVGAN vocoder (MiniMax H3 32 kHz configuration). + + use_bias_at_final=False, use_tanh_at_final=False (output clamped to [-1, 1]). + """ + + def __init__( + self, + num_mels=2048, + upsample_initial_channel=1024, + upsample_rates=(5, 5, 2, 2, 2, 2, 2), + upsample_kernel_sizes=(9, 9, 4, 4, 4, 4, 4), + resblock_kernel_sizes=(3, 7, 11), + resblock_dilation_sizes=((1, 3, 5), (1, 3, 5), (1, 3, 5)), + ): + super().__init__() + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + + self.conv_pre = ops.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3) + + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + nn.ModuleList( + [ + ops.ConvTranspose1d( + upsample_initial_channel // (2 ** i), + upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ] + ) + ) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = upsample_initial_channel // (2 ** (i + 1)) + for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes): + self.resblocks.append(AMPBlock1(ch, k, d)) + + self.activation_post = Activation1d(activation=SnakeBeta(ch)) + self.conv_post = ops.Conv1d(ch, 1, 7, 1, padding=3, bias=False) + + def forward(self, x): + x = self.conv_pre(x) + + for i in range(self.num_upsamples): + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs.div_(self.num_kernels) + + x = self.activation_post(x) + return self.conv_post(x).clamp_(-1.0, 1.0) + + +# Top-level VAE + +class MiniMaxH3AudioVAE(nn.Module): + """MiniMax H3 stereo audio VAE at 32 kHz. + + Latents are [B, 32, 2, T]: 32 channels, 2 stereo channels, T frames at + 40 latent frames per second (800 audio samples per latent frame). The + stereo channels are processed independently by the mono encoder/decoder. + Latents are normalized with the stored per-channel latents_mean/std. + """ + + def __init__( + self, + encoder_dim=64, + encoder_rates=(2, 4, 4, 5, 5), + latent_dim=2048, + decoder_dim=1024, + vae_latent_channels=32, + ): + super().__init__() + self.sample_rate = 32000 + + self.hop_length = 1 + for r in encoder_rates: + self.hop_length *= r + self.samples_per_latent = self.hop_length # 800 + self.latents_per_second = self.sample_rate // self.hop_length # 40 + self.output_sample_rate = self.sample_rate # read by LTXVAudioVAEDecode + + self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim) + + self.pre_block = AttnProjection(latent_dim, vae_latent_channels, num_heads=8) + + self.mean_proj = ops.Conv1d(vae_latent_channels, vae_latent_channels, 1) + # logs_proj exists in the checkpoint but is unused at inference + # (encode returns the posterior mean, no sampling). + self.logs_proj = ops.Conv1d(vae_latent_channels, vae_latent_channels, 1) + + self.dec_in_proj = ops.Conv1d(vae_latent_channels, latent_dim, 1) + self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim) + + self.register_buffer("latents_mean", torch.empty(vae_latent_channels)) + self.register_buffer("latents_std", torch.empty(vae_latent_channels)) + + def decode(self, z): + """Decode normalized latents [B, 32, 2, T] to stereo waveforms [B, 2, L] at 32 kHz.""" + b, c, s, t = z.shape + z = z.permute(0, 2, 1, 3).reshape(b * s, c, t) + mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype) + std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype) + z = z * std + mean + x = self.dec_in_proj(z) + x = self.decoder(x) # [b * s, 1, L], already clamped to [-1, 1] + return x.reshape(b, s, -1) + + def encode(self, waveform): + """Encode stereo waveforms [B, 2, L] at 32 kHz (in [-1, 1]) to normalized latents [B, 32, 2, T]. + + L is right-padded with zeros to a multiple of 800 samples; the returned + posterior mean is used directly (no sampling). + """ + b, s, length = waveform.shape + right_pad = math.ceil(length / self.hop_length) * self.hop_length - length + waveform = F.pad(waveform, (0, right_pad)) + x = waveform.reshape(b * s, 1, -1) + x = self.encoder(x) # [b * s, latent_dim, T] + x = self.pre_block(x.transpose(1, 2)).transpose(1, 2) # [b * s, 32, T] + z = self.mean_proj(x) + mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype) + std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype) + z = (z - mean) / std + return z.reshape(b, s, z.shape[1], z.shape[2]).permute(0, 2, 1, 3) diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py new file mode 100644 index 000000000..494350d40 --- /dev/null +++ b/comfy/ldm/minimax/model.py @@ -0,0 +1,646 @@ +"""MiniMax H3 audio-video DiT. + +Single-stream packed-token transformer denoising video (24ch, patch 1x2x2) and +stereo audio (32ch, 40 Hz) latents jointly, conditioned on Qwen3-VL layer-50 hidden states. +The packed sequence is: +[text | cond rows | audio | video] for t2va/fl2va +[text | reference blocks | audio | video] for ref2va + +Timestep domain: the model receives the *video* sigma from the sampler and +derives per-token timesteps t = 1 - sigma internally; the audio stream runs on +its own shifted schedule (sigma_shift video 12.0 / audio 3.0), mapped from the +video sigma in closed form. The audio velocity is returned scaled by the +schedule map's derivative d(sigma_a)/d(sigma_v). +""" + +import math + +import torch +import torch.nn as nn + +import comfy.ldm.common_dit +import comfy.model_management +import comfy.model_prefetch +import comfy.ops +import comfy.patcher_extension +import comfy.quant_ops +from comfy.ldm.modules.attention import optimized_attention + +FRAME_PER_TOKEN = (1, 4, 4, 4, 4) +FRAME_RESCALE = 5.0 / 3.0 +VISUAL_COND_TIMESTEP = 0.999 +AUDIO_COND_TIMESTEP = 1.0 + + +def time_shift_sigma(sigma, from_shift, to_shift): + # invert sigma = s*b/(1+(s-1)*b) to the base grid, re-apply the other shift + base = sigma / (from_shift + sigma * (1.0 - from_shift)) + 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. + + Scaling a stream's returned velocity by this slope makes the flat ODE that + any sampler integrates on the from-schedule equal to that stream's true ODE + on its own schedule. + """ + base = sigma / (from_shift + sigma * (1.0 - from_shift)) + return (to_shift * (1.0 + (from_shift - 1.0) * base) ** 2) / (from_shift * (1.0 + (to_shift - 1.0) * base) ** 2) + + +def patchify_video(latent, patch_size=(1, 2, 2)): + # [B, C, T, H, W] -> [B*t*h*w, C*pt*ph*pw] + b, c, t_full, h_full, w_full = latent.shape + pt, ph, pw = patch_size + t, h, w = t_full // pt, h_full // ph, w_full // pw + x = latent.reshape(b, c, t, pt, h, ph, w, pw) + x = torch.einsum("nctrhpwq->nthwcrpq", x) + return x.reshape(b * t * h * w, c * pt * ph * pw) + + +def unpatchify_video(rows, t, h, w, c=24, patch_size=(1, 2, 2)): + pt, ph, pw = patch_size + x = rows.reshape(-1, t, h, w, c, pt, ph, pw) + x = torch.einsum("nthwcrpq->nctrhpwq", x) + return x.reshape(-1, c, t * pt, h * ph, w * pw) + + +def pack_audio(latent): + # [B, C=32, ch=2, T] -> [ch*T, 32] channel-major (ch0 t0..T-1, ch1 t0..T-1) + b, c, ch, t = latent.shape + return latent[0].permute(1, 2, 0).reshape(ch * t, c) + + +def unpack_audio(rows, ch=2): + t = rows.shape[0] // ch + return rows.reshape(ch, t, rows.shape[-1]).permute(2, 0, 1).unsqueeze(0) + + +def _axis_from_sqrt_area(dim, patch, sqrt_area): + # linspace((1 - ratio) / 2, (1 + ratio) / 2, dim // patch, endpoint=False) * 32 + ratio = dim / sqrt_area + n = dim // patch + return (torch.arange(n, dtype=torch.float64) * (ratio / n) + (1.0 - ratio) / 2.0) * 32.0 + + +def _frame_grid(h, w): + # area-normalized (h, w) coordinates of one latent frame's 2x2-patch rows + area = math.sqrt(h * w) + hh, ww = torch.meshgrid(_axis_from_sqrt_area(h, 2, area), _axis_from_sqrt_area(w, 2, area), indexing="ij") + return torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1), _axis_from_sqrt_area(w, 2, area) + + +def _video_t_spans(n): + return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(n)] + + +def _video_t_grid(n, origin): + # origin + exclusive cumsum + spans = torch.tensor(_video_t_spans(n), dtype=torch.float64) + return float(origin) + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)]) + + +def _audio_grid(cursor, t, w_low, w_high): + # channel-major stereo rows: t advances per latent frame, w pinned to the grid extremes per stereo channel, h stays 0 + g = torch.zeros(t * 2, 3, dtype=torch.float64) + g[:, 0] = (cursor + torch.arange(t, dtype=torch.float64)).repeat(2) + g[:t, 2] = w_low + g[t:, 2] = w_high + return g + + +def _video_grid(vt, frame, cursor): + g = torch.empty(vt, frame.shape[0], 3, dtype=torch.float64) + g[:, :, 0] = _video_t_grid(vt, cursor)[:, None] + g[:, :, 1:] = frame[None] + return g.reshape(-1, 3) + + +class TimeEmbedder(nn.Module): + def __init__(self, freq_dim, hidden, out, dtype=None, device=None, operations=None): + super().__init__() + self.freq_dim = freq_dim + self.proj_in = operations.Linear(freq_dim, hidden, bias=True, dtype=dtype, device=device) + self.proj_out = operations.Linear(hidden, out, bias=True, dtype=dtype, device=device) + + def forward(self, t): + # t: [M] in [0, 1]; fp32 throughout, cos before sin + half = self.freq_dim // 2 + freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=t.device) / half) + args = t.to(torch.float32)[:, None] * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.proj_out(nn.functional.silu(self.proj_in(emb))) + + +def rope_rotation_table(angles, dtype): + """[S, rot_dim] pair angles -> [1, S, 1, rot_dim/2, 2, 2] rotation matrices.""" + half = angles.shape[-1] // 2 + ang = angles[:, :half] # duplicated halves: [:, :half] == [:, half:] + c, s = torch.cos(ang), torch.sin(ang) + table = torch.stack([c, -s, s, c], dim=-1).reshape(1, angles.shape[0], 1, half, 2, 2) + return table.to(dtype) + + +class Attention(nn.Module): + def __init__(self, hidden, heads, head_dim, eps, dtype=None, device=None, operations=None): + super().__init__() + self.heads = heads + self.head_dim = head_dim + inner = heads * head_dim + self.qkv_proj = operations.Linear(hidden, inner * 3, bias=False, dtype=dtype, device=device) + self.q_norm = operations.RMSNorm(head_dim, eps=eps, dtype=dtype, device=device) + self.k_norm = operations.RMSNorm(head_dim, eps=eps, dtype=dtype, device=device) + self.out_proj = operations.Linear(inner, hidden, bias=False, dtype=dtype, device=device) + + def forward(self, x, rope_freqs=None, transformer_options={}): + s = x.shape[0] + q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1) + v = v.view(s, self.heads, self.head_dim) + if rope_freqs is not None: + # fused per-head RMSNorm + partial split-half rope, in place on the qkv buffer + q = q.view(1, s, self.heads, self.head_dim) + k = k.view(1, s, self.heads, self.head_dim) + qw = comfy.model_management.cast_to(self.q_norm.weight, device=x.device) + kw = comfy.model_management.cast_to(self.k_norm.weight, device=x.device) + rot = rope_freqs.shape[-3] * 2 + if comfy.model_management.in_training: + q, k = comfy.quant_ops.ck.rms_rope_split_half( + q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot) + else: + comfy.quant_ops.ck.rms_rope_split_half_( + q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot) + q = q[0] + k = k[0] + else: + q = self.q_norm(q.view(s, self.heads, self.head_dim)) + k = self.k_norm(k.view(s, self.heads, self.head_dim)) + q = q.transpose(0, 1).unsqueeze(0) + k = k.transpose(0, 1).unsqueeze(0) + v = v.transpose(0, 1).unsqueeze(0) + out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options) + return self.out_proj(out.squeeze(0)) + + +class MLP(nn.Module): + def __init__(self, hidden, ffn, dtype=None, device=None, operations=None): + super().__init__() + self.fc1 = operations.Linear(hidden, ffn * 2, bias=False, dtype=dtype, device=device) + self.fc2 = operations.Linear(ffn, hidden, bias=False, dtype=dtype, device=device) + + def forward(self, x): + return comfy.ops.linear_input_act(self.fc2, self.fc1(x), "swiglu") + + +class AdalnProj(nn.Module): + def __init__(self, t_dim, hidden, expand, modalities, apply_silu=True, + dtype=None, device=None, operations=None): + super().__init__() + self.expand = expand + self.modalities = modalities + self.hidden = hidden + self.apply_silu = apply_silu + self.linear = operations.Linear(t_dim, expand * hidden * modalities, bias=True, dtype=dtype, device=device) + + def forward(self, t_emb): + # [M, t_dim] -> expand tensors of [M*modalities, hidden] + x = self.linear(nn.functional.silu(t_emb) if self.apply_silu else t_emb) + x = x.view(x.shape[0] * self.modalities, self.expand * self.hidden) + return x.chunk(self.expand, dim=-1) + + +def _mod_scale_shift(h, shift, scale, segments): + # segments: [(start, stop, mod_row)] covering h contiguously. + for a, b, row in segments: + h[a:b].mul_(1.0 + scale[row].to(h.dtype)).add_(shift[row].to(h.dtype)) + return h + + +def _mod_gate(x, gate, other, segments): + # other is the fresh attn/mlp output: accumulate the gated residual into the stream in place, one fused kernel per segment + for a, b, row in segments: + x[a:b].addcmul_(other[a:b], gate[row].to(x.dtype)) + return x + + +class RefinerBlock(nn.Module): + def __init__(self, hidden, heads, head_dim, ffn, eps, qk_eps, dtype=None, device=None, operations=None): + super().__init__() + self.norm1 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device) + self.norm2 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device) + self.attn = Attention(hidden, heads, head_dim, qk_eps, dtype=dtype, device=device, operations=operations) + self.mlp = MLP(hidden, ffn, dtype=dtype, device=device, operations=operations) + + def forward(self, x, transformer_options={}): + # attn/mlp outputs are fresh: accumulate residuals in place + x = self.attn(self.norm1(x), transformer_options=transformer_options).add_(x) + return self.mlp(self.norm2(x)).add_(x) + + +class TokenRefiner(nn.Module): + def __init__(self, num_layers, hidden, heads, head_dim, ffn, eps, qk_eps, final_eps, + dtype=None, device=None, operations=None): + super().__init__() + self.blocks = nn.ModuleList([ + RefinerBlock(hidden, heads, head_dim, ffn, eps, qk_eps, dtype=dtype, device=device, operations=operations) + for _ in range(num_layers)]) + self.final_norm = operations.RMSNorm(hidden, eps=final_eps, dtype=dtype, device=device) + + def forward(self, x, transformer_options={}): + for block in self.blocks: + x = block(x, transformer_options=transformer_options) + return self.final_norm(x) + + +class DiTBlock(nn.Module): + def __init__(self, hidden, heads, head_dim, ffn, t_dim, eps, qk_eps, + apply_silu=True, adaln_dtype=None, dtype=None, device=None, operations=None): + super().__init__() + self.norm1 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device) + self.norm2 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device) + self.attn = Attention(hidden, heads, head_dim, qk_eps, dtype=dtype, device=device, operations=operations) + self.mlp = MLP(hidden, ffn, dtype=dtype, device=device, operations=operations) + self.adaln_proj = AdalnProj(t_dim, hidden, 6, 3, apply_silu=apply_silu, + dtype=adaln_dtype if adaln_dtype is not None else dtype, + device=device, operations=operations) + + def forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={}): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb) + h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments) + x = _mod_gate(x, gate_msa, self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments) + h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments) + return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments) + + +class FinalLayer(nn.Module): + def __init__(self, hidden, t_dim, video_dim, audio_dim, eps, apply_silu=True, adaln_dtype=None, + dtype=None, device=None, operations=None): + super().__init__() + self.norm = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device) + self.adaln_proj = AdalnProj(t_dim, hidden, 2, 1, apply_silu=apply_silu, + dtype=adaln_dtype if adaln_dtype is not None else dtype, + device=device, operations=operations) + # output heads are the checkpoint's fp32 island; norm/adaln are stored at model dtype + self.video_out = operations.Linear(hidden, video_dim, bias=True, dtype=torch.float32, device=device) + self.audio_out = operations.Linear(hidden, audio_dim, bias=True, dtype=torch.float32, device=device) + + def forward(self, x, t_emb, video_seg, audio_seg): + # video_seg / audio_seg: (start, stop, timestep_row) of the target streams + shift, scale = self.adaln_proj(t_emb) + va, vb, vrow = video_seg + aa, ab, arow = audio_seg + hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32) + ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32) + return self.video_out(hv), self.audio_out(ha) + + +class PackedLayout: + """Static packed-sequence structure for one shape/conditioning signature.""" + + def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None, frame_count=None): + frame, w_grid = _frame_grid(latent_h, latent_w) + frame_rows = frame.shape[0] + + segments = [("text", text_len)] # (kind, n_rows) + g = torch.zeros(text_len, 3, dtype=torch.float64) + g[:, 0] = torch.arange(text_len, dtype=torch.float64) + pos = [g] # per segment: [n, 3] float64 (t, h, w) + + img_pos, img_update = [], [] + audio_pos, audio_update = [], [] + cursor = text_len + row = text_len + + if keyframes: + # fl2va: keyframe cond rows right after text, sharing the target spatial grid + for kf in keyframes: + pixel_index = kf["resolved_frame_index"] + if pixel_index == 0: + cond_t = float(text_len) + elif frame_count is not None and pixel_index == frame_count - 1: + cond_t = float(text_len) + sum(_video_t_spans(latent_t)) - FRAME_RESCALE + else: + raise ValueError("only first/last keyframe anchors are supported") + g = torch.empty(frame_rows, 3, dtype=torch.float64) + g[:, 0] = cond_t + g[:, 1:] = frame + segments.append(("cond", frame_rows)) + pos.append(g) + img_pos.append(torch.arange(row, row + frame_rows)) + img_update.append(torch.zeros(frame_rows, dtype=torch.bool)) + row += frame_rows + + target_audio_w = (float(w_grid[0]), float(w_grid[-1])) + if refs: + cursor = float(text_len) + for blk in refs: + kind = blk["kind"] + if kind == "image": + r_frame, _ = _frame_grid(blk["latent_h"], blk["latent_w"]) + n = r_frame.shape[0] + g = torch.empty(n, 3, dtype=torch.float64) + g[:, 0] = cursor + g[:, 1:] = r_frame + segments.append(("ref_img", n)) + pos.append(g) + img_pos.append(torch.arange(row, row + n)) + img_update.append(torch.zeros(n, dtype=torch.bool)) + row += n + cursor += 1.0 + elif kind == "audio": + rt = blk["ref_audio_t"] + if rt > 0: + segments.append(("ref_audio", rt * 2)) + pos.append(_audio_grid(cursor, rt, *target_audio_w)) + audio_pos.append(torch.arange(row, row + rt * 2)) + audio_update.append(torch.zeros(rt * 2, dtype=torch.bool)) + row += rt * 2 + cursor += float(rt) + elif kind in ("video", "video_audio"): + # the block's audio rows pack immediately before its video + # rows, both sharing the cursor origin + rt = blk["ref_audio_t"] + vt = blk["latent_t"] + r_frame, r_w_grid = _frame_grid(blk["latent_h"], blk["latent_w"]) + if rt > 0: + segments.append(("ref_audio", rt * 2)) + pos.append(_audio_grid(cursor, rt, float(r_w_grid[0]), float(r_w_grid[-1]))) + audio_pos.append(torch.arange(row, row + rt * 2)) + audio_update.append(torch.zeros(rt * 2, dtype=torch.bool)) + row += rt * 2 + n = vt * r_frame.shape[0] + segments.append(("ref_img", n)) + pos.append(_video_grid(vt, r_frame, cursor)) + img_pos.append(torch.arange(row, row + n)) + img_update.append(torch.zeros(n, dtype=torch.bool)) + row += n + cursor += max(float(rt), sum(_video_t_spans(vt))) + + # target audio then target video, always the last two segments + segments.append(("audio", audio_t * 2)) + pos.append(_audio_grid(cursor, audio_t, *target_audio_w)) + audio_pos.append(torch.arange(row, row + audio_t * 2)) + audio_update.append(torch.ones(audio_t * 2, dtype=torch.bool)) + row += audio_t * 2 + + n_video = latent_t * frame_rows + segments.append(("video", n_video)) + pos.append(_video_grid(latent_t, frame, cursor)) + img_pos.append(torch.arange(row, row + n_video)) + img_update.append(torch.ones(n_video, dtype=torch.bool)) + row += n_video + + self.seq_len = row + self.position_ids = torch.cat(pos) # [S, 3] float64 + self.img_pos = torch.cat(img_pos) + self.img_update = torch.cat(img_update) + self.audio_pos = torch.cat(audio_pos) + self.audio_update = torch.cat(audio_update) + self.signature = (text_len, latent_t, latent_h, latent_w, audio_t) + # contiguous segment table (start, stop, kind) + # kinds: text / cond / ref_img / ref_audio / audio / video + # the packed sequence is uniform per segment in (modality tag, timestep class), + # except the text span (tag runs resolved at forward time from the presentation tags) + seg_abs = [] + off = 0 + for kind, n in segments: + seg_abs.append((off, off + n, kind)) + off += n + self.segments = seg_abs + + +class MiniMaxH3Model(nn.Module): + def __init__(self, hidden_size=5376, num_layers=50, token_refiner_num_layers=2, + num_attention_heads=56, attention_head_dim=128, ffn_hidden_size=14336, + latents_dim=24, audio_latents_dim=32, patch_size=(1, 2, 2), text_dim=5120, + timestep_input_dim=256, time_embed_hidden_size=5376, time_embed_dim=2688, + rope_inv_freq_len=16, norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5, + sigma_shift_video=12.0, sigma_shift_audio=3.0, + adaln_curve_grid=None, + image_model=None, dtype=None, device=None, operations=None, **kwargs): + super().__init__() + self.dtype = dtype + self.hidden_size = hidden_size + self.patch_size = tuple(patch_size) + self.latents_dim = latents_dim + self.audio_latents_dim = audio_latents_dim + self.sigma_shift_video = sigma_shift_video + self.sigma_shift_audio = sigma_shift_audio + self.use_adaln_curves = adaln_curve_grid is not None + # curve-form checkpoints replace the time embedder and full-width adaln weights with a small shared basis of the time-embedding curve + curve = {"apply_silu": not self.use_adaln_curves, + "adaln_dtype": torch.float32 if self.use_adaln_curves else dtype} + video_patch_dim = latents_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2] + + self.video_patch_proj = operations.Linear(video_patch_dim, hidden_size, bias=True, dtype=torch.float32, device=device) + self.audio_patch_proj = operations.Linear(audio_latents_dim, hidden_size, bias=True, dtype=torch.float32, device=device) + self.condition_proj = operations.Linear(text_dim, hidden_size, bias=True, dtype=dtype, device=device) + if self.use_adaln_curves: + self.register_buffer("adaln_t_table", torch.empty(adaln_curve_grid, time_embed_dim, dtype=torch.float32)) + else: + self.time_embedder = TimeEmbedder(timestep_input_dim, time_embed_hidden_size, time_embed_dim, + dtype=torch.float32, device=device, operations=operations) + self.rope = nn.Module() + self.rope.register_buffer("inv_freq", torch.empty(rope_inv_freq_len, dtype=torch.float32)) + self.token_refiner = TokenRefiner(token_refiner_num_layers, hidden_size, num_attention_heads, + attention_head_dim, ffn_hidden_size, norm_eps, qk_norm_eps, + final_norm_eps, dtype=dtype, device=device, operations=operations) + self.blocks = nn.ModuleList([ + DiTBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, + time_embed_dim, norm_eps, qk_norm_eps, **curve, dtype=dtype, device=device, operations=operations) + for _ in range(num_layers)]) + self.final_layer = FinalLayer(hidden_size, time_embed_dim, video_patch_dim, audio_latents_dim, + final_norm_eps, **curve, dtype=dtype, device=device, operations=operations) + + def preprocess_text_embeds(self, text_states): + """[B, L, text_dim] Qwen states -> [B, L, hidden] refined text embeds.""" + if text_states.shape[-1] == self.hidden_size: + return text_states + return self.token_refiner(self.condition_proj(text_states[0])).unsqueeze(0) + + def rope_freqs(self, position_ids, device): + # [S, 3] float64 -> [S, 96] fp32 + pos = position_ids.to(torch.float32).to(device) + inv = comfy.model_management.cast_to(self.rope.inv_freq, device=device) + per_axis = pos.unsqueeze(-1) * inv.view(1, 1, -1) # [S, 3, 16] + t_f, h_f, w_f = per_axis.unbind(dim=1) + half = torch.cat((t_f, h_f, w_f), dim=-1) # [S, 48] + return torch.cat((half, half), dim=-1) # [S, 96] + + def _cond_video_rows(self, payload, device): + """Concatenated visual condition rows (normalized latents -> patchified), with condition noise augmentation.""" + rows = [] + aug = payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP) + seed = int(payload.get("seed", 0)) + # every condition intentionally restarts the same RNG stream + for z in payload.get("cond_video_latents", []): + r = patchify_video(z.to(torch.float32), self.patch_size) + if aug < 1.0: + gen = torch.Generator("cpu").manual_seed(seed) + noise = torch.randn(r.shape, generator=gen, dtype=torch.float32) + r = aug * r + (1.0 - aug) * noise.to(r.device) + rows.append(r.to(device)) + return torch.cat(rows, dim=0) if rows else None + + def _cond_audio_rows(self, payload, device): + rows = [] + aug = payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP) + seed = int(payload.get("seed", 0)) + 1 + for z in payload.get("cond_audio_latents", []): + r = pack_audio(z.to(torch.float32)) + if aug < 1.0: + gen = torch.Generator("cpu").manual_seed(seed) + noise = torch.randn(r.shape, generator=gen, dtype=torch.float32) + r = aug * r + (1.0 - aug) * noise.to(r.device) + rows.append(r.to(device)) + return torch.cat(rows, dim=0) if rows else None + + def forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs): + return comfy.patcher_extension.WrapperExecutor.new_class_executor( + self._forward, + self, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options) + ).execute(x, timestep, context, transformer_options, minimax_payload=minimax_payload, **kwargs) + + def _forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs): + video_x, audio_x = x[0], x[1] + orig_t, orig_h, orig_w = video_x.shape[2], video_x.shape[3], video_x.shape[4] + video_x = comfy.ldm.common_dit.pad_to_patch_size(video_x, self.patch_size) + if video_x.shape[0] != 1: + raise ValueError("MiniMax H3 supports batch size 1") + payload = minimax_payload or {} + device = video_x.device + dtype = context.dtype # compute dtype + + latent_t, lat_h, lat_w = video_x.shape[2], video_x.shape[3], video_x.shape[4] + audio_t = audio_x.shape[-1] + text_len = context.shape[1] + # extra_conds prebuilds the layout once per sampling run + layout = payload.get("layout") + if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t): + layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t, + keyframes=payload.get("keyframes"), + refs=payload.get("refs"), + frame_count=payload.get("frame_count")) + + # model_base passes model_sampling.timestep(sigma) = sigma * 1000 + shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video)) + shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio)) + sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6) + t_v = float(1.0 - sigma_v) + t_a = float(1.0 - time_shift_sigma(sigma_v, shift_v, shift_a)) + + # distinct timesteps are known analytically: text/pad follow video, cond rows pin near 1 + vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP)) + aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP)) + has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments) + has_aud_cond = any(k == "ref_audio" for _, _, k in layout.segments) + seg_t = {"text": t_v, "video": t_v, "audio": t_a, + "cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug), + "ref_audio": max(t_a, aud_aug)} + unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set()) + | ({seg_t["ref_audio"]} if has_aud_cond else set())) + t_row = {t: i for i, t in enumerate(unique_t)} + seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2} + + text_tags = payload.get("text_token_tags") + mod_segments = [] + for a, b, kind in layout.segments: + row_base = t_row[seg_t[kind]] * 3 + if kind == "text" and text_tags is not None: + # the presentation text span mixes tags (vision pads carry the video modality) split into tag runs + tags = text_tags.view(-1).tolist() + run_start = 0 + for i in range(1, b - a + 1): + if i == b - a or tags[i] != tags[run_start]: + mod_segments.append((a + run_start, a + i, row_base + int(tags[run_start]))) + run_start = i + else: + mod_segments.append((a, b, row_base + seg_tag[kind])) + + # embed + img_update = layout.img_update.to(device) + audio_update = layout.audio_update.to(device) + video_rows = patchify_video(video_x.to(torch.float32), self.patch_size) + audio_rows = pack_audio(audio_x.to(torch.float32)) + cond_video_rows = self._cond_video_rows(payload, device) + cond_audio_rows = self._cond_audio_rows(payload, device) + + all_video_rows = video_rows + if cond_video_rows is not None: + all_video_rows = torch.empty(img_update.shape[0], video_rows.shape[1], dtype=torch.float32, device=device) + all_video_rows[~img_update] = cond_video_rows + all_video_rows[img_update] = video_rows + all_audio_rows = audio_rows + if cond_audio_rows is not None: + all_audio_rows = torch.empty(audio_update.shape[0], audio_rows.shape[1], dtype=torch.float32, device=device) + all_audio_rows[~audio_update] = cond_audio_rows + all_audio_rows[audio_update] = audio_rows + + video_embed = self.video_patch_proj(all_video_rows).to(dtype) + audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype) + text_states = context[0] + if text_states.shape[-1] != self.hidden_size: + text_states = self.token_refiner(self.condition_proj(text_states), + transformer_options=transformer_options) + + # segments are contiguous: assemble by slices, embed rows follow segment order + h = torch.empty(layout.seq_len, self.hidden_size, dtype=dtype, device=device) + voff = aoff = 0 + for a, b, kind in layout.segments: + n = b - a + if kind == "text": + h[a:b] = text_states + elif kind in ("cond", "ref_img", "video"): + h[a:b] = video_embed[voff:voff + n] + voff += n + else: # ref_audio / audio + h[a:b] = audio_embed[aoff:aoff + n] + aoff += n + + t_vals = torch.tensor(unique_t, dtype=torch.float32, device=device) + if self.use_adaln_curves: + # adaln projections consume interpolated coordinates of the time-embedding curve + table = comfy.model_management.cast_to(self.adaln_t_table, device=device) + pos = t_vals.clamp(0.0, 1.0) * (table.shape[0] - 1) # t in [0,1] -> fractional grid index, out-of-range t clamps to the curve ends + i0 = pos.floor().long().clamp(max=table.shape[0] - 2) # lower grid row, max-clamp keeps t=1.0 on the last interval instead of reading past the table + t_emb = torch.lerp(table[i0], table[i0 + 1], (pos - i0).unsqueeze(1)) # blend the two rows by the fractional part + else: + t_emb = self.time_embedder(t_vals).to(dtype) + + # rotation table computed once per forward, consumed by the kitchen split-half rope + rope_freqs = rope_rotation_table(self.rope_freqs(layout.position_ids, device), dtype) + + # blocks + patches_replace = transformer_options.get("patches_replace", {}) + blocks_replace = patches_replace.get("dit", {}) + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), device, transformer_options) + for i, block in enumerate(self.blocks): + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block) + if ("double_block", i) in blocks_replace: + def block_wrap(args): + return {"img": block(args["img"], args["t_emb"], args["mod_segments"], args["rope_freqs"], + transformer_options=args["transformer_options"])} + h = blocks_replace[("double_block", i)]( + {"img": h, "t_emb": t_emb, "mod_segments": mod_segments, "rope_freqs": rope_freqs, + "transformer_options": transformer_options}, + {"original_block": block_wrap})["img"] + else: + h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options) + if prefetch_queue is not None: + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None) + + # target streams are single contiguous segments (audio then video, last two) + video_seg = next((a, b, t_row[seg_t["video"]]) for a, b, k in layout.segments if k == "video") + audio_seg = next((a, b, t_row[seg_t["audio"]]) for a, b, k in layout.segments if k == "audio") + v, a = self.final_layer(h, t_emb, video_seg, audio_seg) + + video_out = unpatchify_video(v, latent_t, lat_h // 2, lat_w // 2, self.latents_dim, self.patch_size) + video_out = video_out[:, :, :orig_t, :orig_h, :orig_w] + audio_out = unpack_audio(a) + + # The sampler integrates the flat ODE dX/dsigma_v = (X - denoised)/sigma_v. + # Scaling the audio velocity by d(sigma_a)/d(sigma_v) makes that ODE equal + # to the audio stream's true ODE on its own shifted schedule. + slope_a = time_shift_slope(sigma_v, shift_v, shift_a).to(audio_out.dtype) + return [-video_out.to(video_x.dtype), (-slope_a) * audio_out.to(audio_x.dtype)] diff --git a/comfy/ldm/minimax/vae.py b/comfy/ldm/minimax/vae.py new file mode 100644 index 000000000..aeb3421a2 --- /dev/null +++ b/comfy/ldm/minimax/vae.py @@ -0,0 +1,694 @@ +# MiniMax H3 video VAE: 3D causal CNN encoder + ViT3D decoder. + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import comfy.ops +import comfy.quant_ops +import comfy.rmsnorm +from comfy.ldm.modules.attention import optimized_attention + +ops = comfy.ops.disable_weight_init + +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) + +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.68317747116088865, 1.7549455165863037, + 1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875, + 0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647, + 0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775, + 2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265, + 3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525, +] + + +# 3D causal CNN encoder + +class CausalConv3d(ops.Conv3d): + # Reflect spatial padding, causal (zeros, front-only) temporal padding. + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0): + super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride) + self.causal_padding = (padding,) * 3 if isinstance(padding, int) else tuple(padding) + + def forward(self, x): + if sum(self.causal_padding) == 0: + return super().forward(x) + + x = F.pad(x, (self.causal_padding[2], self.causal_padding[2], self.causal_padding[1], self.causal_padding[1], 0, 0), mode="reflect") + if x.shape[2] == 1: + # single frame: the causal front padding is all zeros truncate the temporal taps instead of convolving zero frames + return super().forward(x, autopad="causal_zero") + x = F.pad(x, (0, 0, 0, 0, self.causal_padding[0] * 2, 0), mode="constant") + return super().forward(x) + + +class TemporalIsolatedGroupNorm(ops.GroupNorm): + # GroupNorm with statistics computed per frame (time merged into batch). + def forward(self, x): + if x.dim() == 5: + b, c, t, h, w = x.shape + x = x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w) + x = super().forward(x) + return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous() + return super().forward(x) + + +def group_norm_3d(num_channels): + return TemporalIsolatedGroupNorm(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True) + + +class Downsample3D(nn.Module): + def __init__(self, in_channels, out_channels, time_stride=1, space_stride=2): + super().__init__() + self.space_stride = space_stride + self.conv = CausalConv3d( + in_channels, + out_channels, + kernel_size=3, + padding=(1, 0, 0), + stride=(time_stride, space_stride, space_stride), + ) + + def forward(self, x): + if self.space_stride == 2: + x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect") + return self.conv(x) + + +class ResnetBlock3D(nn.Module): + def __init__(self, in_channels, out_channels=None): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = group_norm_3d(in_channels) + self.norm2 = group_norm_3d(out_channels) + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1) + self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3, padding=1) + if in_channels != out_channels: + self.nin_shortcut = CausalConv3d(in_channels, out_channels, kernel_size=1) + + def forward(self, x): + h = self.conv1(F.silu(self.norm1(x), inplace=True)) + h = self.conv2(F.silu(self.norm2(h), inplace=True)) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return h.add_(x) + + +class EncoderFCN3D(nn.Module): + def __init__(self, ch, ch_mult, space_down, time_down, num_res_blocks, in_channels, z_channels, double_z=True): + super().__init__() + self.num_levels = len(ch_mult) + if isinstance(num_res_blocks, int): + num_res_blocks = [num_res_blocks] * self.num_levels + self.num_res_blocks = num_res_blocks + + block_mid = [ch * ch_mult[i] for i in range(self.num_levels)] + block_in = [block_mid[0]] + block_mid[:-1] + block_out = block_mid + + self.conv_in = CausalConv3d(in_channels, block_in[0], kernel_size=3, padding=1) + + self.down = nn.ModuleList() + for i_level in range(self.num_levels): + down = nn.Module() + down.block = nn.ModuleList() + for i in range(self.num_res_blocks[i_level]): + down.block.append( + ResnetBlock3D( + in_channels=block_in[i_level] if i == 0 else block_mid[i_level], + out_channels=block_mid[i_level], + ) + ) + if space_down[i_level] * time_down[i_level] > 1: + down.downsample = Downsample3D( + block_mid[i_level], + block_out[i_level], + time_stride=time_down[i_level], + space_stride=space_down[i_level], + ) + self.down.append(down) + + self.norm_out = group_norm_3d(block_out[-1]) + self.conv_out = CausalConv3d( + block_out[-1], + 2 * z_channels if double_z else z_channels, + kernel_size=3, + padding=1, + ) + + def forward(self, x): + h = self.conv_in(x) + for i_level in range(self.num_levels): + for i_block in range(self.num_res_blocks[i_level]): + h = self.down[i_level].block[i_block](h) + if hasattr(self.down[i_level], "downsample"): + h = self.down[i_level].downsample(h) + h = F.silu(self.norm_out(h)) + return self.conv_out(h) + + +# ViT3D decoder + +def create_token_ids(patch_dims, device, dtype): + coords_list = [] + for dim_size in patch_dims: + coords = torch.arange(0.5, dim_size, dtype=dtype, device=device) + coords = coords / dim_size + coords = 2.0 * coords - 1.0 + coords_list.append(coords) + coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1) + return coords.flatten(0, len(patch_dims) - 1).unsqueeze(0) + + +class RotaryEmbeddingND(nn.Module): + def __init__(self, dim, rotary_base=100.0, n_dim=3): + super().__init__() + self.n_dim = n_dim + self.angle_scale = 2.0 * math.pi + inv_freq = 1 / rotary_base ** torch.arange(0, 1, 2 * n_dim / dim, dtype=torch.float32) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, img_ids): + # [B, S, n_dim] -> [B, S, 1, pairs, 2, 2] rotation table for the kitchen split-half rope + angles = ( + self.angle_scale + * img_ids[:, :, :, None].float() + * self.inv_freq.to(img_ids.device)[None, None, None, :] + ) + angles = angles.flatten(2, 3) + c, s = torch.cos(angles), torch.sin(angles) + table = torch.stack([c, -s, s, c], dim=-1).reshape(*angles.shape[:2], 1, angles.shape[-1], 2, 2) + return table.to(img_ids.dtype) + + +class FeedForward(nn.Module): + # Gated SiLU FFN. + def __init__(self, dim, mult=4, bias=True): + super().__init__() + inner_dim = dim * mult + self.w1 = ops.Linear(dim, inner_dim * 2, bias=bias) + self.w2 = ops.Linear(inner_dim, dim, bias=bias) + + def forward(self, x): + gate, x = self.w1(x).chunk(2, dim=-1) + return self.w2(F.silu(gate).mul_(x)) + + +class Attention(nn.Module): + def __init__(self, heads, dim_head, bias=True, eps=1e-5): + super().__init__() + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + self.norm_q = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False) + self.norm_k = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False) + self.to_qkv = ops.Linear(inner_dim, inner_dim * 3, bias=bias) + self.to_out = ops.Linear(inner_dim, inner_dim, bias=bias) + + def forward(self, x, rotary_pos_emb=None): + batch_size, seq_len, _ = x.shape + + qkv = self.to_qkv(x) + qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head) + query, key, value = torch.chunk(qkv, 3, dim=-1) + + query = comfy.rmsnorm.rms_norm(query, self.norm_q.weight, self.norm_q.eps) + key = comfy.rmsnorm.rms_norm(key, self.norm_k.weight, self.norm_k.eps) + + if rotary_pos_emb is not None: + rot = rotary_pos_emb.shape[-3] * 2 + query[..., :rot], key[..., :rot] = comfy.quant_ops.ck.apply_rope_split_half( + query[..., :rot], key[..., :rot], rotary_pos_emb) + + out = optimized_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), + self.heads, skip_reshape=True).nan_to_num_(0.0) + return self.to_out(out) + + +class TransformerBlock(nn.Module): + def __init__(self, heads, dim_head, bias=True, eps=1e-5): + super().__init__() + dim = heads * dim_head + self.norm1 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps) + self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps) + self.scale1 = nn.Parameter(torch.empty(dim)) + self.norm2 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps) + self.ff = FeedForward(dim=dim, bias=bias) + self.scale2 = nn.Parameter(torch.empty(dim)) + + def forward(self, x, rotary_pos_emb=None): + x = x.addcmul_(self.attn(comfy.rmsnorm.rms_norm(x, self.norm1.weight, self.norm1.eps), rotary_pos_emb), self.scale1) + return x.addcmul_(self.ff(comfy.rmsnorm.rms_norm(x, self.norm2.weight, self.norm2.eps)), self.scale2) + + +class ViT3DDecoder(nn.Module): + def __init__(self, patch_size=16, patch_size_t=4, in_channels=24, out_channels=3, num_layers=36, heads=32, dim_head=64, rope_theta=100.0, + rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4): + super().__init__() + dim = heads * dim_head + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.out_channels = out_channels + self.num_register_tokens = num_register_tokens + + self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, n_dim=3) + self.x_embedder = ops.Linear(in_channels, dim) + self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim)) + # unused at inference; kept so the checkpoint loads without leftover keys + self.register_buffer("mask_token", torch.empty(1, 1, dim)) + + self.transformer_blocks = nn.ModuleList( + [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps) + for _ in range(num_layers)] + ) + + self.norm_out = ops.LayerNorm(dim, elementwise_affine=True, eps=eps) + self.proj_out = ops.Linear(dim, out_channels * patch_size_t * patch_size * patch_size) + + def forward(self, x): + B, C, latent_T, latent_H, latent_W = x.shape + + h = self.x_embedder(x.flatten(2).transpose(1, 2)) # [B, T*H*W, C] + + num_patches = h.shape[1] + num_suffix = 1 + self.num_register_tokens + + h = torch.cat([h, self.register_tokens.expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1) + + img_ids = create_token_ids((latent_T, latent_H, latent_W), x.device, x.dtype).expand(B, -1, -1) + suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype) + img_ids = torch.cat([img_ids, suffix_ids], dim=1) + + rotary_pos_emb = self.pos_embed(img_ids) + + for block in self.transformer_blocks: + h = block(h, rotary_pos_emb) + + output = self.proj_out(self.norm_out(h)) + + output = output[:, :num_patches, :] + + output = output.view( + B, latent_T, latent_H, latent_W, + self.out_channels, self.patch_size_t, self.patch_size, self.patch_size, + ) + output = output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous() + output = output.reshape( + B, self.out_channels, + latent_T * self.patch_size_t, + latent_H * self.patch_size, + latent_W * self.patch_size, + ) + return output + + +# Full VAE + +class MiniMaxH3VideoVAE(nn.Module): + def __init__( + self, + in_channels=3, + out_ch=3, + ch=128, + embed_dim=24, + z_channels=24, + ch_mult=(1, 2, 2, 4, 4, 8), + num_res_blocks=2, + space_down=(2, 2, 2, 2, 1, 1), + time_down=(1, 2, 2, 1, 1, 1), + clip_length=17, + token_drop=3, + tile_size=256, + tile_overlap_min=64, + tiling=True, + ): + super().__init__() + self.vae_ratio = int(math.prod(space_down)) + self.vae_ratio_t = int(math.prod(time_down)) + + # temporal chunking parameters + self.clip_length = clip_length + self.token_drop = token_drop + self.frame_pre_padding = (-clip_length) % self.vae_ratio_t + self.tokens_chunk_size = math.ceil(clip_length / self.vae_ratio_t) + self.token_overlap = (-token_drop) % self.tokens_chunk_size + self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0) + + # spatial tiling parameters + self.tiling = tiling + self.tile_size = tile_size + self.tile_overlap_min = tile_overlap_min + + self.encoder = EncoderFCN3D( + ch=ch, + ch_mult=list(ch_mult), + space_down=list(space_down), + time_down=list(time_down), + num_res_blocks=num_res_blocks, + in_channels=in_channels, + z_channels=z_channels, + double_z=True, + ) + self.quant_conv = ops.Conv3d(z_channels * 2, 2 * embed_dim, 1) + self.post_quant_conv = ops.Conv3d(embed_dim, z_channels, 1) + self.decoder = ViT3DDecoder( + patch_size=self.vae_ratio, + patch_size_t=self.vae_ratio_t, + in_channels=z_channels, + out_channels=out_ch, + ) + + self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN)) + self.register_buffer("latents_std", torch.tensor(LATENTS_STD)) + self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1, 1), persistent=False) + self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False) + + # single-shot forward + + def _encode_moments(self, x): + return self.quant_conv(self.encoder(x)) + + def _decode_pixels(self, z): + return self.decoder(self.post_quant_conv(z)) + + def _adaptive_encode(self, x): + if self.tiling: + return self.tiled_encode(x) + return self._encode_moments(x) + + def _adaptive_decode(self, z): + if self.tiling: + return self.tiled_decode(z) + return self._decode_pixels(z) + + # spatial tiling + + def split_tiles(self, input_len): + tile_size = self.tile_size + if tile_size >= input_len: + return [0], [input_len], [] + + N = math.ceil(input_len / tile_size) + while True: + overlaps = [self.tile_overlap_min] * (N - 1) + remaining = tile_size * N - sum(overlaps) - input_len + if remaining < 0: + N += 1 + else: + break + + remaining_units = remaining // self.vae_ratio + for i in range(remaining_units): + overlaps[i % (N - 1)] += self.vae_ratio + + tile_start_idx = [0] + for i in range(N - 1): + tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i]) + + return tile_start_idx, [tile_size] * N, overlaps + + def blend(self, a, b, blend_extent, dim): + blend_extent = min(a.shape[dim], b.shape[dim], blend_extent) + + positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype) + weight_a = 1 - positions / blend_extent + weight_b = positions / blend_extent + + shape = [1] * a.ndim + shape[dim] = blend_extent + weight_a = weight_a.view(shape) + weight_b = weight_b.view(shape) + + slice_a = [slice(None)] * a.ndim + slice_a[dim] = slice(-blend_extent, None) + slice_b = [slice(None)] * b.ndim + slice_b[dim] = slice(0, blend_extent) + + blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b + + if blend_extent < b.shape[dim]: + slice_b_rest = [slice(None)] * b.ndim + slice_b_rest[dim] = slice(blend_extent, None) + return torch.cat([blended, b[tuple(slice_b_rest)]], dim=dim) + return blended + + def tiled_encode(self, x): + height, width = x.shape[-2], x.shape[-1] + y_idx, y_len, y_overlap = self.split_tiles(height) + x_idx, x_len, x_overlap = self.split_tiles(width) + + rows = [] + for i_pos, i_len in zip(y_idx, y_len): + row = [] + for j_pos, j_len in zip(x_idx, x_len): + tile = x[..., i_pos:i_pos + i_len, j_pos:j_pos + j_len] + row.append(self._encode_moments(tile)) + rows.append(row) + + latent_y_overlap = [o // self.vae_ratio for o in y_overlap] + latent_x_overlap = [o // self.vae_ratio for o in x_overlap] + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2) + if j > 0: + tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1) + if i < len(rows) - 1: + tile = tile[..., :-latent_y_overlap[i], :] + if j < len(row) - 1: + tile = tile[..., :, :-latent_x_overlap[j]] + result_row.append(tile) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def tiled_decode(self, z): + height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio + y_idx, y_len, y_overlap = self.split_tiles(height) + x_idx, x_len, x_overlap = self.split_tiles(width) + + # Blended tiles are written straight into a pre-allocated canvas. + canvas = None + row_tails = [] + out_y = 0 + for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)): + zi, zl = i_pos // self.vae_ratio, i_len // self.vae_ratio + new_tails = [] + left_tail = None + out_x = 0 + for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)): + zj, zw = j_pos // self.vae_ratio, j_len // self.vae_ratio + tile = self._decode_pixels(z[..., zi:zi + zl, zj:zj + zw]) + if i < len(y_idx) - 1: + new_tails.append(tile[..., -y_overlap[i]:, :].clone()) + next_left_tail = tile[..., :, -x_overlap[j]:].clone() if j < len(x_idx) - 1 else None + if i > 0: + tile = self.blend(row_tails[j], tile, y_overlap[i - 1], dim=-2) + if j > 0: + tile = self.blend(left_tail, tile, x_overlap[j - 1], dim=-1) + left_tail = next_left_tail + if i < len(y_idx) - 1: + tile = tile[..., :-y_overlap[i], :] + if j < len(x_idx) - 1: + tile = tile[..., :, :-x_overlap[j]] + if canvas is None: + canvas = torch.empty(*tile.shape[:-2], height, width, dtype=tile.dtype, device=tile.device) + canvas[..., out_y:out_y + tile.shape[-2], out_x:out_x + tile.shape[-1]].copy_(tile) + out_x += tile.shape[-1] + row_tails = new_tails + out_y += tile.shape[-2] + return canvas + + # temporal chunking + + def encode_temporal(self, x): + if x.shape[2] % self.clip_length != 0: + pad_size = (-x.shape[2]) % self.clip_length + pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1) + x = torch.cat([x, pad_frames], dim=2) + + num_chunks = x.shape[2] // self.clip_length + + z_list = [] + for i in range(num_chunks): + clip_x = x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :] + z_list.append(self._adaptive_encode(clip_x)) + + z = torch.cat(z_list, dim=2) + if self.token_drop > 0: + z = z[:, :, :-self.token_drop] + return z + + def _decode_temporal_pad_frames(self, z_len, pad_tokens): + if pad_tokens <= 0: + return 0 + intra_tail = self.clip_length % self.vae_ratio_t + if intra_tail == 0: + return pad_tokens * self.vae_ratio_t + + z_len_before_pad = z_len - pad_tokens + return sum( + (intra_tail if (z_len_before_pad + k) % self.tokens_chunk_size == 0 + else self.vae_ratio_t) + for k in range(pad_tokens) + ) + + def _decode_temporal_frame_plan(self, z_len, num_chunks, pad_tokens): + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + split_count = int(self.token_drop > 0) + 1 + total_frames = 0 + final_overlap_frames = 0 + + for i in range(num_chunks): + t_start_idx = i * self.tokens_chunk_size + t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap + clip_token_len = max(0, min(t_end_idx, z_len) - min(t_start_idx, z_len)) + clip_frame_len = clip_token_len * self.vae_ratio_t + + for j in range(split_count): + f_start_idx = j * chunk_dec + f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len) + chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding) + if j == 0: + total_frames += chunk_frames + else: + final_overlap_frames = chunk_frames + + total_frames += final_overlap_frames + return total_frames - self._decode_temporal_pad_frames(z_len, pad_tokens) + + def decode_temporal(self, z): + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + split_count = int(self.token_drop > 0) + 1 + + pseudo_total_tokens = z.shape[2] + self.token_drop + + pad_tokens = 0 + remainder = pseudo_total_tokens % self.tokens_chunk_size + if remainder != 0: + pad_tokens = self.tokens_chunk_size - remainder + pseudo_total_tokens += pad_tokens + + num_chunks = pseudo_total_tokens // self.tokens_chunk_size - int(self.token_drop > 0) + if num_chunks < 1: + # too few tokens for one chunk (e.g. T_lat == 2): pad one extra chunk + pad_tokens += self.tokens_chunk_size + num_chunks += 1 + + if pad_tokens > 0: + pad_z = z[:, :, -1:, :, :].repeat(1, 1, pad_tokens, 1, 1) + z = torch.cat([z, pad_z], dim=2) + + output_frames = self._decode_temporal_frame_plan(z.shape[2], num_chunks, pad_tokens) + + dec = None + dec_overlap = None + write_pos = 0 + + def write_part(part): + nonlocal dec, write_pos + part_frames = part.shape[2] + if part_frames <= 0: + return + if dec is None: + out_shape = list(part.shape) + out_shape[2] = output_frames + dec = torch.empty(out_shape, dtype=part.dtype, device=part.device) + copy_frames = min(part_frames, max(0, dec.shape[2] - write_pos)) + if copy_frames > 0: + dec[:, :, write_pos:write_pos + copy_frames, :, :].copy_( + part[:, :, :copy_frames, :, :] + ) + write_pos += copy_frames + + for i in range(num_chunks): + t_start_idx = i * self.tokens_chunk_size + t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap + clip_z = z[:, :, t_start_idx:t_end_idx, :, :] + + clip_dec = self._adaptive_decode(clip_z) + + for j in range(split_count): + f_start_idx = j * chunk_dec + f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2]) + clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :] + clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding:, :, :] + + if j == 0: + if dec_overlap is not None: + clip_dec_chunk = self.blend( + dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3 + ) + dec_overlap = None + write_part(clip_dec_chunk) + else: + dec_overlap = clip_dec_chunk.contiguous() + + if i == num_chunks - 1 and dec_overlap is not None: + write_part(dec_overlap) + dec_overlap = None + + del clip_dec, clip_z + + return dec + + + def encode(self, x): + # x: [B, 3, T, H, W] in [-1, 1] -> normalized latents [B, 24, T_lat, H/16, W/16] + if x.ndim == 4: + x = x.unsqueeze(2) + + x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x)) + + if x.shape[2] == 1: + moments = self._adaptive_encode(x) + moments = moments[:, :, -1:, :, :] + else: + moments = self.encode_temporal(x) + + mean = torch.chunk(moments.float(), 2, dim=1)[0] + + latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(mean) + latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(mean) + return (mean - latents_mean) / latents_std + + def encode_tiled(self, x, **kwargs): + # tiling is always on internally with the reference's semantic tile sizes, ignore tiling fallbacks + return self.encode(x) + + def decode_tiled(self, z, **kwargs): + return self.decode(z) + + def decode(self, z): + # z: [B, 24, T_lat, H_lat, W_lat] normalized latents -> pixels [B, 3, T, H, W] in [-1, 1] + latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(z) + latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(z) + z = z * latents_std + latents_mean + + if z.shape[2] == 1: + dec = self._adaptive_decode(z) + dec = dec[:, :, -1:, :, :] + else: + dec = self.decode_temporal(z) + + dec = dec.float() + dec.mul_(self.pixel_std.to(dec)).add_(self.pixel_mean.to(dec)).clamp_(0.0, 1.0).mul_(2.0).sub_(1.0) + return dec diff --git a/comfy/model_base.py b/comfy/model_base.py index 7c68c6eba..6631c9eb0 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -21,6 +21,7 @@ import comfy.ldm.hunyuan3dv2_1.hunyuandit import torch import logging import comfy.ldm.lightricks.av_model +import comfy.ldm.minimax.model import comfy.ldm.lightricks.symmetric_patchifier import comfy.context_windows from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep @@ -2063,6 +2064,57 @@ class Hunyuan3Dv2_1(BaseModel): out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance])) return out +class MiniMaxH3(BaseModel): + def __init__(self, model_config, model_type=ModelType.FLOW, device=None): + super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax.model.MiniMaxH3Model) + + def extra_conds(self, **kwargs): + out = super().extra_conds(**kwargs) + cross_attn = kwargs.get("cross_attn", None) + if cross_attn is not None: + # run condition_proj + token refiner once per sampling instead of per step + cross_attn = self.diffusion_model.preprocess_text_embeds( + cross_attn.to(device=kwargs["device"], dtype=self.get_dtype_inference())) + out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn) + + latent_shapes = kwargs.get("latent_shapes", None) + if latent_shapes is not None: + out['latent_shapes'] = comfy.conds.CONDConstant(latent_shapes) + + # Everything H3-specific rides in one dict so _apply_model's dtype cast + # (which would flatten fp32 cond latents and long tags to bf16) skips it. + payload = {} + tags = kwargs.get("minimax_token_tags", None) + if tags is not None: + payload["text_token_tags"] = tags + keyframes = kwargs.get("minimax_keyframes", None) + if keyframes is not None: + payload["keyframes"] = keyframes + payload["frame_count"] = kwargs.get("minimax_frame_count", None) + payload["cond_video_latents"] = [kf["latent"] for kf in keyframes] + refs = kwargs.get("minimax_refs", None) + if refs is not None: + payload["refs"] = refs + payload["cond_video_latents"] = [r["latent"] for r in refs if "latent" in r] + payload["cond_audio_latents"] = [r["audio_latent"] for r in refs if r.get("audio_latent") is not None] + if kwargs.get("minimax_visual_cond_noise_aug", None) is not None: + payload["visual_cond_noise_aug"] = kwargs["minimax_visual_cond_noise_aug"] + if kwargs.get("minimax_audio_cond_noise_aug", None) is not None: + payload["audio_cond_noise_aug"] = kwargs["minimax_audio_cond_noise_aug"] + payload["seed"] = kwargs.get("seed", 0) + if cross_attn is not None and latent_shapes is not None and len(latent_shapes) > 1: + # packed layout built once per sampling run, h/w rounded up to the DiT's 2x2 patch + vs = latent_shapes[0] + payload["layout"] = comfy.ldm.minimax.model.PackedLayout( + cross_attn.shape[1], vs[2], (vs[3] + 1) // 2 * 2, (vs[4] + 1) // 2 * 2, + latent_shapes[1][-1], keyframes=payload.get("keyframes"), + refs=payload.get("refs"), frame_count=payload.get("frame_count")) + out['minimax_payload'] = comfy.conds.CONDConstant(payload) + return out + + def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs): + return latent_image + class TripoSplat(BaseModel): def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.triposplat.model.LatentSeqMMFlowModel) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 39e973d36..103680fd1 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -359,6 +359,35 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): # PixArt diffusers return None + if '{}video_patch_proj.weight'.format(key_prefix) in state_dict_keys and '{}audio_patch_proj.weight'.format(key_prefix) in state_dict_keys: # MiniMax H3 + dit_config = {} + dit_config["image_model"] = "minimax_h3" + dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.') + dit_config["token_refiner_num_layers"] = count_blocks(state_dict_keys, '{}token_refiner.blocks.'.format(key_prefix) + '{}.') + dit_config["hidden_size"] = state_dict['{}video_patch_proj.weight'.format(key_prefix)].shape[0] + dit_config["latents_dim"] = state_dict['{}final_layer.video_out.weight'.format(key_prefix)].shape[0] // 4 # patch 1x2x2 + dit_config["audio_latents_dim"] = state_dict['{}final_layer.audio_out.weight'.format(key_prefix)].shape[0] + dit_config["attention_head_dim"] = state_dict['{}blocks.0.attn.q_norm.weight'.format(key_prefix)].shape[0] + qkv = state_dict['{}blocks.0.attn.qkv_proj.weight'.format(key_prefix)] + dit_config["num_attention_heads"] = qkv.shape[0] // (3 * dit_config["attention_head_dim"]) + dit_config["ffn_hidden_size"] = state_dict['{}blocks.0.mlp.fc1.weight'.format(key_prefix)].shape[0] // 2 + dit_config["text_dim"] = state_dict['{}condition_proj.weight'.format(key_prefix)].shape[1] + table_key = '{}adaln_t_table'.format(key_prefix) + if table_key in state_dict_keys: + # adaln shipped over a precomputed curve basis: the adaln linears span a small shared basis of the time-embedding curve (no time embedder) + table = state_dict[table_key].shape # [grid, k] + dit_config["adaln_curve_grid"] = table[0] + dit_config["time_embed_dim"] = table[1] + else: + te = state_dict['{}time_embedder.proj_in.weight'.format(key_prefix)] + dit_config["timestep_input_dim"] = te.shape[1] + dit_config["time_embed_hidden_size"] = te.shape[0] + dit_config["time_embed_dim"] = state_dict['{}time_embedder.proj_out.weight'.format(key_prefix)].shape[0] + dit_config["rope_inv_freq_len"] = state_dict['{}rope.inv_freq'.format(key_prefix)].shape[0] + if metadata is not None and "config" in metadata: + dit_config.update(json.loads(metadata["config"]).get("transformer", {})) + return dit_config + if '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: #Lightricks ltxv dit_config = {} dit_config["image_model"] = "ltxav" if f'{key_prefix}audio_adaln_single.linear.weight' in state_dict_keys else "ltxv" diff --git a/comfy/ops.py b/comfy/ops.py index 6c3845eef..077a42351 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -943,13 +943,61 @@ if CUBLAS_IS_AVAILABLE: # ============================================================================== # Mixed Precision Operations # ============================================================================== +from . import quant_ops from .quant_ops import ( QuantizedTensor, QUANT_ALGOS, TensorCoreFP8Layout, + TensorWiseINT8Layout, get_layout_class, ) +def _swiglu_eager(x): + gate, up = x.chunk(2, dim=-1) + return torch.nn.functional.silu(gate).mul_(up) + + +INPUT_ACT_EAGER = { + "gelu_tanh": lambda x: torch.nn.functional.gelu(x, approximate="tanh"), + "swiglu": _swiglu_eager, +} + + +def linear_input_act(linear, x, input_act): + """``linear(act(x))``, with ``act`` folded into an INT8 activation quantizer. + + An INT8 linear quantizes its input anyway, so an elementwise activation can + ride along inside that kernel instead of writing a full-size intermediate to + HBM and reading it straight back. Worth it for an MLP's down-projection, + where the intermediate is several times the hidden size. + + """ + weight = linear.weight + if (comfy.model_management.in_training + or not isinstance(weight, QuantizedTensor) + or weight._layout_cls != "TensorWiseINT8Layout" + or getattr(weight._params, "transposed", False)): + return linear(INPUT_ACT_EAGER[input_act](x)) + + # want_requant keeps a vbar-streamed layer on the INT8 path when a LoRA is + # patched in on the fly; without it the cast hands back a dequantized weight. + weight, bias, offload_stream = cast_bias_weight( + linear, x, offloadable=True, compute_dtype=x.dtype, want_requant=True) + try: + if not isinstance(weight, QuantizedTensor): + # A LoRA weight_function, or activations whose dtype differs from the + # weight's, make the cast hand back a dequantized tensor. + return torch.nn.functional.linear(INPUT_ACT_EAGER[input_act](x), weight, bias) + qdata, scale = TensorWiseINT8Layout.get_plain_tensors(weight) + return quant_ops.ck.int8_linear( + x, qdata, scale, bias, x.dtype, + convrot=getattr(weight._params, "convrot", False), + convrot_groupsize=getattr(weight._params, "convrot_groupsize", 256), + input_act=input_act, + ) + finally: + uncast_bias_weight(linear, weight, bias, offload_stream) + class QuantLinearFunc(torch.autograd.Function): """Custom autograd function for quantized linear: quantized forward, optionally FP8 backward. @@ -1259,7 +1307,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec def state_dict(self, *args, destination=None, prefix="", **kwargs): sd = destination if destination is not None else {} - return _quantized_weight_state_dict(self, sd, prefix, extra_quant_params=("input_scale",)) + return _quantized_weight_state_dict(self, sd, prefix, extra_quant_params=("input_scale", "pre_quant_scale")) def _forward(self, input, weight, bias): return torch.nn.functional.linear(input, weight, bias) @@ -1298,6 +1346,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec def forward(self, input, *args, **kwargs): run_every_op() + # ModelOpt AWQ-style smoothing + pre_quant_scale = getattr(self, 'pre_quant_scale', None) + if pre_quant_scale is not None: + input = input * comfy.model_management.cast_to_device(pre_quant_scale, input.device, input.dtype) + input_shape = input.shape reshaped_nd = False #If cast needs to apply lora, it should be done in the compute dtype diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 15f9b1fdb..53586956a 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -240,7 +240,7 @@ QUANT_ALGOS = { }, "nvfp4": { "storage_t": torch.uint8, - "parameters": {"weight_scale", "weight_scale_2", "input_scale"}, + "parameters": {"weight_scale", "weight_scale_2", "input_scale", "pre_quant_scale"}, "comfy_tensor_layout": "TensorCoreNVFP4Layout", "group_size": 16, }, diff --git a/comfy/sd.py b/comfy/sd.py index caf78222d..8d670106e 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -72,6 +72,9 @@ import comfy.text_encoders.ace15 import comfy.text_encoders.longcat_image import comfy.text_encoders.qwen35 import comfy.text_encoders.qwen3vl +import comfy.text_encoders.minimax +import comfy.ldm.minimax.vae +import comfy.ldm.minimax.audio_vae import comfy.text_encoders.boogu import comfy.text_encoders.ernie import comfy.text_encoders.gemma4 @@ -936,6 +939,50 @@ class VAE: #Force cast it for --disable-dynamic-vram users until there is a true core fix. if not comfy.memory_management.aimdo_enabled: self.disable_offload = True + elif "decoder.transformer_blocks.0.scale1" in sd and "encoder.down.5.block.0.conv1.weight" in sd: # MiniMax H3 video VAE + self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE() + self.latent_channels = 24 + self.latent_dim = 3 + # frames 17k+5 <-> latents 5k+2, 16x spatial + self.upscale_ratio = (lambda a: max(1, (a - 2) // 5 * 17 + 5), 16, 16) + self.upscale_index_formula = (4, 16, 16) + self.downscale_ratio = (lambda a: max(1, (a - 5) // 17 * 5 + 2) if a > 1 else 1, 16, 16) + self.downscale_index_formula = (4, 16, 16) + self.working_dtypes = [torch.float16, torch.float32] + # the model tiles internally (256px spatial, 17-frame temporal chunks) + self.handles_tiling = True + def estimate_encode_memory(frames, height, width, dtype): + fixed = 110_000_000 if frames == 1 else 1_300_000_000 + elements_per_pixel = 7 if frames == 1 else 9.5 + return (elements_per_pixel * frames * height * width + fixed) * model_management.dtype_size(dtype) * 1.03 + + def estimate_decode_memory(frames, height, width, dtype): + fixed = 110_000_000 if frames <= 22 else 270_000_000 + return (9.5 * frames * height * width + fixed) * model_management.dtype_size(dtype) * 1.03 + + self.memory_used_encode = lambda shape, dtype: estimate_encode_memory(shape[2], shape[3], shape[4], dtype) + self.memory_used_decode = lambda shape, dtype: estimate_decode_memory(self.upscale_ratio[0](shape[2]), shape[3] * self.upscale_ratio[1], shape[4] * self.upscale_ratio[2], dtype) + elif "pre_block.attn.zero_k_bias" in sd: # MiniMax H3 audio VAE (DAC encoder + BigVGAN decoder) + self.first_stage_model = comfy.ldm.minimax.audio_vae.MiniMaxH3AudioVAE() + self.latent_channels = 32 + self.output_channels = 2 + self.pad_channel_value = "replicate" + self.audio_sample_rate = 32000 + self.upscale_ratio = 800 + self.downscale_ratio = 800 + self.latent_dim = 2 # [B, 32, stereo 2, T] + self.process_output = lambda audio: audio + self.process_input = lambda audio: audio + self.working_dtypes = [torch.float32] + # encode gets the waveform shape [B, 2, samples], decode the latent shape [B, 32, 2, T] + def estimate_encode_memory(samples, dtype): + return (900 * samples + 105_000_000) * model_management.dtype_size(dtype) * 1.03 + + def estimate_decode_memory(samples, dtype): + return max(42_000_000, 220 * samples + 20_000_000) * model_management.dtype_size(dtype) * 1.03 + + self.memory_used_encode = lambda shape, dtype: estimate_encode_memory(shape[2], dtype) + self.memory_used_decode = lambda shape, dtype: estimate_decode_memory(shape[-1] * self.upscale_ratio, dtype) elif "gs.base_offset_scale" in sd and "octree.out_proj.weight" in sd: # TripoSplat octree gaussian decoder self.first_stage_model = comfy.ldm.triposplat.vae.OctreeGaussianDecoder() self.latent_channels = 16 @@ -1393,6 +1440,7 @@ class CLIPType(Enum): KREA2 = 32 JOYIMAGE = 33 MAGE = 34 + MINIMAX = 35 @@ -1449,6 +1497,7 @@ class TEModel(Enum): QWEN3VL_4B = 34 QWEN3VL_8B = 35 GEMMA_4_12B = 36 + QWEN3VL_32B = 37 def detect_te_model(sd): @@ -1515,6 +1564,9 @@ def detect_te_model(sd): return TEModel.QWEN35_2B if "model.visual.deepstack_merger_list.0.norm.weight" in sd: # DeepStack is unique to Qwen3-VL return TEModel.QWEN3VL_4B if sd["model.visual.merger.linear_fc2.weight"].shape[0] == 2560 else TEModel.QWEN3VL_8B + if "visual.deepstack_merger_list.0.norm.weight" in sd and "model.layers.49.self_attn.q_proj.weight" in sd: + # MiniMax H3 conditioning encoder: Qwen3-VL-32B, truncated to 50 layers + return TEModel.QWEN3VL_32B if "model.layers.0.post_attention_layernorm.weight" in sd: weight = sd['model.layers.0.post_attention_layernorm.weight'] if 'model.layers.0.self_attn.q_norm.weight' in sd: @@ -1744,6 +1796,9 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip qwen3vl_type = {TEModel.QWEN3VL_4B: "qwen3vl_4b", TEModel.QWEN3VL_8B: "qwen3vl_8b"}[te_model] clip_target.clip = comfy.text_encoders.qwen3vl.te(**llama_detect(clip_data), model_type=qwen3vl_type) clip_target.tokenizer = comfy.text_encoders.qwen3vl.tokenizer(model_type=qwen3vl_type) + elif te_model == TEModel.QWEN3VL_32B: + clip_target.clip = comfy.text_encoders.minimax.te(**llama_detect(clip_data)) + clip_target.tokenizer = comfy.text_encoders.minimax.MiniMaxH3Tokenizer elif te_model == TEModel.QWEN3_06B: clip_target.clip = comfy.text_encoders.anima.te(**llama_detect(clip_data)) clip_target.tokenizer = comfy.text_encoders.anima.AnimaTokenizer diff --git a/comfy/supported_models.py b/comfy/supported_models.py index ca89850a5..51b58ed1e 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -15,6 +15,7 @@ import comfy.text_encoders.flux import comfy.text_encoders.genmo import comfy.text_encoders.lt import comfy.text_encoders.hunyuan_video +import comfy.text_encoders.minimax import comfy.text_encoders.cosmos import comfy.text_encoders.lumina2 import comfy.text_encoders.wan @@ -955,6 +956,33 @@ class LTXAV(LTXV): out = model_base.LTXAV(self, device=device) return out +class MiniMaxH3(supported_models_base.BASE): + unet_config = { + "image_model": "minimax_h3", + } + + sampling_settings = { + "shift": 12.0, + } + + unet_extra_config = {} + latent_format = latent_formats.MiniMaxH3AV + + memory_usage_factor = 0.114 + + supported_inference_dtypes = [torch.bfloat16, torch.float32] + + vae_key_prefix = ["vae."] + text_encoder_key_prefix = ["text_encoders."] + + def get_model(self, state_dict, prefix="", device=None): + return model_base.MiniMaxH3(self, device=device) + + def clip_target(self, state_dict={}, prefix=""): + pref = self.text_encoder_key_prefix[0] + detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_32b.transformer.".format(pref)) + return supported_models_base.ClipTarget(comfy.text_encoders.minimax.MiniMaxH3Tokenizer, comfy.text_encoders.minimax.te(**detect)) + class HunyuanVideo(supported_models_base.BASE): unet_config = { "image_model": "hunyuan_video", @@ -2407,6 +2435,7 @@ models = [ GenmoMochi, LTXV, LTXAV, + MiniMaxH3, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index 40d04007e..f5c5597ef 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -264,6 +264,17 @@ class Qwen3VL_4BConfig(Qwen3VL_8BConfig): intermediate_size: int = 9728 lm_head: bool = False # 4B ties word embeddings +@dataclass +class Qwen3VL_32BConfig(Qwen3VL_8BConfig): + # MiniMax H3 conditioning checkpoint: truncated to the first 50 of 64 layers, + # consumed as the unnormalized hidden state after layer 50 (no final norm, no lm_head) + hidden_size: int = 5120 + intermediate_size: int = 25600 + num_hidden_layers: int = 50 + num_attention_heads: int = 64 + lm_head: bool = False + final_norm: bool = False + @dataclass class Ovis25_2BConfig: vocab_size: int = 151936 diff --git a/comfy/text_encoders/minimax.py b/comfy/text_encoders/minimax.py new file mode 100644 index 000000000..c2dc47f7f --- /dev/null +++ b/comfy/text_encoders/minimax.py @@ -0,0 +1,201 @@ +"""MiniMax H3 text/vision conditioning: Qwen3-VL-32B (truncated to 50 layers). + +The H3 presentation is NOT chat-templated: token ids are raw prompt/label text +(no special tokens) with explicit vision blocks spliced in: + + t2va: + fl2va: ": " [": " ] + ref2va: per condition in request order (1-based ordinals per type): + image -> ": " + audio -> "