Merge branch 'master' into jaeone94/native-lora-stack

Author: alexisrolland

.coderabbit.yaml

@@ -140,5 +140,16 @@ knowledge_base:
     filePatterns:
       - files: "AGENTS.md"
         applyTo: "**"
+  linked_repositories:
+    - repository: "Comfy-Org/ComfyUI_frontend"
+      instructions: |
+        Frontend consumer of Core APIs, routes, feature flags, asset behavior,
+        node metadata, and package-version rollouts. Check for compatibility
+        issues and changes that require coordinated rollout or merge ordering.
+    - repository: "Comfy-Org/comfy-kitchen"
+      instructions: |
+        Owns optimized inference operations and kernels used by ComfyUI. Check
+        whether new model math or backend-specific operations duplicate an
+        existing implementation or belong in comfy-kitchen instead of Core.
   learnings:
     scope: "auto"

comfy/configurable.py

@@ -0,0 +1,7 @@
+from torch import nn
+
+
+class ConfigurableModule(nn.Module):
+    def with_config(self, encoded_config):
+        """Return a new module configured from a uint8 JSON tensor, without modifying this module."""
+        raise NotImplementedError

comfy/ldm/lightricks/model.py

@@ -327,7 +327,7 @@ def forward(self, x):
         # Dropout, so leave it to the stock path whenever it could be active.
         if comfy.model_management.in_training:
             return self.net(x)
-        return self.net[2](self.net[0].proj(x), input_act="gelu_tanh")
+        return comfy.ops.linear_input_act(self.net[2], self.net[0].proj(x), "gelu_tanh")
 
 def apply_rotary_emb(input_tensor, freqs_cis):
     rotation_matrix, split_pe = freqs_cis

comfy/ldm/minimax/model.py

@@ -208,7 +208,7 @@ def __init__(self, hidden, ffn, dtype=None, device=None, operations=None):
         self.fc2 = operations.Linear(ffn, hidden, bias=False, dtype=dtype, device=device)
 
     def forward(self, x):
-        return self.fc2(self.fc1(x), input_act="swiglu")
+        return comfy.ops.linear_input_act(self.fc2, self.fc1(x), "swiglu")
 
 
 class AdalnProj(nn.Module):

comfy/ldm/minimax/vae.py

@@ -269,8 +269,8 @@ def __init__(self, dim, mult=4, bias=True, operations=ops):
 
     def forward(self, x, pre_norm, residual, residual_scale):
         # norm, gated silu and residual addcmul fold into the INT8 kernels
-        h = self.w1(x, input_act="rms_norm", act_weight=pre_norm)
-        return self.w2(h, input_act="swiglu", residual=residual, residual_scale=residual_scale)
+        h = comfy.ops.linear_input_act(self.w1, x, "rms_norm", act_weight=pre_norm)
+        return comfy.ops.linear_input_act(self.w2, h, "swiglu", residual=residual, residual_scale=residual_scale)
 
 
 class Attention(nn.Module):
@@ -289,7 +289,7 @@ def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
     def forward(self, x, rotary_pos_emb, pre_norm, residual, residual_scale):
         batch_size, seq_len, _ = x.shape
 
-        qkv = self.to_qkv(x, input_act="rms_norm", act_weight=pre_norm)
+        qkv = comfy.ops.linear_input_act(self.to_qkv, x, "rms_norm", act_weight=pre_norm)
         qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head)
         query, key, value = torch.chunk(qkv, 3, dim=-1)
 
@@ -313,7 +313,7 @@ def forward(self, x, rotary_pos_emb, pre_norm, residual, residual_scale):
             out = out.transpose(1, 2).reshape(batch_size, seq_len, -1)
         else:
             out = optimized_attention(query, key, value, self.heads, skip_reshape=True)
-        return self.to_out(torch.nan_to_num(out), residual=residual, residual_scale=residual_scale)
+        return comfy.ops.linear_input_act(self.to_out, torch.nan_to_num(out), None, residual=residual, residual_scale=residual_scale)
 
 
 class TransformerBlock(nn.Module):

