[pull] master from Comfy-Org:master

comfy/cli_args.py

@@ -112,7 +112,7 @@ def from_string(cls, value: str):
 parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
 
 cache_group = parser.add_mutually_exclusive_group()
-cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 10%% of system RAM (min 2GB, max 10GB), inactive 100%% of system RAM (max 96GB).")
+cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 10%% of system RAM (min 2GB, max 10GB), inactive 100%% of system RAM (max 128GB).")
 cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
 cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
 cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")

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)