[pull] master from Comfy-Org:master

comfy/ldm/modules/attention.py

@@ -90,22 +90,6 @@ def default(val, d):
         return val
     return d
 
-def _gqa_repeat_factor(query_heads, key_heads, value_heads):
-    if key_heads != value_heads:
-        raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
-    if query_heads == key_heads:
-        return 1
-    if query_heads % key_heads != 0:
-        raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
-    return query_heads // key_heads
-
-def _repeat_kv_for_gqa(k, v, query_heads, head_dim):
-    n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
-    if n_rep > 1:
-        k = k.repeat_interleave(n_rep, dim=head_dim)
-        v = v.repeat_interleave(n_rep, dim=head_dim)
-    return k, v
-
 def _heads_from_dim(tensor, dim_head, name):
     inner_dim = tensor.shape[-1]
     if inner_dim % dim_head != 0:
@@ -122,10 +106,8 @@ def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_
         value_heads = heads
     k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head)
     v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head)
-    if enable_gqa:
-        _gqa_repeat_factor(heads, key_heads, value_heads)
-        if expand_kv:
-            k, v = _repeat_kv_for_gqa(k, v, heads, -2)
+    if enable_gqa and expand_kv:
+        k, v = comfy.ops.repeat_kv_for_gqa(k, v, heads, -2)
     return q, k, v
 
 
@@ -196,7 +178,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
     h = heads
     if skip_reshape:
         if kwargs.get("enable_gqa", False):
-            k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
+            k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-3], -3)
         q, k, v = map(
             lambda t: t.reshape(b * heads, -1, dim_head),
             (q, k, v),
@@ -262,7 +244,7 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
 
     if skip_reshape:
         if kwargs.get("enable_gqa", False):
-            key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3)
+            key, value = comfy.ops.repeat_kv_for_gqa(key, value, query.shape[-3], -3)
         query = query.reshape(b * heads, -1, dim_head)
         value = value.reshape(b * heads, -1, dim_head)
         key = key.reshape(b * heads, -1, dim_head).movedim(1, 2)
@@ -338,7 +320,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
 
     if skip_reshape:
         if kwargs.get("enable_gqa", False):
-            k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
+            k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-3], -3)
         q, k, v = map(
             lambda t: t.reshape(b * heads, -1, dim_head),
             (q, k, v),
@@ -476,7 +458,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
             (q, k, v),
         )
         if kwargs.get("enable_gqa", False):
-            k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2)
+            k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-2], -2)
     # actually do the reshaping
     else:
         dim_head //= heads
@@ -573,7 +555,7 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
         b, _, _, dim_head = q.shape
         tensor_layout = "HND"
         if kwargs.get("enable_gqa", False):
-            k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
+            k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-3], -3)
     else:
         b, _, dim_head = q.shape
         dim_head //= heads
@@ -671,7 +653,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
     if skip_reshape:
         q_s = q
         if kwargs.get("enable_gqa", False):
-            k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3)
+            k_s, v_s = comfy.ops.repeat_kv_for_gqa(k, v, H, -3)
         else:
             k_s, v_s = k, v
     else:

comfy/model_patcher.py

@@ -558,12 +558,9 @@ def match_multigpu_clones(self):
             new_multigpu_models = []
             for mm in multigpu_models:
                 # clone main model, but bring over relevant props from existing multigpu clone
-                n = self.clone()
+                n = self.clone(model_override=mm.get_clone_model_override())
                 n.load_device = mm.load_device
-                n.backup = mm.backup
-                n.object_patches_backup = mm.object_patches_backup
                 n.hook_backup = mm.hook_backup
-                n.model = mm.model
                 n.is_multigpu_base_clone = mm.is_multigpu_base_clone
                 n.remove_additional_models("multigpu")
                 orig_additional_models: dict[str, list[ModelPatcher]] = comfy.patcher_extension.copy_nested_dicts(n.additional_models)
