From 45ffd5430beeccf63682b5f8b569faad45fd60e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Sat, 25 Jul 2026 06:14:01 +0300 Subject: [PATCH] feat: Support MageFlow (CORE-372) (#15026) --- comfy/ldm/mage_flow/model.py | 187 ++++++++++++ comfy/ldm/mage_flow/vae.py | 477 +++++++++++++++++++++++++++++++ comfy/model_base.py | 16 +- comfy/model_detection.py | 7 + comfy/sd.py | 18 ++ comfy/supported_models.py | 31 ++ comfy/text_encoders/mage_flow.py | 94 ++++++ comfy/text_encoders/qwen3vl.py | 4 +- comfy_extras/nodes_mage.py | 103 +++++++ nodes.py | 3 +- 10 files changed, 935 insertions(+), 5 deletions(-) create mode 100644 comfy/ldm/mage_flow/model.py create mode 100644 comfy/ldm/mage_flow/vae.py create mode 100644 comfy/text_encoders/mage_flow.py create mode 100644 comfy_extras/nodes_mage.py diff --git a/comfy/ldm/mage_flow/model.py b/comfy/ldm/mage_flow/model.py new file mode 100644 index 000000000..92a5faa52 --- /dev/null +++ b/comfy/ldm/mage_flow/model.py @@ -0,0 +1,187 @@ +# Mage-Flow (https://github.com/microsoft/Mage) native-resolution MMDiT (MIT) +# Architecture is a 12-layer variant of the Qwen-Image double-stream block with +# patch_size=1 (no 2x2 packing), unrotated text tokens and a bf16-rounded +# timestep frequency table. +import math +import torch +import torch.nn as nn +from typing import Optional, Tuple + +from comfy.ldm.lightricks.model import TimestepEmbedding +from comfy.ldm.flux.layers import EmbedND +from comfy.ldm.qwen_image.model import QwenImageTransformerBlock, LastLayer +import comfy.patcher_extension + + +class MageTimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim, dtype=None, device=None, operations=None): + super().__init__() + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim, + dtype=dtype, device=device, operations=operations + ) + + def forward(self, timestep, hidden_states): + timestep = timestep.to(hidden_states.dtype) + half_dim = 128 + exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) / half_dim + emb = torch.exp(exponent).to(timestep.dtype) + emb = timestep[:, None].float() * emb[None, :] + emb = 1000.0 * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) # flip_sin_to_cos + return self.timestep_embedder(emb.to(dtype=hidden_states.dtype)) + + +class MageFlowTransformer2DModel(nn.Module): + def __init__( + self, + in_channels: int = 128, + out_channels: Optional[int] = 128, + num_layers: int = 12, + attention_head_dim: int = 128, + num_attention_heads: int = 24, + joint_attention_dim: int = 2560, + axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), + image_model=None, + dtype=None, + device=None, + operations=None, + ): + super().__init__() + self.dtype = dtype + self.patch_size = 1 + self.in_channels = in_channels + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.pe_embedder = EmbedND(dim=attention_head_dim, theta=10000, axes_dim=list(axes_dims_rope)) + + self.time_text_embed = MageTimestepProjEmbeddings(embedding_dim=self.inner_dim, dtype=dtype, device=device, operations=operations) + + self.txt_norm = operations.RMSNorm(joint_attention_dim, eps=1e-6, dtype=dtype, device=device) + self.img_in = operations.Linear(in_channels, self.inner_dim, dtype=dtype, device=device) + self.txt_in = operations.Linear(joint_attention_dim, self.inner_dim, dtype=dtype, device=device) + + self.transformer_blocks = nn.ModuleList([ + QwenImageTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + dtype=dtype, + device=device, + operations=operations + ) + for _ in range(num_layers) + ]) + + self.norm_out = LastLayer(self.inner_dim, self.inner_dim, dtype=dtype, device=device, operations=operations) + self.proj_out = operations.Linear(self.inner_dim, self.out_channels, bias=True, dtype=dtype, device=device) + + def process_img(self, x, index=0): + # patch_size=1: tokens are raw latent pixels, no 2x2 packing. + bs, c, h, w = x.shape + hidden_states = x.movedim(1, -1).reshape(bs, h * w, c) + + img_ids = torch.zeros((h, w, 3), device=x.device) + # Frame axis: positive image index (0 = target, 1..N = reference images). + img_ids[:, :, 0] = index + # Mage scale_rope centering: positions [-ceil(n/2), floor(n/2)), i.e. + # offset by (n - n//2). Differs from Qwen-Image's -(n//2) for odd sizes. + img_ids[:, :, 1] = img_ids[:, :, 1] + torch.arange(h, device=x.device)[:, None] - (h - h // 2) + img_ids[:, :, 2] = img_ids[:, :, 2] + torch.arange(w, device=x.device)[None, :] - (w - w // 2) + return hidden_states, img_ids.reshape(h * w, 3).unsqueeze(0).expand(bs, -1, -1), (h, w) + + def forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, **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, attention_mask, ref_latents, transformer_options, **kwargs) + + def _forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, control=None, **kwargs): + if attention_mask is not None and not torch.is_floating_point(attention_mask): + attention_mask = (attention_mask - 1).to(x.dtype) * torch.finfo(x.dtype).max + + hidden_states, img_ids, orig_shape = self.process_img(x) + num_embeds = hidden_states.shape[1] + + if ref_latents is not None: + ref_num_tokens = [] + index = 0 + for ref in ref_latents: + index += 1 + kontext, kontext_ids, _ = self.process_img(ref, index=index) + hidden_states = torch.cat([hidden_states, kontext], dim=1) + img_ids = torch.cat([img_ids, kontext_ids], dim=1) + ref_num_tokens.append(kontext.shape[1]) + transformer_options = transformer_options.copy() + transformer_options["reference_image_num_tokens"] = ref_num_tokens + + # Text tokens are not rotated in Mage-Flow: RoPE at position 0 is the + # identity rotation. + txt_ids = torch.zeros((x.shape[0], context.shape[1], 3), device=x.device) + + hidden_states = self.img_in(hidden_states) + context = self.txt_norm(context) + context = self.txt_in(context) + + temb = self.time_text_embed(timestep, hidden_states) + + patches_replace = transformer_options.get("patches_replace", {}) + patches = transformer_options.get("patches", {}) + blocks_replace = patches_replace.get("dit", {}) + + if "post_input" in patches: + for p in patches["post_input"]: + out = p({"img": hidden_states, "txt": context, "img_ids": img_ids, "txt_ids": txt_ids, "transformer_options": transformer_options}) + hidden_states = out["img"] + context = out["txt"] + img_ids = out["img_ids"] + txt_ids = out["txt_ids"] + + ids = torch.cat((txt_ids, img_ids), dim=1) + image_rotary_emb = self.pe_embedder(ids).contiguous() + del ids, txt_ids, img_ids + + transformer_options["total_blocks"] = len(self.transformer_blocks) + transformer_options["block_type"] = "double" + for i, block in enumerate(self.transformer_blocks): + transformer_options["block_index"] = i + if ("double_block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=attention_mask, temb=args["vec"], image_rotary_emb=args["pe"], transformer_options=args["transformer_options"]) + return out + out = blocks_replace[("double_block", i)]({"img": hidden_states, "txt": context, "vec": temb, "pe": image_rotary_emb, "transformer_options": transformer_options}, {"original_block": block_wrap}) + hidden_states = out["img"] + context = out["txt"] + else: + context, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=context, + encoder_hidden_states_mask=attention_mask, + temb=temb, + image_rotary_emb=image_rotary_emb, + transformer_options=transformer_options, + ) + + if "double_block" in patches: + for p in patches["double_block"]: + out = p({"img": hidden_states, "txt": context, "x": x, "block_index": i, "transformer_options": transformer_options}) + hidden_states = out["img"] + context = out["txt"] + + if control is not None: # Controlnet + control_i = control.get("input") + if i < len(control_i): + add = control_i[i] + if add is not None: + hidden_states[:, :add.shape[1]] += add + + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states[:, :num_embeds] + h, w = orig_shape + return hidden_states.reshape(x.shape[0], h, w, self.out_channels).movedim(-1, 1) diff --git a/comfy/ldm/mage_flow/vae.py b/comfy/ldm/mage_flow/vae.py new file mode 100644 index 000000000..e6e21b99f --- /dev/null +++ b/comfy/ldm/mage_flow/vae.py @@ -0,0 +1,477 @@ +# Mage-VAE (https://github.com/microsoft/Mage) (MIT) +# Symmetric one-step diffusion codec: DConvEncoder (image -> 128ch latent) and +# DConvDenoiser + CoD Decoder (latent -> image). 16x downsample, latents in the +# Flux.2-VAE-anchored space (no patch packing, no BN normalization). +# Both encode and decode are single forward passes at t=0. +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import comfy.ops +from comfy.ldm.modules.diffusionmodules.model import vae_attention + +ops = comfy.ops.disable_weight_init + + +def nonlinearity(x): + return torch.nn.functional.silu(x) + + +def Normalize(in_channels): + return ops.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + +def modulate(x, shift, scale): + if x.dim() == 4: + b, c = x.shape[:2] + return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1) + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +class LayerNorm2d(ops.LayerNorm): + def __init__(self, num_channels, eps=1e-6, affine=True): + super().__init__(num_channels, eps=eps, elementwise_affine=affine) + + def forward(self, x): + x = x.permute(0, 2, 3, 1).contiguous() + x = super().forward(x) + return x.permute(0, 3, 1, 2).contiguous() + + +class TimestepEmbedder(nn.Module): + """DConv-style timestep MLP (max_period=10000, freq_size=256).""" + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + ops.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + ops.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half + ).to(t.device) + args = t[:, None].float() * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) + return emb + + def forward(self, t, dtype): + emb = self.timestep_embedding(t, self.frequency_embedding_size) + return self.mlp(emb.to(dtype)) + + +class BottleneckPatchEmbed(nn.Module): + """Image patch embed concatenated with a per-patch conditioning vector.""" + + def __init__(self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True): + super().__init__() + self.proj1 = ops.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False) + self.proj2 = ops.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias) + + def forward(self, x, cond): + return self.proj2(torch.cat([self.proj1(x), cond], dim=1)) + + +class DiCoBlock(nn.Module): + """DConv block with adaLN modulation.""" + + def __init__(self, hidden_size, mlp_ratio=4.0): + super().__init__() + self.conv1 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True) + self.conv2 = ops.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True) + self.conv3 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True) + + self.ca = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + ops.Conv2d(hidden_size, hidden_size, 1, bias=True), + nn.Sigmoid(), + ) + + ffn = int(mlp_ratio * hidden_size) + self.conv4 = ops.Conv2d(hidden_size, ffn, 1, bias=True) + self.conv5 = ops.Conv2d(ffn, hidden_size, 1, bias=True) + + self.norm1 = LayerNorm2d(hidden_size, affine=False) + self.norm2 = LayerNorm2d(hidden_size, affine=False) + + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + ops.Linear(hidden_size, 6 * hidden_size, bias=True), + ) + + def forward(self, inp, c): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) + x = modulate(self.norm1(inp), shift_msa, scale_msa) + x = F.gelu(self.conv2(self.conv1(x))) + x = x * self.ca(x) + x = self.conv3(x) + x = inp + gate_msa[..., None, None] * x + x = x + gate_mlp[..., None, None] * self.conv5( + F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp))) + ) + return x + + +class EncoderDiCoBlock(nn.Module): + """DiCoBlock without adaLN, for the encoder head pathway.""" + + def __init__(self, hidden_size, mlp_ratio=4.0): + super().__init__() + self.conv1 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True) + self.conv2 = ops.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True) + self.conv3 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True) + self.ca = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + ops.Conv2d(hidden_size, hidden_size, 1, bias=True), + nn.Sigmoid(), + ) + ffn = int(mlp_ratio * hidden_size) + self.conv4 = ops.Conv2d(hidden_size, ffn, 1, bias=True) + self.conv5 = ops.Conv2d(ffn, hidden_size, 1, bias=True) + self.norm1 = LayerNorm2d(hidden_size) + self.norm2 = LayerNorm2d(hidden_size) + + def forward(self, inp): + x = self.norm1(inp) + x = F.gelu(self.conv2(self.conv1(x))) + x = x * self.ca(x) + x = self.conv3(x) + x = inp + x + return x + self.conv5(F.gelu(self.conv4(self.norm2(x)))) + + +class NerfEmbedder(nn.Module): + """Patch-position embedder used by the DConv decoder x-pathway.""" + + def __init__(self, in_channels, hidden_size_input, max_freqs=8): + super().__init__() + self.max_freqs = max_freqs + self.embedder = nn.Sequential( + ops.Linear(in_channels + max_freqs ** 2, hidden_size_input, bias=True), + ) + + def fetch_pos(self, patch_size, device, dtype): + pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype) + pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij") + pos_x = pos_x.reshape(-1, 1, 1) + pos_y = pos_y.reshape(-1, 1, 1) + freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device) + fx = freqs[None, :, None] + fy = freqs[None, None, :] + coeffs = (1 + fx * fy) ** -1 + dct_x = torch.cos(pos_x * fx * torch.pi) + dct_y = torch.cos(pos_y * fy * torch.pi) + return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2) + + def forward(self, x): + B, P2, _ = x.shape + ps = int(P2 ** 0.5) + dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1) + return self.embedder(torch.cat([x, dct], dim=-1)) + + +class NerfFinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm = ops.RMSNorm(hidden_size, eps=1e-6) + self.linear = ops.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x): + return self.linear(self.norm(x)) + + +class MLPResBlock(nn.Module): + def __init__(self, channels): + super().__init__() + self.in_ln = ops.LayerNorm(channels, eps=1e-6) + self.mlp = nn.Sequential( + ops.Linear(channels, channels, bias=True), + nn.SiLU(), + ops.Linear(channels, channels, bias=True), + ) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + ops.Linear(channels, 3 * channels, bias=True), + ) + + def forward(self, x, y): + shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1) + h = self.in_ln(x) * (1 + scale) + shift + return x + gate * self.mlp(h) + + +class SimpleMLPAdaLN(nn.Module): + """Final small MLP that maps NerfEmbedder features to per-patch RGB.""" + + def __init__(self, in_channels, model_channels, out_channels, z_channels, num_res_blocks, patch_size): + super().__init__() + self.in_channels = in_channels + self.model_channels = model_channels + self.out_channels = out_channels + self.num_res_blocks = num_res_blocks + self.patch_size = patch_size + + self.cond_embed = ops.Linear(z_channels, patch_size ** 2 * model_channels) + self.input_proj = ops.Linear(in_channels, model_channels) + + self.res_blocks = nn.ModuleList(MLPResBlock(model_channels) for _ in range(num_res_blocks)) + + def forward(self, x, c): + x = self.input_proj(x) + c = self.cond_embed(c).reshape(c.shape[0], self.patch_size ** 2, -1) + for block in self.res_blocks: + x = block(x, c) + return x + + +class ResnetBlock(nn.Module): + """GroupNorm + Conv ResBlock used by the CoD Decoder.""" + + def __init__(self, *, in_channels, out_channels=None): + super().__init__() + out_channels = out_channels or in_channels + self.in_channels = in_channels + self.out_channels = out_channels + + self.norm1 = Normalize(in_channels) + self.conv1 = ops.Conv2d(in_channels, out_channels, 3, padding=1) + self.norm2 = Normalize(out_channels) + self.conv2 = ops.Conv2d(out_channels, out_channels, 3, padding=1) + if in_channels != out_channels: + self.nin_shortcut = ops.Conv2d(in_channels, out_channels, 1) + + def forward(self, x): + h = self.conv1(nonlinearity(self.norm1(x))) + h = self.conv2(nonlinearity(self.norm2(h))) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class AttnBlock(nn.Module): + """Patched (windowed) self-attention used by the CoD Decoder.""" + + def __init__(self, in_channels, patch_size=32): + super().__init__() + self.in_channels = in_channels + self.patch_size = patch_size + self.norm = Normalize(in_channels) + self.q = ops.Conv2d(in_channels, in_channels, 1) + self.k = ops.Conv2d(in_channels, in_channels, 1) + self.v = ops.Conv2d(in_channels, in_channels, 1) + self.proj_out = ops.Conv2d(in_channels, in_channels, 1) + # VAE attention selection: full-precision backends only (no sage/quantized attention) + self.optimized_attention = vae_attention() + + def forward(self, x): + h_ = self.norm(x) + Q = self.q(h_) + K = self.k(h_) + V = self.v(h_) + + d = self.patch_size + b, c, H, W = Q.shape + pad_h = (d - H % d) % d + pad_w = (d - W % d) % d + if pad_h or pad_w: + Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate") + K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate") + V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate") + _, _, H_pad, W_pad = Q.shape + nph, npw = H_pad // d, W_pad // d + np_ = nph * npw + + def to_patches(t): + return (t.reshape(b, c, nph, d, npw, d) + .permute(0, 2, 4, 1, 3, 5) + .reshape(b * np_, c, d * d)) + + # [b*np, c, d*d]: attention over the d*d spatial positions of each window + Q = to_patches(Q) + K = to_patches(K) + V = to_patches(V) + + h_ = self.optimized_attention(Q, K, V) + h_ = h_.reshape(b, nph, npw, c, d, d).permute(0, 3, 1, 4, 2, 5).reshape(b, c, H_pad, W_pad) + if pad_h or pad_w: + h_ = h_[:, :, :H, :W] + return x + self.proj_out(h_) + + +class CoDDecoder(nn.Module): + """CoD Decoder: latent -> conditioning features for the denoiser (ds=16, light).""" + + def __init__(self, out_ch=384, z_ch=128): + super().__init__() + self.conv_in = ops.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1) + self.block = nn.Sequential( + ResnetBlock(in_channels=out_ch, out_channels=out_ch), + AttnBlock(out_ch, patch_size=32), + ResnetBlock(in_channels=out_ch, out_channels=out_ch), + AttnBlock(out_ch, patch_size=32), + ResnetBlock(in_channels=out_ch, out_channels=out_ch), + ) + self.norm_out = Normalize(out_ch) + self.conv_out = ops.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1) + self.ada = nn.Identity() + + def forward(self, z): + h = self.block(self.conv_in(z)) + h = self.conv_out(nonlinearity(self.norm_out(h))) + return self.ada(h) + + +class DConvEncoder(nn.Module): + """DConvEncoder: image -> packed (mean, logvar) latent.""" + + def __init__( + self, + z_ch=128, + hidden_size=384, + num_blocks=21, + patch_size=16, + mlp_ratio=4.0, + head_size=768, + num_head_blocks=2, + out_ch_mult=2, + ): + super().__init__() + self.z_ch = z_ch + self.patch_size = patch_size + self.patch_cond_embed = ops.Conv2d(3, head_size, kernel_size=patch_size, stride=patch_size, bias=True) + self.head_blocks = nn.ModuleList([ + EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) for _ in range(num_head_blocks) + ]) + self.proj_down = ops.Conv2d(head_size, hidden_size, kernel_size=1, bias=True) + self.z_proj = ops.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True) + self.fuse_proj = ops.Conv2d(hidden_size * 2, hidden_size, kernel_size=1, bias=True) + self.t_embedder = TimestepEmbedder(hidden_size) + self.blocks = nn.ModuleList([ + DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks) + ]) + self.norm_out = LayerNorm2d(hidden_size) + self.proj_out = ops.Conv2d(hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True) + + def forward_pred(self, z_t, t, y): + cond = self.patch_cond_embed(y) + for block in self.head_blocks: + cond = block(cond) + cond = self.proj_down(cond) + + s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1)) + c = self.t_embedder(t.view(-1), y.dtype) + for block in self.blocks: + s = block(s, c) + return self.proj_out(self.norm_out(s)) + + +class YEmbedder(nn.Module): + """Holds only the CoD decoder (the original Flux2-VAE encoder side is dropped at load).""" + + def __init__(self, ch=384, z_ch=128): + super().__init__() + self.decoder = CoDDecoder(out_ch=ch, z_ch=z_ch) + + +class DConvDenoiser(nn.Module): + """One-step DConv denoiser: latent (via cond) + zero noise -> reconstructed image.""" + + def __init__( + self, + patch_size=16, + in_channels=3, + hidden_size=384, + hidden_size_x=32, + mlp_ratio=4.0, + num_blocks=24, + num_cond_blocks=21, + bottleneck_dim=128, + ): + super().__init__() + self.in_channels = in_channels + self.patch_size = patch_size + self.hidden_size = hidden_size + self.num_cond_blocks = num_cond_blocks + + self.t_embedder = TimestepEmbedder(hidden_size) + self.y_embedder_x = ops.Conv2d(hidden_size, hidden_size_x * patch_size ** 2, 1, 1, 0) + self.x_embedder = NerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8) + self.s_embedder = BottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True) + self.blocks = nn.ModuleList([ + DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks) + ]) + self.dec_net = SimpleMLPAdaLN( + in_channels=hidden_size_x, + model_channels=hidden_size_x, + out_channels=in_channels, + z_channels=hidden_size, + num_res_blocks=num_blocks - num_cond_blocks, + patch_size=patch_size, + ) + self.final_layer = NerfFinalLayer(hidden_size_x, in_channels) + self.y_embedder = YEmbedder(ch=hidden_size, z_ch=bottleneck_dim) + + def forward(self, x, t, cond): + b, _, h, w = x.shape + c = self.t_embedder(t.view(-1), x.dtype) + + s = self.s_embedder(x, cond) + for block in self.blocks: + s = block(s, c) + + length = s.shape[-2] * s.shape[-1] + s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size) + + x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size) + x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1) + x = x.reshape(b, -1, self.patch_size ** 2, length).permute(0, 3, 2, 1).flatten(0, 1) + x = self.x_embedder(x) + + x = self.dec_net(x, s) + x = self.final_layer(x) + x = x.transpose(1, 2).reshape(b, length, -1) + return torch.nn.functional.fold( + x.transpose(1, 2).contiguous(), (h, w), + kernel_size=self.patch_size, stride=self.patch_size, + ) + + +class MageVAE(nn.Module): + """ + Encode: DConvEncoder (one-step at t=0) -> posterior mean [B, 128, H/16, W/16] + Decode: DConvDenoiser + CoD Decoder -> image [B, 3, H, W] in [-1, 1] + """ + + latent_channels = 128 + downsample_factor = 16 + + def __init__(self): + super().__init__() + self.dconv_encoder = DConvEncoder() + self.decoder_model = DConvDenoiser() + + def encode(self, x): + B, _, H, W = x.shape + ps = self.dconv_encoder.patch_size + z_t = torch.zeros(B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype) + t = torch.zeros(B, device=x.device, dtype=x.dtype) + out = self.dconv_encoder.forward_pred(z_t, t, x) + return out[:, : self.latent_channels] # posterior mean (sample_posterior=False) + + def decode(self, z): + cond = self.decoder_model.y_embedder.decoder(z) + B = z.shape[0] + H = z.shape[2] * self.downsample_factor + W = z.shape[3] * self.downsample_factor + noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype) + t = torch.zeros(B, device=z.device, dtype=z.dtype) + return self.decoder_model.forward(noise, t, cond) diff --git a/comfy/model_base.py b/comfy/model_base.py index 3494925be..50c73a431 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -58,6 +58,7 @@ import comfy.ldm.omnigen.omnigen2 import comfy.ldm.seedvr.model import comfy.ldm.boogu.model import comfy.ldm.qwen_image.model +import comfy.ldm.mage_flow.model import comfy.ldm.joyimage.model import comfy.ldm.ideogram4.model import comfy.ldm.krea2.model @@ -2243,8 +2244,8 @@ class Boogu(Omnigen2): self.memory_usage_factor_conds = ("ref_latents",) class QwenImage(BaseModel): - def __init__(self, model_config, model_type=ModelType.FLUX, device=None): - super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel) + def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel): + super().__init__(model_config, model_type, device=device, unet_model=unet_model) self.memory_usage_factor_conds = ("ref_latents",) def extra_conds(self, **kwargs): @@ -2274,6 +2275,17 @@ class QwenImage(BaseModel): out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16]) return out +class MageFlow(QwenImage): + def __init__(self, model_config, model_type=ModelType.FLOW, device=None): + super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.mage_flow.model.MageFlowTransformer2DModel) + + def extra_conds_shapes(self, **kwargs): + out = {} + ref_latents = kwargs.get("reference_latents", None) + if ref_latents is not None: + out['ref_latents'] = list([1, 128, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 128]) + return out + class JoyImage(BaseModel): def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index a1bf047f8..39e973d36 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -884,6 +884,13 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): "selected_layer_index": selected_layer_index, } + if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys and '{}proj_out.weight'.format(key_prefix) in state_dict_keys and state_dict['{}txt_norm.weight'.format(key_prefix)].shape[0] == 2560 and state_dict['{}proj_out.weight'.format(key_prefix)].shape[0] == 128: # Mage-Flow (Qwen Image txt_norm/proj_out are 3584/64) + dit_config = {} + dit_config["image_model"] = "mage_flow" + dit_config["in_channels"] = 128 + dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.') + return dit_config + if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys: # Qwen Image dit_config = {} dit_config["image_model"] = "qwen_image" diff --git a/comfy/sd.py b/comfy/sd.py index e15e0a9fd..caf78222d 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -17,6 +17,7 @@ import comfy.ldm.wan.vae import comfy.ldm.wan.vae2_2 import comfy.ldm.hunyuan3d.vae import comfy.ldm.seedvr.vae +import comfy.ldm.mage_flow.vae import comfy.ldm.triposplat.vae import comfy.ldm.ace.vae.music_dcae_pipeline import comfy.ldm.cogvideo.vae @@ -60,6 +61,7 @@ import comfy.text_encoders.qwen_image import comfy.text_encoders.hunyuan_image import comfy.text_encoders.z_image import comfy.text_encoders.krea2 +import comfy.text_encoders.mage_flow import comfy.text_encoders.ideogram4 import comfy.text_encoders.ovis import comfy.text_encoders.kandinsky5 @@ -567,6 +569,17 @@ class VAE: self.upscale_index_formula = (4, 8, 8) self.process_input = lambda image: image * 2.0 - 1.0 self.crop_input = False + elif "student.dconv_encoder.proj_out.weight" in sd: # Mage-VAE (one-step diffusion codec, Flux2-anchored 128ch/16x latents) + sd = comfy.utils.state_dict_prefix_replace(sd, {"student.dconv_encoder.": "dconv_encoder.", "pipeline.": "decoder_model."}) + # Drop the unused Flux2-VAE anchor encoder carried in the checkpoint. + sd = {k: v for k, v in sd.items() if not k.startswith("decoder_model.y_embedder.encoder.") and not k.startswith("decoder_model.y_embedder.bottleneck.")} + self.first_stage_model = comfy.ldm.mage_flow.vae.MageVAE() + self.latent_channels = 128 + self.downscale_ratio = 16 + self.upscale_ratio = 16 + self.working_dtypes = [torch.bfloat16, torch.float32] + self.memory_used_encode = lambda shape, dtype: (400 * shape[2] * shape[3]) * model_management.dtype_size(dtype) + self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * 16 * 16) * model_management.dtype_size(dtype) elif "decoder.conv_in.weight" in sd: if sd['decoder.conv_in.weight'].shape[1] == 64: ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True} @@ -1379,6 +1392,7 @@ class CLIPType(Enum): BOOGU = 31 KREA2 = 32 JOYIMAGE = 33 + MAGE = 34 @@ -1713,6 +1727,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."}) clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data)) clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer + elif clip_type == CLIPType.MAGE and te_model == TEModel.QWEN3VL_4B: # Mage-Flow: full Qwen3-VL-4B, last hidden state, Qwen-Image-style templates. + clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."}) + clip_target.clip = comfy.text_encoders.mage_flow.te(**llama_detect(clip_data)) + clip_target.tokenizer = comfy.text_encoders.mage_flow.MageFlowTokenizer elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B: # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx. clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."}) clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data)) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index e7c8983aa..ca89850a5 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -27,6 +27,7 @@ import comfy.text_encoders.z_image import comfy.text_encoders.ideogram4 import comfy.text_encoders.boogu import comfy.text_encoders.krea2 +import comfy.text_encoders.mage_flow import comfy.text_encoders.joyimage import comfy.text_encoders.anima import comfy.text_encoders.ace15 @@ -1883,6 +1884,35 @@ class Krea2(supported_models_base.BASE): hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref)) return supported_models_base.ClipTarget(comfy.text_encoders.krea2.Krea2Tokenizer, comfy.text_encoders.krea2.te(**hunyuan_detect)) +class MageFlow(supported_models_base.BASE): + unet_config = { + "image_model": "mage_flow", + } + + sampling_settings = { + "multiplier": 1.0, + "shift": 6.0, + } + + memory_usage_factor = 6.5 + + unet_extra_config = {} + latent_format = latent_formats.Flux2 + + 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): + out = model_base.MageFlow(self, device=device) + return out + + def clip_target(self, state_dict={}): + pref = self.text_encoder_key_prefix[0] + hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref)) + return supported_models_base.ClipTarget(comfy.text_encoders.mage_flow.MageFlowTokenizer, comfy.text_encoders.mage_flow.te(**hunyuan_detect)) + class QwenImage(supported_models_base.BASE): unet_config = { "image_model": "qwen_image", @@ -2421,6 +2451,7 @@ models = [ ACEStep15, Omnigen2, Boogu, + MageFlow, QwenImage, JoyImage, Ideogram4, diff --git a/comfy/text_encoders/mage_flow.py b/comfy/text_encoders/mage_flow.py new file mode 100644 index 000000000..6542ad315 --- /dev/null +++ b/comfy/text_encoders/mage_flow.py @@ -0,0 +1,94 @@ +"""Mage-Flow text encoder: Qwen3-VL-4B, last hidden state (2560-dim). + +Mage-Flow conditions on the final hidden state of Qwen3-VL-4B with the leading +system + user-opening template tokens stripped (reference start_idx 34 for t2i, +64 for edit). The t2i template is identical to Qwen-Image's; the edit template +uses the same system prompt as Qwen-Image-Edit with "Image N: " reference +prefixes and no block. +""" + +import numbers + +import torch + +import comfy.text_encoders.qwen3vl +from comfy import sd1_clip + +MAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>" + +MAGE_T2I_TEMPLATE = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" +MAGE_EDIT_TEMPLATE = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + + +class MageFlowTokenizer(comfy.text_encoders.qwen3vl.Qwen3VLTokenizer): + def __init__(self, embedding_directory=None, tokenizer_data={}): + super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type="qwen3vl_4b") + self.llama_template = MAGE_T2I_TEMPLATE + self.llama_template_images = MAGE_EDIT_TEMPLATE + + def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=True, **kwargs): + image = kwargs.get("image", None) + if image is not None and len(images) == 0: + images = [image[i:i + 1] for i in range(image.shape[0])] + if llama_template is None: + if len(images) > 0: + # Training-time multi-reference body: "Image 1: Image 2: ...{instruction}" + prefix = "".join("Image {}: {}".format(j + 1, MAGE_VISION_BLOCK) for j in range(len(images))) + llama_template = self.llama_template_images.replace("{}", prefix + "{}", 1) + else: + llama_template = self.llama_template + # thinking=True: Mage templates end at "<|im_start|>assistant\n" with no block. + return super().tokenize_with_weights(text, return_word_ids=return_word_ids, llama_template=llama_template, images=images, prevent_empty_text=prevent_empty_text, thinking=thinking, **kwargs) + + +class MageFlowQwen3VLClipModel(comfy.text_encoders.qwen3vl.Qwen3VLClipModel): + def __init__(self, device="cpu", dtype=None, attention_mask=True, model_options={}, model_type="qwen3vl_4b"): + super().__init__(device=device, dtype=dtype, attention_mask=attention_mask, model_options=model_options, model_type=model_type) + # apply the final RMSNorm to the tapped last layer (HF last_hidden_state) + self.layer_norm_hidden_state = True + + +class MageFlowTEModel(sd1_clip.SD1ClipModel): + def __init__(self, device="cpu", dtype=None, model_options={}): + clip_model = lambda **kw: MageFlowQwen3VLClipModel(**kw, model_type="qwen3vl_4b") # noqa: E731 + super().__init__(device=device, dtype=dtype, name="qwen3vl_4b", clip_model=clip_model, model_options=model_options) + + def encode_token_weights(self, token_weight_pairs, template_end=-1): + # Strip the system + user-opening prefix (reference drop_idx: 34 t2i / 64 edit). + out, pooled, extra = super().encode_token_weights(token_weight_pairs) + tok_pairs = token_weight_pairs["qwen3vl_4b"][0] + count_im_start = 0 + if template_end == -1: + for i, v in enumerate(tok_pairs): + elem = v[0] + if not torch.is_tensor(elem): + if isinstance(elem, numbers.Integral): + if elem == 151644 and count_im_start < 2: # <|im_start|> + template_end = i + count_im_start += 1 + + if out.shape[1] > (template_end + 3): + if tok_pairs[template_end + 1][0] == 872: # "user" + if tok_pairs[template_end + 2][0] == 198: # "\n" + template_end += 3 + + out = out[:, template_end:] + + if "attention_mask" in extra: + extra["attention_mask"] = extra["attention_mask"][:, template_end:] + if extra["attention_mask"].sum() == torch.numel(extra["attention_mask"]): + extra.pop("attention_mask") # attention mask is useless if no masked elements + + return out, pooled, extra + + +def te(dtype_llama=None, llama_quantization_metadata=None): + class MageFlowTEModel_(MageFlowTEModel): + def __init__(self, device="cpu", dtype=None, model_options={}): + if dtype_llama is not None: + dtype = dtype_llama + if llama_quantization_metadata is not None: + model_options = model_options.copy() + model_options["quantization_metadata"] = llama_quantization_metadata + super().__init__(device=device, dtype=dtype, model_options=model_options) + return MageFlowTEModel_ diff --git a/comfy/text_encoders/qwen3vl.py b/comfy/text_encoders/qwen3vl.py index 7a329d2d6..2dd60d4e6 100644 --- a/comfy/text_encoders/qwen3vl.py +++ b/comfy/text_encoders/qwen3vl.py @@ -158,12 +158,12 @@ class Qwen3VLTokenizer(sd1_clip.SD1Tokenizer): self.llama_template = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" self.llama_template_images = "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" - def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=False, **kwargs): + def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=False, skip_template=False, **kwargs): image = kwargs.get("image", None) if image is not None and len(images) == 0: images = [image[i:i + 1] for i in range(image.shape[0])] - skip_template = text.startswith('<|im_start|>') + skip_template = skip_template or text.startswith('<|im_start|>') if prevent_empty_text and text == '': text = ' ' diff --git a/comfy_extras/nodes_mage.py b/comfy_extras/nodes_mage.py new file mode 100644 index 000000000..a3b0d394c --- /dev/null +++ b/comfy_extras/nodes_mage.py @@ -0,0 +1,103 @@ +from typing_extensions import override + +import comfy.utils +import node_helpers +import torch +import comfy.model_management +from comfy_api.latest import ComfyExtension, io + + +class TextEncodeMageFlowEdit(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncodeMageFlowEdit", + category="model/conditioning/mage", + description="Encode an edit instruction with one or more reference images for Mage-Flow-Edit. Reference latents are resized to the output resolution (width/height, or the first image's size when 0). Use the latent output for sampling so the sizes always match.", + inputs=[ + io.Clip.Input("clip"), + io.String.Input("prompt", multiline=True, dynamic_prompts=True), + io.String.Input("negative_prompt", multiline=True, dynamic_prompts=True, advanced=True), + io.Vae.Input("vae", optional=True), + io.Autogrow.Input( + "images", + template=io.Autogrow.TemplateNames( + io.Image.Input("image"), + names=[f"image_{i}" for i in range(1, 17)], + min=0, + ), + tooltip="Reference image(s) to edit. All references are resized to the output resolution before encoding.", + ), + io.Int.Input("width", default=0, min=0, max=8192, step=16, tooltip="Output width. 0 = use the first reference image's size."), + io.Int.Input("height", default=0, min=0, max=8192, step=16, tooltip="Output height. 0 = use the first reference image's size."), + io.Int.Input("batch_size", default=1, min=1, max=4096), + ], + outputs=[ + io.Conditioning.Output(display_name="positive"), + io.Conditioning.Output(display_name="negative"), + io.Latent.Output(display_name="latent"), + ], + ) + + @classmethod + def execute(cls, clip, prompt, negative_prompt="", vae=None, images: io.Autogrow.Type = None, width=0, height=0, batch_size=1) -> io.NodeOutput: + ref_latents = [] + images = images or {} + images = [images[name] for name in sorted(images, key=lambda n: int(n.rsplit("_", 1)[-1])) if images[name] is not None] + images_vl = [] + + # Output resolution: explicit width/height, else the primary reference's own size, floored to /16. + # Each dimension falls back independently so a 0 on one axis keeps an explicit value on the other. + if width == 0 or height == 0: + if len(images) > 0: + ref_h, ref_w = images[0].shape[1], images[0].shape[2] + else: + ref_h, ref_w = 1024, 1024 + height = height or ref_h + width = width or ref_w + width = max(16, (width // 16) * 16) + height = max(16, (height // 16) * 16) + + for image in images: + samples = image.movedim(-1, 1) + + # VL conditioning copy: cap the long edge at 384 (training preprocessing). + long_edge = max(samples.shape[3], samples.shape[2]) + if long_edge > 384: + scale_by = 384 / long_edge + s = comfy.utils.common_upscale(samples, max(1, round(samples.shape[3] * scale_by)), max(1, round(samples.shape[2] * scale_by)), "bicubic", "disabled") + images_vl.append(s.movedim(1, -1)) + else: + images_vl.append(image) + + if vae is not None: + # All references are resized to the output resolution before encoding, because Mage's RoPE aligns reference and target content by position + if samples.shape[3] != width or samples.shape[2] != height: + s = comfy.utils.common_upscale(samples, width, height, "bicubic", "disabled") + else: + s = samples + ref_latents.append(vae.encode(s.movedim(1, -1)[:, :, :, :3])) + + # Negative branch keeps the same reference images (VL tokens + ref latents), only the instruction differs. + positive = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=images_vl)) + negative = clip.encode_from_tokens_scheduled(clip.tokenize(negative_prompt if negative_prompt else " ", images=images_vl)) + + if len(ref_latents) > 0: + positive = node_helpers.conditioning_set_values(positive, {"reference_latents": ref_latents}, append=True) + negative = node_helpers.conditioning_set_values(negative, {"reference_latents": ref_latents}, append=True) + + latent = torch.zeros([batch_size, 128, height // 16, width // 16], device=comfy.model_management.intermediate_device()) + return io.NodeOutput(positive, negative, {"samples": latent}) + + +class MageExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + TextEncodeMageFlowEdit, + ] + + +async def comfy_entrypoint() -> MageExtension: + return MageExtension() diff --git a/nodes.py b/nodes.py index b03d6c603..243a55bf2 100644 --- a/nodes.py +++ b/nodes.py @@ -992,7 +992,7 @@ class CLIPLoader: @classmethod def INPUT_TYPES(s): return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ), - "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage"], ), + "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage", "mage"], ), }, "optional": { "device": (["default", "cpu"], {"advanced": True}), @@ -2462,6 +2462,7 @@ async def init_builtin_extra_nodes(): "nodes_seedvr.py", "nodes_context_windows.py", "nodes_qwen.py", + "nodes_mage.py", "nodes_joyimage.py", "nodes_boogu.py", "nodes_chroma_radiance.py",