Merge branch 'master' into update-comfyui-embedded-docs-0.5.12-20260918-171917

comfy/latent_formats.py

@@ -928,6 +928,59 @@ class HunyuanImage21(LatentFormat):
 
     latent_rgb_factors_bias = [0.0007, -0.0256, -0.0206]
 
+class QwenImage21(LatentFormat):
+    latent_channels = 64
+    latent_dimensions = 2
+    spacial_downscale_ratio = 16
+
+    latent_rgb_factors = [
+        [-0.0158, -0.0115, -0.0174], [ 0.0030,  0.0120,  0.0027], [ 0.0637,  0.0470, -0.0127], [ 0.0360,  0.0661, -0.0030],
+        [ 0.0159,  0.0181,  0.0082], [ 0.0132,  0.0326,  0.0169], [ 0.0191,  0.0261,  0.0136], [-0.0146, -0.0276, -0.0361],
+        [ 0.0187, -0.0024, -0.0072], [-0.1059, -0.0090,  0.0350], [-0.0195, -0.0226, -0.0138], [-0.0295,  0.0024, -0.0215],
+        [ 0.0191, -0.0393, -0.0001], [-0.0144, -0.0166, -0.0272], [ 0.0389,  0.0430,  0.0445], [-0.0153, -0.0336,  0.0031],
+        [ 0.0339,  0.0122,  0.0220], [-0.0136, -0.0078, -0.0120], [-0.0340, -0.0282, -0.0245], [-0.0133, -0.0176, -0.0133],
+        [ 0.0109, -0.0087,  0.0096], [-0.0010,  0.0044,  0.0016], [ 0.0301,  0.0053,  0.0361], [-0.0281, -0.0205, -0.0032],
+        [-0.0725,  0.0002,  0.0160], [-0.0036,  0.0158,  0.0807], [ 0.0087,  0.0040, -0.0053], [-0.0260,  0.0183, -0.0077],
+        [-0.0039, -0.0035, -0.0107], [-0.0026,  0.0172,  0.0237], [ 0.0088,  0.0078,  0.0078], [-0.0087, -0.0310, -0.0122],
+        [-0.0027,  0.0018,  0.0094], [-0.0064,  0.0292, -0.0256], [ 0.0594,  0.1049,  0.1180], [ 0.0103, -0.0103, -0.0026],
+        [-0.0091,  0.0025, -0.0015], [ 0.0178,  0.0243,  0.0292], [-0.0063, -0.0012,  0.0202], [ 0.0452,  0.0246,  0.0143],
+        [ 0.0149,  0.0270,  0.0052], [ 0.1484,  0.0801,  0.0804], [-0.0120,  0.0040,  0.0010], [ 0.0181,  0.0051, -0.0021],
+        [ 0.0132,  0.0050,  0.0019], [ 0.0291,  0.0020,  0.0092], [ 0.0066, -0.0410, -0.1314], [-0.1153, -0.0629, -0.0802],
+        [ 0.0258,  0.0378,  0.0298], [ 0.0375,  0.1139,  0.0468], [-0.0142, -0.0126, -0.0276], [ 0.0339,  0.0153,  0.0138],
+        [ 0.0346,  0.0211,  0.0267], [ 0.0369, -0.0431, -0.0993], [-0.0052, -0.0092,  0.0056], [-0.0279,  0.0410, -0.0357],
+        [ 0.0036,  0.0017, -0.0083], [-0.0441, -0.0367, -0.0454], [-0.0001, -0.0092, -0.0001], [-0.0222, -0.0183, -0.0051],
+        [ 0.0039,  0.0053, -0.0184], [-0.0094, -0.0075, -0.0143], [-0.0066, -0.0088, -0.0063], [ 0.0220,  0.0074,  0.0100],
+    ]
+    latent_rgb_factors_bias = [-0.1228, -0.1869, -0.3083]
+
+    def __init__(self):
+        self.latents_mean = torch.tensor([
+            0.5126, 0.7721, -0.0631, 1.3506, -0.7855, -2.1025, -0.3458, 1.3722,
+            1.8873, -1.7177, -0.6510, 0.2732, 0.7562, -0.6163, -1.0277, 3.8363,
+            2.0210, 0.0472, 0.9320, 2.0087, 2.4954, -0.1391, -1.4249, 1.8464,
+            -0.5236, 1.2826, 3.7046, -1.3035, 2.7286, -1.4518, -1.9036, -1.9955,
+            -0.0342, -1.0265, -0.7636, 3.0555, 0.0746, -3.0751, -0.1076, 1.7376,
+            -1.0914, -1.9435, -0.2784, -1.3680, 0.4809, -0.4433, 0.3764, 0.5729,
+            -2.0595, 1.0960, -1.3260, -2.0211, -5.0179, 0.5275, 4.0162, 1.8505,
+            0.3026, 1.9373, 1.4937, 0.2632, 0.5547, -1.7121, -0.1562, 0.0304,
+        ]).view(1, self.latent_channels, 1, 1)
+        self.latents_std = torch.tensor([
+            3.2001, 3.2936, 3.4321, 3.0091, 3.1061, 4.0379, 4.0705, 3.7910,
+            3.0785, 3.6500, 3.9308, 3.0904, 2.8778, 3.7675, 3.7320, 5.0756,
+            3.2864, 4.0397, 3.1317, 4.0443, 2.9249, 3.9454, 3.0988, 4.2489,
+            3.4896, 3.8513, 3.9323, 3.4719, 3.7498, 4.2830, 3.5694, 4.2467,
+            3.9037, 3.2947, 5.0770, 3.5075, 3.2700, 3.4767, 2.8063, 5.1125,
+            3.5327, 4.7833, 3.1286, 4.1819, 3.8527, 3.8312, 3.5605, 4.3875,
+            3.9624, 4.0168, 3.5643, 4.0550, 5.5614, 4.2963, 4.4080, 3.4959,
+            3.8747, 3.7608, 3.5735, 3.1490, 3.7662, 3.6746, 3.4563, 3.8161,
+        ]).view(1, self.latent_channels, 1, 1)
+
+    def process_in(self, latent):
+        return (latent - self.latents_mean.to(latent.device, latent.dtype)) / self.latents_std.to(latent.device, latent.dtype)
+
+    def process_out(self, latent):
+        return latent * self.latents_std.to(latent.device, latent.dtype) + self.latents_mean.to(latent.device, latent.dtype)
+
 class HunyuanImage21Refiner(LatentFormat):
     latent_channels = 64
     latent_dimensions = 3

comfy/ldm/qwen_image21/model.py