@@ -1758,6 +1755,9 @@ def __init__(self, model, load_device, offload_device, size=0, weight_inplace_up
         self.register_load_device(self.load_device)
         self.non_dynamic_delegate_model = None
         assert load_device is not None
+        if not hasattr(self.model, "dynamic_patchers"):
+            self.model.dynamic_patchers = set()
+        self.model.dynamic_patchers.add(id(self))
 
     def register_load_device(self, device):
         """Ensure dynamic_pins has an entry for *device*.
@@ -1813,6 +1813,18 @@ def unpin_weight(self, key):
     def unpin_all_weights(self):
         self.partially_unload_ram(1e32)
 
+    def __del__(self):
+        model = getattr(self, "model", None)
+        dynamic_patchers = getattr(model, "dynamic_patchers", None)
+        if dynamic_patchers is None or id(self) not in dynamic_patchers:
+            return
+        dynamic_patchers.discard(id(self))
+        try:
+            if not dynamic_patchers:
+                self.unpin_all_weights()
+        finally:
+            self.detach(unpatch_all=False)
+
     def memory_required(self, input_shape):
         #Pad this significantly. We are trying to get away from precise estimates. This
         #estimate is only used when using the ModelPatcherDynamic after ModelPatcher. If you

comfy/ops.py

@@ -19,6 +19,7 @@
 import torch
 import logging
 import contextlib
+import inspect
 import comfy.model_management
 from comfy.cli_args import args, PerformanceFeature
 import comfy.float
@@ -36,30 +37,59 @@ def run_every_op():
 
     comfy.model_management.throw_exception_if_processing_interrupted()
 
+def gqa_repeat_factor(query_heads, key_heads, value_heads):
+    if key_heads != value_heads:
+        raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
+    if query_heads == key_heads:
+        return 1
+    if query_heads % key_heads != 0:
+        raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
+    return query_heads // key_heads
+
+def repeat_kv_for_gqa(k, v, query_heads, head_dim):
+    n_rep = gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
+    if n_rep > 1:
+        k = k.repeat_interleave(n_rep, dim=head_dim)
+        v = v.repeat_interleave(n_rep, dim=head_dim)
+    return k, v
+
 def scaled_dot_product_attention(q, k, v, *args, **kwargs):
+    attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
+    if kwargs.get("enable_gqa", False) and attn_mask is not None:
+        k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
+        kwargs["enable_gqa"] = False
     return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
 
 
 try:
     if torch.cuda.is_available():
         from torch.nn.attention import SDPBackend, sdpa_kernel
-        import inspect
         if "set_priority" in inspect.signature(sdpa_kernel).parameters:
             SDPA_BACKEND_PRIORITY = [
                 SDPBackend.FLASH_ATTENTION,
+                SDPBackend.CUDNN_ATTENTION,
                 SDPBackend.EFFICIENT_ATTENTION,
                 SDPBackend.MATH,
             ]
 
-            if comfy.model_management.WINDOWS:
-                SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION)
-            else:
-                SDPA_BACKEND_PRIORITY.insert(1, SDPBackend.CUDNN_ATTENTION)
-
             def scaled_dot_product_attention(q, k, v, *args, **kwargs):
-                if q.nelement() < 1024 * 128:  # arbitrary number, for small inputs cudnn attention seems slower
-                    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
+                attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
+                if kwargs.get("enable_gqa", False) and attn_mask is not None and not comfy.model_management.is_nvidia():
+                    k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
+                    kwargs["enable_gqa"] = False
                 with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
+                    if kwargs.get("enable_gqa", False) and attn_mask is not None and q.shape[-3] != k.shape[-3]:
+                        dropout_p = args[1] if len(args) > 1 else kwargs.get("dropout_p", 0.0)
+                        is_causal = args[2] if len(args) > 2 else kwargs.get("is_causal", False)
+                        params = torch.backends.cuda.SDPAParams(q, k, v, attn_mask, dropout_p, is_causal, True)
+                        supports_native_gqa = (
+                            torch.backends.cuda.can_use_flash_attention(params)
+                            or torch.backends.cuda.can_use_cudnn_attention(params)
+                            or torch.backends.cuda.can_use_efficient_attention(params)
+                        )
+                        if not supports_native_gqa:
+                            k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
+                            kwargs["enable_gqa"] = False
                     return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
         else:
             logging.warning("Torch version too old to set sdpa backend priority.")
@@ -464,8 +494,7 @@ class Linear(torch.nn.Linear, CastWeightBiasOp):
 
         def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
             # don't trust subclasses that BYO state dict loader to call us.
-            if (not comfy.model_management.WINDOWS
-                or not comfy.memory_management.aimdo_enabled
+            if (not comfy.memory_management.aimdo_enabled
                 or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):
                 super().__init__(in_features, out_features, bias, device, dtype)
                 return
@@ -487,8 +516,7 @@ def __init__(self, in_features, out_features, bias=True, device=None, dtype=None
         def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                                 strict, missing_keys, unexpected_keys, error_msgs):
 
-            if (not comfy.model_management.WINDOWS
-                or not comfy.memory_management.aimdo_enabled
+            if (not comfy.memory_management.aimdo_enabled
                 or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):
                 return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
                                                      missing_keys, unexpected_keys, error_msgs)
@@ -716,8 +744,7 @@ def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=Non
                      norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None,
                      _freeze=False, device=None, dtype=None):
             # don't trust subclasses that BYO state dict loader to call us.
-            if (not comfy.model_management.WINDOWS
-                or not comfy.memory_management.aimdo_enabled
+            if (not comfy.memory_management.aimdo_enabled
                 or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):
                 super().__init__(num_embeddings, embedding_dim, padding_idx, max_norm,
                                  norm_type, scale_grad_by_freq, sparse, _weight,
@@ -744,8 +771,7 @@ def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=Non
         def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                                 strict, missing_keys, unexpected_keys, error_msgs):
 
-            if (not comfy.model_management.WINDOWS
-                or not comfy.memory_management.aimdo_enabled
+            if (not comfy.memory_management.aimdo_enabled
                 or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):
                 return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
                                                      missing_keys, unexpected_keys, error_msgs)

comfy_api_nodes/nodes_bria.py

@@ -357,7 +357,7 @@ def define_schema(cls):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                expr="""{"type":"usd","usd":0.0042,"format":{"suffix":"/second"}}""",
+                expr="""{"type":"usd","usd":0.005,"format":{"suffix":"/second"}}""",
             ),
         )
 
@@ -433,7 +433,7 @@ def define_schema(cls):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                expr="""{"type":"usd","usd":0.0042,"format":{"suffix":"/second"}}""",
+                expr="""{"type":"usd","usd":0.005,"format":{"suffix":"/second"}}""",
             ),
         )
 

comfy_api_nodes/util/conversions.py

@@ -266,16 +266,14 @@ def audio_tensor_to_contiguous_ndarray(waveform: torch.Tensor) -> np.ndarray:
         waveform: a tensor of shape (1, channels, samples) derived from a Comfy `AUDIO` type.
 
     Returns:
-        Contiguous numpy array of the audio waveform. If the audio was batched,
-            the first item is taken.
+        Contiguous numpy array of the audio waveform.
+
+    Raises:
+        ValueError: If the waveform is not shaped (1, channels, samples).
     """
     if waveform.ndim != 3 or waveform.shape[0] != 1:
         raise ValueError("Expected waveform tensor shape (1, channels, samples)")
 
-    # If batch is > 1, take first item
-    if waveform.shape[0] > 1:
-        waveform = waveform[0]
-
     # Prepare for av: remove batch dim, move to CPU, make contiguous, convert to numpy array
     audio_data_np = waveform.squeeze(0).cpu().contiguous().numpy()
     if audio_data_np.dtype != np.float32:
@@ -285,20 +283,21 @@ def audio_tensor_to_contiguous_ndarray(waveform: torch.Tensor) -> np.ndarray:
 
 
 def audio_input_to_mp3(audio: Input.Audio) -> BytesIO:
-    waveform = audio["waveform"].cpu()
+    audio_data_np = audio_tensor_to_contiguous_ndarray(audio["waveform"])
+    sample_rate = int(audio["sample_rate"])
 
     output_buffer = BytesIO()
     output_container = av.open(output_buffer, mode="w", format="mp3")
 
-    out_stream = output_container.add_stream("libmp3lame", rate=audio["sample_rate"])
+    out_stream = output_container.add_stream("libmp3lame", rate=sample_rate)
     out_stream.bit_rate = 320000
 
     frame = av.AudioFrame.from_ndarray(
-        waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
-        format="flt",
-        layout="mono" if waveform.shape[0] == 1 else "stereo",
+        audio_data_np,
+        format="fltp",
+        layout="stereo" if audio_data_np.shape[0] > 1 else "mono",
     )
-    frame.sample_rate = audio["sample_rate"]
+    frame.sample_rate = sample_rate
     frame.pts = 0
     output_container.mux(out_stream.encode(frame))
     output_container.mux(out_stream.encode(None))

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.47.11
-comfyui-workflow-templates==0.11.20
+comfyui-workflow-templates==0.11.23
 comfyui-embedded-docs==0.5.9
 torch
 torchsde