diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index e6500cff48b..2c549e09545 100644 --- a/comfy/ldm/modules/attention.py +++ b/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: diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index e44322e7277..6b698f767df 100644 --- a/comfy/model_patcher.py +++ b/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 diff --git a/comfy/ops.py b/comfy/ops.py index 9d692dcc719..6c3845eef98 100644 --- a/comfy/ops.py +++ b/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) diff --git a/comfy_api_nodes/nodes_bria.py b/comfy_api_nodes/nodes_bria.py index 8077f139857..9e5a9333052 100644 --- a/comfy_api_nodes/nodes_bria.py +++ b/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"}}""", ), ) diff --git a/comfy_api_nodes/util/conversions.py b/comfy_api_nodes/util/conversions.py index 9cd644fc061..f46cac3f815 100644 --- a/comfy_api_nodes/util/conversions.py +++ b/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)) diff --git a/requirements.txt b/requirements.txt index 35dd5e613ad..af414d413bb 100644 --- a/requirements.txt +++ b/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 diff --git a/tests-unit/comfy_api_nodes_test/audio_conversions_test.py b/tests-unit/comfy_api_nodes_test/audio_conversions_test.py new file mode 100644 index 00000000000..d7c9b3899df --- /dev/null +++ b/tests-unit/comfy_api_nodes_test/audio_conversions_test.py @@ -0,0 +1,78 @@ +import math + +import av +import numpy as np +import pytest +import torch + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +from comfy_api_nodes.util.conversions import audio_input_to_mp3 # noqa: E402 + +SAMPLE_RATE = 48000 +DURATION = 2.0 +LEFT_HZ = 440.0 +RIGHT_HZ = 880.0 + + +def tone(freq, duration=DURATION, sample_rate=SAMPLE_RATE): + t = torch.arange(int(sample_rate * duration), dtype=torch.float32) / sample_rate + return 0.5 * torch.sin(2 * math.pi * freq * t) + + +@pytest.fixture +def stereo_audio(): + """Comfy AUDIO with two tones that stay distinguishable through mp3.""" + waveform = torch.stack([tone(LEFT_HZ), tone(RIGHT_HZ)]).unsqueeze(0) + return {"waveform": waveform, "sample_rate": SAMPLE_RATE} + + +@pytest.fixture +def mono_audio(): + return {"waveform": tone(LEFT_HZ).unsqueeze(0).unsqueeze(0), "sample_rate": SAMPLE_RATE} + + +def decode(buffer): + """(planes[C, N], sample_rate, channels) of an encoded mp3 buffer""" + buffer.seek(0) + with av.open(buffer, mode="r") as container: + stream = container.streams.audio[0] + planes = [] + for frame in container.decode(audio=0): + array = frame.to_ndarray() + if frame.format.is_planar: + planes.append(array) + else: + planes.append(array.reshape(-1, len(frame.layout.channels)).T) + return np.concatenate(planes, axis=1), stream.codec_context.sample_rate, len(stream.layout.channels) + + +def dominant_hz(signal, sample_rate): + """Peak frequency, ignoring the encoder's padding at either edge""" + edge = int(0.2 * sample_rate) + window = signal[edge:-edge] + spectrum = np.abs(np.fft.rfft(window * np.hanning(window.size))) + return np.fft.rfftfreq(window.size, 1.0 / sample_rate)[int(np.argmax(spectrum))] + + +def test_stereo_duration_is_preserved(stereo_audio): + planes, sample_rate, channels = decode(audio_input_to_mp3(stereo_audio)) + assert channels == 2 + assert sample_rate == SAMPLE_RATE + assert planes.shape[1] / sample_rate == pytest.approx(DURATION, abs=0.15) + + +def test_stereo_channels_are_not_concatenated(stereo_audio): + """The channels must be interleaved; concatenating them plays the clip twice.""" + planes, sample_rate, _ = decode(audio_input_to_mp3(stereo_audio)) + assert dominant_hz(planes[0], sample_rate) == pytest.approx(LEFT_HZ, abs=15) + assert dominant_hz(planes[1], sample_rate) == pytest.approx(RIGHT_HZ, abs=15) + + +def test_mono_duration_is_preserved(mono_audio): + planes, sample_rate, _ = decode(audio_input_to_mp3(mono_audio)) + assert planes.shape[1] / sample_rate == pytest.approx(DURATION, abs=0.15) + assert dominant_hz(planes[0], sample_rate) == pytest.approx(LEFT_HZ, abs=15)