@@ -0,0 +1,353 @@
+# https://github.com/huggingface/diffusers (Apache 2.0) Qwen-Image 2.1
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import comfy.model_management
+import comfy.ops
+import comfy.patcher_extension
+import comfy.quant_ops
+import comfy.rmsnorm
+from comfy.ldm.flux.layers import EmbedND, timestep_embedding
+from comfy.ldm.flux.math import apply_rope1
+from comfy.ldm.lightricks.model import TimestepEmbedding
+from comfy.ldm.modules.attention import optimized_attention
+from comfy.ldm.wan.model_animate2 import PoseBranchCache
+
+
+class ZeroCenteredRMSNorm(nn.Module):
+    # stored weight is scale - 1, applied in fp32
+    def __init__(self, dim, eps=1e-6, dtype=None, device=None):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
+        self.eps = eps
+
+    def forward(self, x):
+        w = comfy.model_management.cast_to(self.weight, dtype=torch.float32, device=x.device) + 1.0
+        return comfy.rmsnorm.rms_norm(x.float(), w, self.eps).to(x.dtype)
+
+
+class TextProjection(nn.Module):
+    def __init__(self, in_dim, hidden_size, eps=1e-6, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.text_norm = ZeroCenteredRMSNorm(in_dim, eps=eps, dtype=dtype, device=device)
+        self.in_layer = operations.Linear(in_dim, hidden_size, bias=False, dtype=dtype, device=device)
+        self.out_layer = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x):
+        return self.out_layer(F.gelu(self.in_layer(self.text_norm(x)), approximate="tanh"))
+
+
+class TimestepProjEmbeddings(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, sample_proj_bias=False, dtype=dtype, device=device, operations=operations)
+
+    def forward(self, timestep, dtype):
+        return self.timestep_embedder(timestep_embedding(timestep.float(), 256).to(dtype))
+
+
+class SwiGLUFeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, fused=True, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.fused = fused
+        if fused:
+            # [gate; up] in one GEMM, the SiLU-gate folded into the down projection's input quantizer
+            self.gate_up = operations.Linear(dim, 2 * hidden_dim, bias=False, dtype=dtype, device=device)
+        else:
+            self.proj = operations.Linear(dim, hidden_dim, bias=False, dtype=dtype, device=device)
+            self.gate_layer = operations.Linear(dim, hidden_dim, bias=False, dtype=dtype, device=device)
+        self.out = operations.Linear(hidden_dim, dim, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x):
+        if self.fused:
+            return comfy.ops.linear_input_act(self.out, self.gate_up(x), "swiglu")
+        return self.out(F.silu(self.gate_layer(x)) * self.proj(x))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads, dim_head, eps=1e-6, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.heads = heads
+        inner_dim = heads * dim_head
+        self.to_q = operations.Linear(dim, inner_dim, bias=False, dtype=dtype, device=device)
+        self.to_k = operations.Linear(dim, inner_dim, bias=False, dtype=dtype, device=device)
+        self.to_v = operations.Linear(dim, inner_dim, bias=False, dtype=dtype, device=device)
+        self.to_out = nn.ModuleList([operations.Linear(inner_dim, dim, bias=False, dtype=dtype, device=device)])
+        self.norm_q = operations.RMSNorm(dim_head, eps=eps, dtype=dtype, device=device)
+        self.norm_k = operations.RMSNorm(dim_head, eps=eps, dtype=dtype, device=device)
+
+    def forward(self, x, pe, attn_fn, prefix_len, transformer_options={}):
+        # (B, N, H, D) throughout: no transposes, the rope table is laid out to match
+        B, N, _ = x.shape
+        q = self.to_q(x).view(B, N, self.heads, -1)
+        k = self.to_k(x).view(B, N, self.heads, -1)
+        v = self.to_v(x).view(B, N, self.heads, -1)
+        patches = transformer_options.get("patches", {}).get("attn1_patch", [])
+        if comfy.model_management.in_training or patches:
+            q, k = self.norm_q(q), self.norm_k(k)
+            if patches:
+                # patches see the Qwen-Image convention: (B, H, N, D) before rope, rope table (1, 1, N, ...), target image rows in img_slice
+                q, k, v, pe = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), pe.transpose(1, 2)
+                extra_options = {**transformer_options, "img_slice": [prefix_len, N]}
+                for p in patches:
+                    out = p(q, k, v, pe=pe, attn_mask=None, extra_options=extra_options)
+                    q, k, v, pe = out.get("q", q), out.get("k", k), out.get("v", v), out.get("pe", pe)
+                q, k, v, pe = q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), pe.transpose(1, 2).contiguous()
+            q = apply_rope1(q, pe)
+            k = apply_rope1(k, pe)
+        else:
+            q_scale, _, q_stream = comfy.ops.cast_bias_weight(self.norm_q, q, offloadable=True)
+            k_scale, _, k_stream = comfy.ops.cast_bias_weight(self.norm_k, k, offloadable=True)
+            q, k = comfy.quant_ops.ck.rms_rope(q, k, pe, q_scale, k_scale, self.norm_q.eps)
+            comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_stream)
+            comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_stream)
+        return self.to_out[0](attn_fn(q, k, v, self.heads))
+
+
+def _split_rows(p):
+    # shared modulation rows: (t = 0 row for text and references, sampled-t rows for the target)
+    return p[-1:].unsqueeze(1), p[:-1].unsqueeze(1)
+
+
+def _modulated_norm(norm, x, scale, prefix_len, zero):
+    # LayerNorm * (1 + scale), fused over every row with the target scale; the prefix rows are then redone with the t = 0 scale
+    s_prefix, s_target = scale
+    if comfy.model_management.in_training:
+        out = norm(x)
+        return torch.cat([out[:, :prefix_len] * (1 + s_prefix), out[:, prefix_len:] * (1 + s_target)], dim=1)
+    out = comfy.quant_ops.ck.adaln(x, s_target, zero, norm.eps)
+    if prefix_len:
+        out[:, :prefix_len] = comfy.quant_ops.ck.adaln(x[:, :prefix_len], s_prefix, zero, norm.eps)
+    return out
+
+
+def _gated_residual(x, y, gate, prefix_len):
+    g_prefix, g_target = gate
+    out = torch.addcmul(x, y, g_target)
+    if prefix_len:
+        out[:, :prefix_len] = torch.addcmul(x[:, :prefix_len], y[:, :prefix_len], g_prefix)
+    return out
+
+
+class QwenImage21TransformerBlock(nn.Module):
+    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=3, eps=1e-6, fused_mlp=True, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.img_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+        self.attn = Attention(dim, num_attention_heads, attention_head_dim, eps=eps, dtype=dtype, device=device, operations=operations)
+        self.img_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+        self.img_mlp = SwiGLUFeedForward(dim, dim * mlp_ratio, fused=fused_mlp, dtype=dtype, device=device, operations=operations)
+
+    def forward(self, x, mod, pe, attn_fn, prefix_len, transformer_options={}):
+        scale1, gate1, scale2, gate2, zero = mod
+        x = _gated_residual(x, self.attn(_modulated_norm(self.img_norm1, x, scale1, prefix_len, zero), pe, attn_fn, prefix_len, transformer_options), gate1, prefix_len)
+        x = _gated_residual(x, self.img_mlp(_modulated_norm(self.img_norm2, x, scale2, prefix_len, zero)), gate2, prefix_len)
+        if x.dtype == torch.float16:
+            x = x.clip(-65504, 65504)
+        return x
+
+
+class LastLayer(nn.Module):
+    # scale only, no shift
+    def __init__(self, dim, eps=1e-6, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.linear = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.norm = operations.LayerNorm(dim, eps, elementwise_affine=False, dtype=dtype, device=device)
+
+    def forward(self, x, temb):
+        scale = self.linear(F.silu(temb)).unsqueeze(1)
+        if comfy.model_management.in_training:
+            return self.norm(x) * (1 + scale)
+        return comfy.quant_ops.ck.adaln(x, scale, torch.zeros_like(scale[:1]), self.norm.eps)
+
+
+def block_causal_attention(segments, transformer_options={}, cache=None, block_index=0, prefix_len=0):
+    # segments: (start, end, mask); text segments get a causal mask, image blocks attend to everything before their end
+    def attn(q, k, v, heads):
+        if cache is not None:
+            # K and V stacked on dim 1 so batch stays first and quantized rows are per token and head
+            cache.put(block_index, torch.stack([k[:, :prefix_len], v[:, :prefix_len]], dim=1))
+        outs = [optimized_attention(q[:, start:end].flatten(2), k[:, :end].flatten(2), v[:, :end].flatten(2), heads, mask=mask, transformer_options=transformer_options)
+                for start, end, mask in segments]
+        return torch.cat(outs, dim=1) if len(outs) > 1 else outs[0]
+    return attn
+
+
+def prefix_cached_attention(prefix_k, prefix_v, transformer_options={}):
+    # target-only queries: block-causal reduces to full attention over [cached prefix, target]
+    def attn(q, k, v, heads):
+        return optimized_attention(q.flatten(2), torch.cat([prefix_k, k], dim=1).flatten(2), torch.cat([prefix_v, v], dim=1).flatten(2), heads, transformer_options=transformer_options)
+    return attn
+
+
+def prefix_cache_key(x, context, refs, slots):
+    # one fp32 tensor per batch row: lengths and slots, then the prompt embedding and reference latents
+    # the target shape is part of it because reference rope ids are centred on the target
+    header = [context.shape[1], len(refs)] + list(x.shape[-2:]) + list(slots) + [s for r in refs for s in r.shape[-2:]]
+    header = torch.tensor(header, dtype=torch.float32, device=context.device).expand(context.shape[0], -1)
+    return torch.cat([header, context.float().flatten(1)] + [r.float().flatten(1) for r in refs], dim=1)
+
+
+class QwenImage21Transformer2DModel(nn.Module):
+    def __init__(
+        self,
+        in_channels=64,
+        out_channels=64,
+        num_layers=32,
+        attention_head_dim=128,
+        num_attention_heads=32,
+        context_in_dim=4096,
+        mlp_ratio=3,
+        axes_dims_rope=(16, 56, 56),
+        eps=1e-6,
+        fused_mlp=True,
+        image_model=None,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.dtype = dtype
+        self.out_channels = out_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 = TimestepProjEmbeddings(self.inner_dim, dtype=dtype, device=device, operations=operations)
+        self.txt_in = TextProjection(context_in_dim, self.inner_dim, eps=eps, dtype=dtype, device=device, operations=operations)
+        self.img_in = operations.Linear(in_channels, self.inner_dim, bias=False, dtype=dtype, device=device)
+
+        # one modulation shared by every block
+        self.modulation = nn.Sequential(nn.SiLU(), operations.Linear(self.inner_dim, 4 * self.inner_dim, bias=False, dtype=dtype, device=device))
+
+        self.transformer_blocks = nn.ModuleList([
+            QwenImage21TransformerBlock(self.inner_dim, num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, eps=eps, fused_mlp=fused_mlp, dtype=dtype, device=device, operations=operations)
+            for _ in range(num_layers)
+        ])
+
+        self.norm_out = LastLayer(self.inner_dim, eps=eps, dtype=dtype, device=device, operations=operations)
+        self.proj_out = operations.Linear(self.inner_dim, out_channels, bias=False, dtype=dtype, device=device)
+
+        # text + reference K/V are step-independent (t = 0 modulation, causal prefix), cached for one sampling run
+        self.prefix_cache = None
+        self.prefix_cache_enabled = False
+
+    def reset_prefix_cache(self, enabled):
+        if self.prefix_cache is not None:
+            self.prefix_cache.free()
+        self.prefix_cache = None
+        self.prefix_cache_enabled = enabled
+
+    def select_prefix_cache(self, key, cache_bytes, device, options):
+        # returns (cache with the slot to read or fill selected, whether the slot is filled), or (None, False) to recompute
+        if options.get("device") == "off":
+            return None, False
+        cache = self.prefix_cache
+        if cache is not None and cache.select(key, create=False):
+            return cache, cache.filled(len(self.transformer_blocks))
+        dtype = options.get("dtype", "default")
+        cache_bytes //= {"int8": 2, "int4": 4}.get(dtype, 1)
+        if cache is None:
+            store = options.get("device", "auto")
+            if store == "auto":
+                # spare VRAM first, then pinned RAM prefetched behind compute; room for cond and uncond
+                store = next((d for d in (device, torch.device("cpu")) if comfy.model_management.get_free_memory(d) > 4 * cache_bytes), None)
+                if store is None:
+                    return None, False
+            else:
+                store = device if store == "gpu" else torch.device("cpu")
+            cache = self.prefix_cache = PoseBranchCache(store_device=store, dtype=dtype)
+        if comfy.model_management.get_free_memory(cache.store_device) < 2 * cache_bytes:
+            # no room for this slot: recompute rather than evict the other cond's slot every step
+            return None, False
+        cache.select(key)
+        return cache, False
+
+    def build_sequence(self, x, context, ref_latents, image_slots):
+        # text with each reference image spliced in at its slot, target image last
+        txt = self.txt_in(context)
+        slots = (image_slots + [txt.shape[1]] * len(ref_latents))[:len(ref_latents)]
+        bounds = [0] + slots + [txt.shape[1]]
+
+        parts, ids, segments = [], [], []
+        pos, length = 0, 0
+        for (start, end), img in zip(zip(bounds[:-1], bounds[1:]), ref_latents + [x]):
+            n = end - start
+            if n > 0:
+                parts.append(txt[:, start:end])
+                ids.append(torch.arange(pos, pos + n, device=x.device, dtype=torch.float32).unsqueeze(1).expand(n, 3))
+                segments.append((length, length + n, torch.ones((n, length + n), dtype=torch.bool, device=x.device).tril(length)))
+                pos += n
+                length += n
+            h, w = img.shape[-2:]
+            parts.append(self.img_in(img.flatten(2).transpose(1, 2)))
+            # half a token where a reference grid has the other parity, so it centres on the target
+            hh = torch.arange(h, device=x.device, dtype=torch.float32) - (h - h // 2) + 0.5 * (h % 2 - x.shape[-2] % 2)
+            ww = torch.arange(w, device=x.device, dtype=torch.float32) - (w - w // 2) + 0.5 * (w % 2 - x.shape[-1] % 2)
+            ids.append(torch.stack([torch.full((h, w), pos, device=x.device, dtype=torch.float32), hh[:, None].expand(h, w), ww[None, :].expand(h, w)], dim=-1).flatten(0, 1))
+            segments.append((length, length + h * w, None))
+            pos += max(h, w)
+            length += h * w
+
+        # (1, N, 1, ...): the layout the fused rms_rope wants for (B, N, H, D) queries
+        pe = self.pe_embedder(torch.cat(ids, dim=0).unsqueeze(0)).transpose(1, 2).contiguous()
+        return torch.cat(parts, dim=1), pe, segments
+
+    def forward(self, x, timestep, context, ref_latents=None, image_slots=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, ref_latents, image_slots, transformer_options, **kwargs)
+
+    def _forward(self, x, timesteps, context, ref_latents=None, image_slots=None, transformer_options={}, **kwargs):
+        B, C, H, W = x.shape
+        dtype = x.dtype
+        ref_latents = list(ref_latents or [])
+        image_slots = list(image_slots or [])
+
+        hidden_states, pe, segments = self.build_sequence(x, context, ref_latents, image_slots)
+        prefix_len = hidden_states.shape[1] - H * W
+        patches = transformer_options.get("patches", {})
+        for p in patches.get("post_input", []):
+            out = p({"img": hidden_states, "pe": pe, "transformer_options": transformer_options})
+            hidden_states, pe = out["img"], out.get("pe", pe)
+
+        # pipeline rounds t*1000 and t to the compute dtype; text and reference tokens modulate from t = 0
+        t = ((timesteps * 1000).to(dtype) / 1000).to(dtype)
+        temb = self.time_text_embed(torch.cat([t, t.new_zeros(1)]), dtype)
+        scale1, gate1, scale2, gate2 = self.modulation(temb).chunk(4, dim=-1)
+        mod = (_split_rows(scale1), _split_rows(gate1.tanh()), _split_rows(scale2), _split_rows(gate2.tanh()), torch.zeros_like(scale1[:1, None]))
+
+        blocks_replace = transformer_options.get("patches_replace", {}).get("dit", {})
+        cache, cached = None, False
+        # a cached step runs target rows only, so anything hooked into a block would see a different sequence from step 2
+        hooked = blocks_replace or patches.get("post_input") or patches.get("single_block") or patches.get("attn1_patch")
+        if self.prefix_cache_enabled and prefix_len > 0 and not hooked:
+            key = prefix_cache_key(x, context, ref_latents, image_slots)
+            cache_bytes = 2 * len(self.transformer_blocks) * B * prefix_len * self.inner_dim * hidden_states.element_size()
+            cache, cached = self.select_prefix_cache(key, cache_bytes, x.device, transformer_options.get("qwen_image21_cache", {}))
+        if cached:
+            hidden_states, pe, prefix_len = hidden_states[:, prefix_len:], pe[:, prefix_len:], 0
+
+        transformer_options["total_blocks"] = len(self.transformer_blocks)
+        transformer_options["block_type"] = "single"
+        for i, block in enumerate(self.transformer_blocks):
+            transformer_options["block_index"] = i
+            if cached:
+                prefix_k, prefix_v = cache.take(i, x.device, dtype, B).unbind(1)
+                cache.prefetch(i + 1, x.device, dtype)  # queue the next block before the compute it should overlap
+                attn_fn = prefix_cached_attention(prefix_k, prefix_v, transformer_options)
+            else:
+                attn_fn = block_causal_attention(segments, transformer_options, cache, i, prefix_len)
+            if ("single_block", i) in blocks_replace:
+                def block_wrap(args):
+                    return {"img": block(args["img"], mod, args["pe"], attn_fn, prefix_len, args["transformer_options"])}
+                hidden_states = blocks_replace[("single_block", i)]({"img": hidden_states, "vec": temb, "pe": pe, "transformer_options": transformer_options}, {"original_block": block_wrap})["img"]
+            else:
+                hidden_states = block(hidden_states, mod, pe, attn_fn, prefix_len, transformer_options)
+            for p in patches.get("single_block", []):
+                hidden_states = p({"img": hidden_states, "x": x, "block_index": i, "transformer_options": transformer_options})["img"]
+
+        hidden_states = self.norm_out(hidden_states[:, prefix_len:], temb[:-1])
+        hidden_states = self.proj_out(hidden_states)
+        return hidden_states.transpose(1, 2).reshape(B, self.out_channels, H, W)

comfy/ldm/wan/model_animate2.py

@@ -116,7 +116,8 @@ class PoseBranchCache:
     Caching the block input rather than its K/V halves the memory; reprojecting K/V on read
     costs ~4% of re-running the block. One slot per distinct pose sequence, so under
     context windows each window keeps its own; least recently used slots are evicted when
-    the store device runs low on memory. Created and freed by WanAnimate2Cache.
+    the store device runs low on memory. Created and freed by WanAnimate2Cache; Qwen-Image 2.1
+    also stores its prefix K/V in it.
     """
 
     CONVROT_GROUPSIZE = 256
@@ -129,21 +130,22 @@ class PoseBranchCache:
         self._pending = {}
         self._staging = {}
 
-    def select(self, pose_latents):
+    def select(self, k, create=True):
         # select runs at a forward boundary: an interrupted forward can leave copies in flight that a different slot's forward would then mistake for its own
         if self._pending:
             for t, stream in self._pending.values():
                 if stream is not None:
                     stream.synchronize()
             self._pending = {}
-        # keyed on batch element 0, so a cond batch size change mid-run stays valid
-        k = pose_latents[:1]
         for s in self.slots:
             if s["key"].shape == k.shape and torch.equal(s["key"], k.to(s["key"].device)):
                 self.slots.remove(s)
                 self.slots.append(s)
                 self.slot = s
-                return
+                return True
+        self.slot = None
+        if not create:
+            return False
         # cache what fits: a filled slot is the size estimate for the next one, and least recently used slots make room when the store device runs low
         est = max((self._slot_bytes(s) for s in self.slots), default=0) * 1.5
         while self.slots and comfy.model_management.get_free_memory(self.store_device) < est:
@@ -169,8 +171,7 @@ class PoseBranchCache:
     def filled(self, num_blocks):
         return self.slot is not None and len(self.slot["blocks"]) == num_blocks
 
-    def put(self, i, x_pose):
-        t = x_pose[:1]
+    def put(self, i, t):
         params = None
         if self.dtype in ("int8", "int4"):
             # convrot is what lets low-bit survive the ~125x per-channel outliers here, and over a [tokens, dim] view per-row scale means per-token. The kernels want 2D and a power-of-4 group that divides dim.
@@ -294,7 +295,7 @@ class WanAnimate2Model(WanModel):
 
         cache = transformer_options.get("animate2_cache", None) if apply_pose else None
         if cache is not None:
-            cache.select(pose_latents)
+            cache.select(pose_latents[:1])  # keyed on batch element 0, so a cond batch size change mid-run stays valid
         cached = cache is not None and cache.filled(len(self.blocks))
 
         x_pose = None
@@ -342,7 +343,7 @@ class WanAnimate2Model(WanModel):
             # pose-only prepass, to avoid inflating dynamic VRAM calibration when using multiple context windows
             for i, block in enumerate(self.blocks):
                 transformer_options["block_index"] = i
-                cache.put(i, x_pose)
+                cache.put(i, x_pose[:1])
                 x_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)[0]
             x_pose = None
             cached = True
@@ -365,7 +366,7 @@ class WanAnimate2Model(WanModel):
                 del x_pose_in
             else:
                 if cache is not None:
-                    cache.put(i, x_pose)
+                    cache.put(i, x_pose[:1])
                 # runs even under a block replace: its state has to reach block i+1
                 x_pose, k_pose, v_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)
             if v_pose is not None and pose_strength != 1.0:

comfy/ldm/wan/vae2_2.py

@@ -1,6 +1,7 @@
 # original version: https://github.com/Wan-Video/Wan2.2/blob/main/wan/modules/vae2_2.py
 # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
 
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -15,7 +16,7 @@ CACHE_T = 2
 
 class Resample(nn.Module):
 
-    def __init__(self, dim, mode):
+    def __init__(self, dim, mode, temporal_kernel=3):
         assert mode in (
             "none",
             "upsample2d",
@@ -40,7 +41,7 @@ class Resample(nn.Module):
                 # ops.Conv2d(dim, dim//2, 3, padding=1)
             )
             self.time_conv = CausalConv3d(
-                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+                dim, dim * 2, (temporal_kernel, 1, 1), padding=(temporal_kernel // 2, 0, 0))
         elif mode == "downsample2d":
             self.resample = nn.Sequential(
                 nn.ZeroPad2d((0, 1, 0, 1)),
@@ -50,7 +51,7 @@ class Resample(nn.Module):
                 nn.ZeroPad2d((0, 1, 0, 1)),
                 ops.Conv2d(dim, dim, 3, stride=(2, 2)))
             self.time_conv = CausalConv3d(
-                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+                dim, dim, (temporal_kernel, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
         else:
             self.resample = nn.Identity()
 
@@ -96,7 +97,10 @@ class Resample(nn.Module):
                     x = x.reshape(b, c, t * 2, h, w)
         t = x.shape[2]
         x = rearrange(x, "b c t h w -> (b t) c h w")
-        x = self.resample(x)
+        if feat_cache is None and self.mode in ("upsample2d", "upsample3d"):
+            x = strip_apply(self.resample, x, scale=2)
+        else:
+            x = self.resample(x)
         x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
 
         if self.mode == "downsample3d":
@@ -114,9 +118,38 @@ class Resample(nn.Module):
         return x
 
 
+STRIP_ELEMS = 2 ** 24
+
+
+def strip_apply(fn, x, scale=1, halo=1, out=None):
+    # strips of rows bound cudnn's conv workspace, a halo row per 3x3 conv keeps them exact
+    n = -(-x.numel() * scale * scale // STRIP_ELEMS)
+    if n <= 1 and out is None:
+        return fn(x)
+    add = out is not None
+    size = x.shape[-2]
+    step = -(-size // n)
+    for a in range(0, size, step):
+        b = min(size, a + step)
+        lo = max(0, a - halo)
+        y = fn(x.narrow(-2, lo, min(size, b + halo) - lo)).narrow(-2, (a - lo) * scale, (b - a) * scale)
+        if out is None:
+            out = y.new_empty(*y.shape[:-2], size * scale, y.shape[-1])
+        dst = out.narrow(-2, a * scale, (b - a) * scale)
+        if add:
+            dst.add_(y)
+        else:
+            dst.copy_(y)
+    return out
+
+
+def conv3x3(in_dim, out_dim, temporal_kernel=3):
+    return CausalConv3d(in_dim, out_dim, (temporal_kernel, 3, 3), padding=(temporal_kernel // 2, 1, 1))
+
+
 class ResidualBlock(nn.Module):
 
-    def __init__(self, in_dim, out_dim, dropout=0.0):
+    def __init__(self, in_dim, out_dim, dropout=0.0, temporal_kernel=3):
         super().__init__()
         self.in_dim = in_dim
         self.out_dim = out_dim
@@ -125,17 +158,20 @@ class ResidualBlock(nn.Module):
         self.residual = nn.Sequential(
             RMS_norm(in_dim, images=False),
             nn.SiLU(),
-            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            conv3x3(in_dim, out_dim, temporal_kernel),
             RMS_norm(out_dim, images=False),
             nn.SiLU(),
             nn.Dropout(dropout),
-            CausalConv3d(out_dim, out_dim, 3, padding=1),
+            conv3x3(out_dim, out_dim, temporal_kernel),
         )
         self.shortcut = (
             CausalConv3d(in_dim, out_dim, 1)
             if in_dim != out_dim else nn.Identity())
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is None:
+            # single image: the whole block runs in strips so its intermediates never exist at full size
+            return strip_apply(lambda s: self.residual(s).add_(self.shortcut(s)), x, halo=2)
         old_x = x
         for layer in self.residual:
             if isinstance(layer, CausalConv3d) and feat_cache is not None:
@@ -302,7 +338,8 @@ class Down_ResidualBlock(nn.Module):
                  dropout,
                  mult,
                  temperal_downsample=False,
-                 down_flag=False):
+                 down_flag=False,
+                 temporal_kernel=3):
         super().__init__()
 
         # Shortcut path with downsample
@@ -316,13 +353,13 @@ class Down_ResidualBlock(nn.Module):
         # Main path with residual blocks and downsample
         downsamples = []
         for _ in range(mult):
-            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout, temporal_kernel=temporal_kernel))
             in_dim = out_dim
 
         # Add the final downsample block
         if down_flag:
             mode = "downsample3d" if temperal_downsample else "downsample2d"
-            downsamples.append(Resample(out_dim, mode=mode))
+            downsamples.append(Resample(out_dim, mode=mode, temporal_kernel=temporal_kernel))
 
         self.downsamples = nn.Sequential(*downsamples)
 
@@ -342,7 +379,8 @@ class Up_ResidualBlock(nn.Module):
                  dropout,
                  mult,
                  temperal_upsample=False,
-                 up_flag=False):
+                 up_flag=False,
+                 temporal_kernel=3):
         super().__init__()
         # Shortcut path with upsample
         if up_flag:
@@ -358,13 +396,13 @@ class Up_ResidualBlock(nn.Module):
         # Main path with residual blocks and upsample
         upsamples = []
         for _ in range(mult):
-            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout, temporal_kernel=temporal_kernel))
             in_dim = out_dim
 
         # Add the final upsample block
         if up_flag:
             mode = "upsample3d" if temperal_upsample else "upsample2d"
-            upsamples.append(Resample(out_dim, mode=mode))
+            upsamples.append(Resample(out_dim, mode=mode, temporal_kernel=temporal_kernel))
 
         self.upsamples = nn.Sequential(*upsamples)
 
@@ -373,6 +411,8 @@ class Up_ResidualBlock(nn.Module):
         for module in self.upsamples:
             x_main = module(x_main, feat_cache, feat_idx)
         if self.avg_shortcut is not None:
+            if feat_cache is None:
+                return strip_apply(lambda s: self.avg_shortcut(s, first_chunk), x, scale=self.avg_shortcut.factor_s, halo=0, out=x_main)
             x_shortcut = self.avg_shortcut(x, first_chunk)
             return x_main + x_shortcut
         else:
@@ -390,6 +430,8 @@ class Encoder3d(nn.Module):
         attn_scales=[],
         temperal_downsample=[True, True, False],
         dropout=0.0,
+        in_channels=12,
+        temporal_kernel=3,
     ):
         super().__init__()
         self.dim = dim
@@ -404,7 +446,7 @@ class Encoder3d(nn.Module):
         scale = 1.0
 
         # init block
-        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+        self.conv1 = conv3x3(in_channels, dims[0], temporal_kernel)
 
         # downsample blocks
         downsamples = []
@@ -420,22 +462,23 @@ class Encoder3d(nn.Module):
                     mult=num_res_blocks,
                     temperal_downsample=t_down_flag,
                     down_flag=i != len(dim_mult) - 1,
+                    temporal_kernel=temporal_kernel,
                 ))
             scale /= 2.0
         self.downsamples = nn.Sequential(*downsamples)
 
         # middle blocks
         self.middle = nn.Sequential(
-            ResidualBlock(out_dim, out_dim, dropout),
+            ResidualBlock(out_dim, out_dim, dropout, temporal_kernel=temporal_kernel),
             AttentionBlock(out_dim),
-            ResidualBlock(out_dim, out_dim, dropout),
+            ResidualBlock(out_dim, out_dim, dropout, temporal_kernel=temporal_kernel),
         )
 
         # # output blocks
         self.head = nn.Sequential(
             RMS_norm(out_dim, images=False),
             nn.SiLU(),
-            CausalConv3d(out_dim, z_dim, 3, padding=1),
+            conv3x3(out_dim, z_dim, temporal_kernel),
         )
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
@@ -506,6 +549,8 @@ class Decoder3d(nn.Module):
         attn_scales=[],
         temperal_upsample=[False, True, True],
         dropout=0.0,
+        out_channels=12,
+        temporal_kernel=3,
     ):
         super().__init__()
         self.dim = dim
@@ -518,13 +563,13 @@ class Decoder3d(nn.Module):
         # dimensions
         dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
         # init block
-        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+        self.conv1 = conv3x3(z_dim, dims[0], temporal_kernel)
 
         # middle blocks
         self.middle = nn.Sequential(
-            ResidualBlock(dims[0], dims[0], dropout),
+            ResidualBlock(dims[0], dims[0], dropout, temporal_kernel=temporal_kernel),
             AttentionBlock(dims[0]),
-            ResidualBlock(dims[0], dims[0], dropout),
+            ResidualBlock(dims[0], dims[0], dropout, temporal_kernel=temporal_kernel),
         )
 
         # upsample blocks
@@ -540,6 +585,7 @@ class Decoder3d(nn.Module):
                     mult=num_res_blocks + 1,
                     temperal_upsample=t_up_flag,
                     up_flag=i != len(dim_mult) - 1,
+                    temporal_kernel=temporal_kernel,
                 ))
         self.upsamples = nn.Sequential(*upsamples)
 
@@ -547,7 +593,7 @@ class Decoder3d(nn.Module):
         self.head = nn.Sequential(
             RMS_norm(out_dim, images=False),
             nn.SiLU(),
-            CausalConv3d(out_dim, 12, 3, padding=1),
+            conv3x3(out_dim, out_channels, temporal_kernel),
         )
 
     def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
@@ -580,7 +626,7 @@ class Decoder3d(nn.Module):
             if feat_cache is not None:
                 x = layer(x, feat_cache, feat_idx, first_chunk)
             else:
-                x = layer(x)
+                x = layer(x, first_chunk=first_chunk)
 
         ## head
         for layer in self.head:
@@ -624,6 +670,9 @@ class WanVAE(nn.Module):
         attn_scales=[],
         temperal_downsample=[True, True, False],
         dropout=0.0,
+        image_channels=3,
+        patch_size=2,
+        temporal_kernel=3,
     ):
         super().__init__()
         self.dim = dim
@@ -633,6 +682,7 @@ class WanVAE(nn.Module):
         self.attn_scales = attn_scales
         self.temperal_downsample = temperal_downsample
         self.temperal_upsample = temperal_downsample[::-1]
+        self.patch_size = patch_size
 
         # modules
         self.encoder = Encoder3d(
@@ -643,6 +693,8 @@ class WanVAE(nn.Module):
             attn_scales,
             self.temperal_downsample,
             dropout,
+            in_channels=image_channels * patch_size * patch_size,
+            temporal_kernel=temporal_kernel,
         )
         self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
         self.conv2 = CausalConv3d(z_dim, z_dim, 1)
@@ -654,12 +706,18 @@ class WanVAE(nn.Module):
             attn_scales,
             self.temperal_upsample,
             dropout,
+            out_channels=image_channels * patch_size * patch_size,
+            temporal_kernel=temporal_kernel,
         )
 
     def encode(self, x):
+        if x.ndim == 4:
+            # single image: no temporal cache, which would keep every conv input alive for the whole pass
+            x = patchify(x.unsqueeze(2), patch_size=self.patch_size)
+            return self.conv1(self.encoder(x)).chunk(2, dim=1)[0].squeeze(2)
         conv_idx = [0]
         feat_map = [None] * count_conv3d(self.encoder)
-        x = patchify(x, patch_size=2)
+        x = patchify(x, patch_size=self.patch_size)
         t = x.shape[2]
         iter_ = 1 + (t - 1) // 4
         for i in range(iter_):
@@ -681,6 +739,9 @@ class WanVAE(nn.Module):
         return mu
 
     def decode(self, z):
+        if z.ndim == 4:
+            out = self.decoder(self.conv2(z.unsqueeze(2)), first_chunk=True)
+            return unpatchify(out, patch_size=self.patch_size).squeeze(2)
         conv_idx = [0]
         feat_map = [None] * count_conv3d(self.decoder)
         iter_ = z.shape[2]
@@ -701,7 +762,7 @@ class WanVAE(nn.Module):
                     feat_idx=conv_idx,
                 )
                 out = torch.cat([out, out_], 2)
-        out = unpatchify(out, patch_size=2)
+        out = unpatchify(out, patch_size=self.patch_size)
         return out
 
     def reparameterize(self, mu, log_var):

comfy/lora.py

@@ -327,11 +327,16 @@ def model_lora_keys_unet(model, key_map={}):
         for k in sdk:
             if k.startswith("diffusion_model.") and k.endswith(".weight"): #QwenImage lora format
                 key_lora = k[len("diffusion_model."):-len(".weight")]
-                # Direct mapping for transformer_blocks format (QwenImage LoRA format)
-                key_map["{}".format(key_lora)] = k
-                # Support transformer prefix format
-                key_map["transformer.{}".format(key_lora)] = k
-                key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = k #SimpleTuner lycoris format
+                targets = [(key_lora, k)]
+                if key_lora.endswith(".img_mlp.gate_up"):  # Qwen Image 2.1 fuses gate_layer/proj at load; LoRAs address the halves
+                    half = sd[k].shape[0] // 2
+                    targets = [(key_lora.replace(".gate_up", ".gate_layer"), (k, (0, 0, half))), (key_lora.replace(".gate_up", ".proj"), (k, (0, half, half)))]
+                for key_lora, to in targets:
+                    # Direct mapping for transformer_blocks format (QwenImage LoRA format)
+                    key_map["{}".format(key_lora)] = to
+                    # Support transformer prefix format
+                    key_map["transformer.{}".format(key_lora)] = to
+                    key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = to #SimpleTuner lycoris format
 
     if isinstance(model, comfy.model_base.Krea2):
         diffusers_keys = comfy.utils.krea2_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")

comfy/model_base.py

@@ -63,6 +63,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.qwen_image21.model
 import comfy.ldm.mage_flow.model
 import comfy.ldm.joyimage.model
 import comfy.ldm.ideogram4.model
@@ -2644,7 +2645,8 @@ class QwenImage(BaseModel):
         out = {}
         ref_latents = kwargs.get("reference_latents", None)
         if ref_latents is not None:
-            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
+            c = self.latent_format.latent_channels
+            out['ref_latents'] = list([1, c, sum(map(lambda a: math.prod(a.size()), ref_latents)) // c])
         return out
 
 class MageFlow(QwenImage):
@@ -2655,11 +2657,27 @@ class MageFlow(QwenImage):
         # Mage runs in bf16 and rounds its timestep frequency table to the timestep dtype, keep that on fp32 devices.
         return timestep.to(torch.bfloat16)
 
-    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])
+class QwenImage21(QwenImage):
+    def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image21.model.QwenImage21Transformer2DModel)
+
+    @property
+    def current_patcher(self):
+        return self._current_patcher
+
+    @current_patcher.setter
+    def current_patcher(self, patcher):
+        # set by pre_run / cleanup: the prefix K/V cache lives for one sampling run, off when hooks can repatch weights mid-run
+        self._current_patcher = patcher
+        diffusion_model = getattr(self, "diffusion_model", None)
+        if diffusion_model is not None:
+            diffusion_model.reset_prefix_cache(patcher is not None and len(patcher.hook_patches) == 0)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        image_slots = kwargs.get("image_slots", None)
+        if image_slots is not None:
+            out['image_slots'] = comfy.conds.CONDConstant(image_slots)
         return out
 
 class JoyImage(BaseModel):

comfy/model_detection.py

@@ -959,6 +959,27 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
         return dit_config
 
+    qwen_image21_keys = ['txt_in.text_norm.weight', 'modulation.1.weight', 'transformer_blocks.0.attn.norm_q.weight', 'img_in.weight', 'proj_out.weight']
+    if all('{}{}'.format(key_prefix, k) in state_dict_keys for k in qwen_image21_keys) and any('{}transformer_blocks.0.img_mlp.{}.weight'.format(key_prefix, k) in state_dict_keys for k in ('gate_up', 'proj')):  # Qwen Image 2.1
+        dit_config = {}
+        dit_config["image_model"] = "qwen_image21"
+        head_dim = state_dict['{}transformer_blocks.0.attn.norm_q.weight'.format(key_prefix)].shape[0]
+        inner_dim = state_dict['{}img_in.weight'.format(key_prefix)].shape[0]
+        dit_config["in_channels"] = state_dict['{}img_in.weight'.format(key_prefix)].shape[1]
+        dit_config["out_channels"] = state_dict['{}proj_out.weight'.format(key_prefix)].shape[0]
+        dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
+        dit_config["attention_head_dim"] = head_dim
+        dit_config["num_attention_heads"] = inner_dim // head_dim
+        dit_config["context_in_dim"] = state_dict['{}txt_in.text_norm.weight'.format(key_prefix)].shape[0]
+        # gate and up projections fuse into one GEMM when their rows can be concatenated: plain weights or per-row scales; a comfy-saved file is already fused
+        gate_up = state_dict.get('{}transformer_blocks.0.img_mlp.gate_up.weight'.format(key_prefix), None)
+        if gate_up is not None:
+            dit_config["mlp_ratio"] = gate_up.shape[0] // 2 // inner_dim
+        else:
+            dit_config["mlp_ratio"] = state_dict['{}transformer_blocks.0.img_mlp.proj.weight'.format(key_prefix)].shape[0] // inner_dim
+        dit_config["fused_mlp"] = gate_up is not None
+        return dit_config
+
     if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys:  # Qwen Image
         dit_config = {}
         dit_config["image_model"] = "qwen_image"

comfy/sd.py

@@ -64,6 +64,7 @@ import comfy.text_encoders.hidream
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
+import comfy.text_encoders.qwen_image21
 import comfy.text_encoders.hunyuan_image
 import comfy.text_encoders.z_image
 import comfy.text_encoders.krea2
@@ -382,6 +383,8 @@ class CLIP:
                     o = self.cond_stage_model.encode_token_weights(tokens)
                     cond, pooled = o[:2]
                     pooled_dict = {"pooled_output": pooled}
+                    if len(o) > 2:
+                        pooled_dict.update(o[2])
                     # add clip_start_percent and clip_end_percent in pooled
                     pooled_dict["clip_start_percent"] = t_range[0]
                     pooled_dict["clip_end_percent"] = t_range[1]
@@ -826,7 +829,20 @@ class VAE:
                 self.memory_used_encode = lambda shape, dtype: (50 * (round((shape[2] + 7) / 8) * 8) * shape[3] * shape[4]) * model_management.dtype_size(dtype)
                 self.working_dtypes = [torch.bfloat16, torch.float32]
             elif "decoder.middle.0.residual.0.gamma" in sd:
-                if "decoder.upsamples.0.upsamples.0.residual.2.weight" in sd:  # Wan 2.2 VAE
+                wan22_layout = "decoder.upsamples.0.upsamples.0.residual.2.weight" in sd
+                head = sd.get("decoder.head.2.weight", None)
+                if wan22_layout and head is not None and head.ndim == 5 and head.shape[2] == 1:  # Qwen Image 2.1 VAE: Wan 2.2 layout, temporal kernel 1, no patchify, RGBA
+                    self.upscale_ratio = 16
+                    self.downscale_ratio = 16
+                    self.latent_channels = 64
+                    self.output_channels = sd["decoder.head.2.weight"].shape[0]
+                    self.pad_channel_value = 1.0  # opaque alpha for RGB input
+                    ddconfig = {"dim": sd["encoder.conv1.weight"].shape[0], "dec_dim": sd["decoder.head.0.gamma"].shape[0], "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 8, 8], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True, True], "dropout": 0.0, "image_channels": self.output_channels, "patch_size": 1, "temporal_kernel": 1}
+                    self.first_stage_model = comfy.ldm.wan.vae2_2.WanVAE(**ddconfig)
+                    self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+                    self.memory_used_encode = lambda shape, dtype: 600 * shape[2] * shape[3] * model_management.dtype_size(dtype)
+                    self.memory_used_decode = lambda shape, dtype: 900 * shape[2] * shape[3] * (16 * 16) * model_management.dtype_size(dtype)
+                elif wan22_layout:  # Wan 2.2 VAE
                     self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
                     self.upscale_index_formula = (4, 16, 16)
                     self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
@@ -1095,6 +1111,9 @@ class VAE:
         self.patcher = mp(self.first_stage_model, load_device=self.device, offload_device=offload_device, fast_disk=fast_disk)
 
         m, u = self.first_stage_model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
+        if not self.patcher.is_dynamic():
+            # Lazy parameters only exist after loading the state dict.
+            self.first_stage_model.to(self.vae_dtype)
         if len(m) > 0:
             logging.warning("Missing VAE keys {}".format(m))
 
@@ -1150,9 +1169,9 @@ class VAE:
 
         decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
         output = self.process_output(
-            (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
-            comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
-             comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar))
+            (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device, pbar = pbar) +
+            comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device, pbar = pbar) +
+             comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device, pbar = pbar))
             / 3.0)
         return output
 
@@ -1933,6 +1952,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.joyimage.te(**llama_detect(clip_data))
                 clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer
+            elif clip_type == CLIPType.QWEN_IMAGE and te_model == TEModel.QWEN3VL_8B:  # Qwen-Image 2.1: full Qwen3-VL-8B, last hidden state, image slots spliced by the DiT.
+                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.qwen_image21.te(**llama_detect(clip_data))
+                clip_target.tokenizer = comfy.text_encoders.qwen_image21.QwenImage21Tokenizer
             elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2):  # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
                 klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
                 clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)

comfy/supported_models.py

@@ -23,6 +23,7 @@ import comfy.text_encoders.wan
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
+import comfy.text_encoders.qwen_image21
 import comfy.text_encoders.hunyuan_image
 import comfy.text_encoders.kandinsky5
 import comfy.text_encoders.z_image
@@ -2051,6 +2052,35 @@ class QwenImage(supported_models_base.BASE):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
 
+class QwenImage21(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "qwen_image21",
+    }
+
+    # scheduler mu at 1024x1024 (base 0.5 @ 256 tokens, max 0.9 @ 8192)
+    sampling_settings = {
+        "multiplier": 1.0,
+        "shift": 0.69,
+    }
+
+    memory_usage_factor = 6.0
+
+    unet_extra_config = {}
+    latent_format = latent_formats.QwenImage21
+
+    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.QwenImage21(self, device=device)
+
+    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_8b.transformer.".format(pref))
+        return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image21.QwenImage21Tokenizer, comfy.text_encoders.qwen_image21.te(**hunyuan_detect))
+
 class JoyImage(supported_models_base.BASE):
     unet_config = {
         "image_model": "joyimage",
@@ -2610,6 +2640,7 @@ models = [
     Boogu,
     MageFlow,
     QwenImage,
+    QwenImage21,
     JoyImage,
     Ideogram4,
     Krea2,

comfy/text_encoders/qwen_image21.py

@@ -0,0 +1,89 @@
+import numbers
+
+import torch
+
+import comfy.text_encoders.qwen3vl
+from comfy import sd1_clip
+
+VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
+SYSTEM_PROMPT = "<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n"
+T2I_TEMPLATE = SYSTEM_PROMPT + "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+
+
+class QwenImage21Tokenizer(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_8b")
+        self.llama_template = T2I_TEMPLATE
+
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=True, keep_vision=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])]
+        if llama_template is None and len(images) > 0:
+            refs = " ".join("<image{}>{}".format(i + 1, VISION_BLOCK) for i in range(len(images)))
+            llama_template = T2I_TEMPLATE.replace("{}", refs + "{}", 1)
+        out = 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)
+        out["keep_vision"] = keep_vision
+        return out
+
+
+class QwenImage21Qwen3VLClipModel(comfy.text_encoders.qwen3vl.Qwen3VLClipModel):
+    def __init__(self, device="cpu", dtype=None, attention_mask=True, model_options={}):
+        super().__init__(device=device, dtype=dtype, attention_mask=attention_mask, model_options=model_options, model_type="qwen3vl_8b")
+        # last layer without the final RMSNorm: transformers 4.57 hidden_states[-1], which Qwen's results are tuned to (5.x norms it)
+        self.layer_norm_hidden_state = False
+        self.image_spans = []
+
+    def process_tokens(self, tokens, device):
+        embeds, attention_mask, num_tokens, embeds_info = super().process_tokens(tokens, device)
+        self.image_spans = [(e["index"], e["size"]) for e in embeds_info if e["type"] == "image"]
+        return embeds, attention_mask, num_tokens, embeds_info
+
+
+class QwenImage21TEModel(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}):
+        super().__init__(device=device, dtype=dtype, name="qwen3vl_8b", clip_model=QwenImage21Qwen3VLClipModel, model_options=model_options)
+
+    def encode_token_weights(self, token_weight_pairs):
+        out, pooled, extra = super().encode_token_weights(token_weight_pairs)
+        tokens = [t[0] for t in token_weight_pairs["qwen3vl_8b"][0]]
+        image_spans = getattr(self, self.clip).image_spans
+
+        # drop the system turn, everything before the second <|im_start|>; positions shift by each image expanded before it
+        im_starts, offset, spans = [], 0, iter(image_spans)
+        for i, t in enumerate(tokens):
+            if isinstance(t, numbers.Integral):
+                if t == 151644:
+                    im_starts.append(i + offset)
+            elif isinstance(t, dict) and t.get("type") == "image":  # a textual embedding is a bare tensor and has no span
+                offset += next(spans, (0, 1))[1] - 1
+        keep = torch.ones(out.shape[1], dtype=torch.bool)
+        keep[:im_starts[1] if len(im_starts) > 1 else 0] = False
+
+        # vision tokens are replaced by reference latents in the DiT: drop them and record where each image goes
+        # with no latents coming (no vae) they stay, and the image conditions through the text encoder alone
+        slots = []
+        if not token_weight_pairs.get("keep_vision", False):
+            for start, size in image_spans:
+                keep[start:start + size] = False
+                slots.append(int(keep[:start].sum()))
+
+        out = out[:, keep.to(out.device)]
+        extra["attention_mask"] = extra["attention_mask"][:, keep.to(extra["attention_mask"].device)]
+        if extra["attention_mask"].sum() == torch.numel(extra["attention_mask"]):
+            extra.pop("attention_mask")
+        if len(slots) > 0:
+            extra["image_slots"] = slots
+        return out, pooled, extra
+
+
+def te(dtype_llama=None, llama_quantization_metadata=None):
+    class QwenImage21TEModel_(QwenImage21TEModel):
+        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 QwenImage21TEModel_

comfy_extras/nodes_qwen.py

@@ -108,6 +108,108 @@ class TextEncodeQwenImageEditPlus(io.ComfyNode):
         return io.NodeOutput(conditioning)
 
 
+class TextEncodeQwenImage21(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncodeQwenImage21",
+            display_name="Text Encode Qwen Image 2.1",
+            category="model/conditioning/qwen image",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
+                io.String.Input("negative_prompt", multiline=True, dynamic_prompts=True),
+                io.Vae.Input("vae", optional=True),
+                io.Int.Input("resolution", default=1024, min=0, max=4096, step=32,
+                             tooltip="Reference images are resized to about resolution x resolution pixels, at multiples of 32, preserving aspect ratio. 0 keeps each reference at its own size, rounded to a multiple of 32. "),
+                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 images, seen by the text encoder and spliced into the sequence as VAE latents.",
+                ),
+            ],
+            outputs=[
+                io.Conditioning.Output(display_name="positive"),
+                io.Conditioning.Output(display_name="negative"),
+                io.Latent.Output(display_name="latent",
+                                 tooltip="Empty latent on the first reference image's size, to match with sampling as any other size shifts the edit."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip, prompt, negative_prompt, vae=None, resolution=1024, images: io.Autogrow.Type = None) -> io.NodeOutput:
+        ref_latents = []
+        images_vl = []
+        images = images or {}
+        latent_w = latent_h = resolution or 1024
+        for name in sorted(images, key=lambda n: int(n.rsplit("_", 1)[-1])):
+            image = images[name]
+            if image is None:
+                continue
+            # same resize for the text encoder and the VAE, so every vision slot covers 2x2 latents; one image per input
+            samples = image[:1].movedim(-1, 1)
+            if resolution > 0:
+                ratio = samples.shape[3] / samples.shape[2]
+                width = round(math.sqrt(resolution * resolution * ratio) / 32) * 32
+                height = round(math.sqrt(resolution * resolution / ratio) / 32) * 32
+            else:
+                width, height = round(samples.shape[3] / 32) * 32, round(samples.shape[2] / 32) * 32
+            width, height = max(32, width), max(32, height)
+            if (width, height) == (samples.shape[3], samples.shape[2]):
+                s = image[:1]
+            else:
+                s = comfy.utils.common_upscale(samples, width, height, "lanczos", "disabled").movedim(1, -1)
+            if not images_vl:
+                latent_w, latent_h = width, height
+            rgb = s[:, :, :, :3]
+            if s.shape[-1] > 3:
+                rgb = rgb * s[:, :, :, 3:] + (1.0 - s[:, :, :, 3:])  # the vision tower sees alpha over white, the vae keeps all four
+            images_vl.append(rgb)
+            if vae is not None:
+                ref_latents.append(vae.encode(s))
+
+        keep_vision = len(ref_latents) == 0
+        positive = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=images_vl, keep_vision=keep_vision, prevent_empty_text=True))
+        negative = clip.encode_from_tokens_scheduled(clip.tokenize(negative_prompt, images=images_vl, keep_vision=keep_vision, prevent_empty_text=True))
+        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([1, 64, latent_h // 16, latent_w // 16], device=comfy.model_management.intermediate_device())
+        return io.NodeOutput(positive, negative, {"samples": latent})
+
+
+class QwenImage21Cache(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="QwenImage21Cache",
+            display_name="Qwen Image 2.1 Cache",
+            category="model/conditioning/qwen image",
+            description=(
+                "Allows setting the KV cache device and quantization, by default the model uses the auto -option. Quantization potentially speeds up edit workflows when memory starved "
+            ),
+            inputs=[
+                io.Model.Input("model"),
+                io.Combo.Input("device", options=["auto", "gpu", "cpu", "off"], default="auto",
+                               tooltip="auto uses spare VRAM, then RAM. cpu (RAM) is prefetched behind compute and costs little speed. off recomputes the prefix every step, which is slower but is the one way to rule the cache out."),
+                io.Combo.Input("dtype", options=["default", "int8", "int4"], default="default",
+                               tooltip="Storage precision. default is lossless. int8 halves the cache at about bf16 accuracy, int4 quarters it but roughly doubles the per-step error."),
+            ],
+            outputs=[io.Model.Output()],
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, model, device, dtype) -> io.NodeOutput:
+        m = model.clone()
+        m.model_options["transformer_options"]["qwen_image21_cache"] = {"device": device, "dtype": dtype}
+        return io.NodeOutput(m)
+
+
 class EmptyQwenImageLayeredLatentImage(io.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -138,6 +240,8 @@ class QwenExtension(ComfyExtension):
         return [
             TextEncodeQwenImageEdit,
             TextEncodeQwenImageEditPlus,
+            TextEncodeQwenImage21,
+            QwenImage21Cache,
             EmptyQwenImageLayeredLatentImage,
         ]