comfy/ldm/minimax_music/ar.py

@@ -95,7 +95,7 @@ def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, op
 
     def forward(self, x):
         if self.merged_mlp:
-            return self.down_proj(self.gate_up_proj(x), input_act="swiglu")
+            return comfy.ops.linear_input_act(self.down_proj, self.gate_up_proj(x), "swiglu")
         return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
 
 

comfy/ldm/modules/attention.py

@@ -17,6 +17,7 @@
 from .sub_quadratic_attention import efficient_dot_product_attention
 
 from comfy import model_management
+from comfy.configurable import ConfigurableModule
 
 if model_management.xformers_enabled():
     import xformers
@@ -73,12 +74,17 @@ def get_attention_function(name: str, default: Any=...) -> Union[Callable, None]
     return REGISTERED_ATTENTION_FUNCTIONS[name]
 
 
-class ComfyAttention(nn.Module):
+class ComfyAttention(ConfigurableModule):
     def __init__(self):
         super().__init__()
         self.config = None
         self.function = None
 
+    def with_config(self, encoded_config):
+        attention = ComfyAttention()
+        attention.load_state_dict({"config": encoded_config})
+        return attention
+
     def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
         self.config = None
         self.function = None

comfy/ldm/qwen_image21/model.py

@@ -62,7 +62,7 @@ def __init__(self, dim, hidden_dim, fused=True, dtype=None, device=None, operati
 
     def forward(self, x):
         if self.fused:
-            return self.out(self.gate_up(x), input_act="swiglu")
+            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))
 
 

comfy/ldm/wan/model.py

@@ -196,7 +196,7 @@ class WanFeedForward(nn.Sequential):
     """[Linear, GELU(tanh), Linear], with the GELU folded into the down-projection."""
 
     def forward(self, x):
-        return self[2](self[0](x), input_act="gelu_tanh")
+        return comfy.ops.linear_input_act(self[2], self[0](x), "gelu_tanh")
 
 
 class WanAttentionBlock(nn.Module):

comfy/ldm/wan/model_animate2.py

@@ -11,6 +11,7 @@
 
 import comfy.ldm.common_dit
 import comfy.model_management
+import comfy.model_prefetch
 import comfy.quant_ops
 import comfy.utils
 from comfy.ldm.flux.math import apply_rope1
@@ -200,16 +201,17 @@ def prefetch(self, i, device, dtype):
         stream = None
         r = None
         if t.device != device:
-            stream = comfy.model_management.get_offload_stream(device)
-            cs = comfy.model_management.current_stream(device)
-            if stream is not None and cs is not None:
-                # the handed-out stream last waited on the main stream a full rotation ago, which does not cover the previous consumer's reads of this slot; wait now so the copy cannot overwrite a slot still being read
-                stream.wait_stream(cs)
             # two persistent staging buffers per tensor shape instead of a fresh allocation per block (~29 GB of churn per pass at 720p); windows of different lengths get their own pair
             buf_key = (tuple(t.shape), cast_dtype if cast_dtype is not None else t.dtype)
             if buf_key not in self._staging:
-                self._staging[buf_key] = [torch.empty(t.shape, dtype=buf_key[1], device=device) for _ in range(2)]
+                with comfy.model_prefetch.pause_malloc_graph():
+                    self._staging[buf_key] = [torch.empty(t.shape, dtype=buf_key[1], device=device) for _ in range(2)]
             r = self._staging[buf_key][i % 2]
+            stream = comfy.model_management.get_offload_stream(device)
+            cs = comfy.model_management.current_stream(device)
+            if stream is not None and cs is not None:
+                # Wait for staging allocation and the previous consumer's reads before overwriting the buffer.
+                stream.wait_stream(cs)
         self._pending[i] = (comfy.model_management.cast_to(t, cast_dtype, device, non_blocking=True, stream=stream, r=r), stream)
 
     def take(self, i, device, dtype, batch_size):