From e4eabfefd540f3a6066c29761cefb8f8b3c80f68 Mon Sep 17 00:00:00 2001 From: Artificial Sweetener Date: Wed, 30 Sep 2026 21:48:07 -0400 Subject: [PATCH] fix(negpip): support Krea attention on ComfyUI 0.28 --- simple_syrup/runtime/negpip/krea2.py | 15 +- simple_syrup/runtime/negpip/krea2_host.py | 144 ++++++++++++++ simple_syrup/services/negpip_model_service.py | 2 + .../negpip/test_krea2_host_compatibility.py | 182 ++++++++++++++++++ .../negpip/test_negpip_model_service.py | 5 + 5 files changed, 344 insertions(+), 4 deletions(-) create mode 100644 simple_syrup/runtime/negpip/krea2_host.py create mode 100644 tests/prompting/negpip/test_krea2_host_compatibility.py diff --git a/simple_syrup/runtime/negpip/krea2.py b/simple_syrup/runtime/negpip/krea2.py index 51896e3..112b809 100644 --- a/simple_syrup/runtime/negpip/krea2.py +++ b/simple_syrup/runtime/negpip/krea2.py @@ -185,9 +185,16 @@ def krea2_diffusion_negpip_wrapper( *args: object, **kwargs: object, ) -> object: - """Move a processed Krea sign mask into call-local transformer options.""" + """Inject call-local signs at either supported Comfy Krea argument boundary.""" - positional_options = args[5] if len(args) > 5 else None + options_index = ( + 5 + if len(args) > 5 + else 4 + if len(args) == 5 and isinstance(args[4], dict) + else None + ) + positional_options = args[options_index] if options_index is not None else None transformer_options = ( positional_options if positional_options is not None @@ -201,9 +208,9 @@ def krea2_diffusion_negpip_wrapper( if not isinstance(multiplier, torch.Tensor): raise TypeError("Krea NegPiP processed mask must be a tensor.") prepared[TRANSFORMER_MASK_KEY] = multiplier - if len(args) > 5: + if options_index is not None: prepared_args = list(args) - prepared_args[5] = prepared + prepared_args[options_index] = prepared return executor(*prepared_args, **kwargs) kwargs["transformer_options"] = prepared return executor(*args, **kwargs) diff --git a/simple_syrup/runtime/negpip/krea2_host.py b/simple_syrup/runtime/negpip/krea2_host.py new file mode 100644 index 0000000..5791fc6 --- /dev/null +++ b/simple_syrup/runtime/negpip/krea2_host.py @@ -0,0 +1,144 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Supply model-local NegPiP attention hooks for Comfy's earlier Krea boundary.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from inspect import signature +from typing import Any, cast + +import torch +from comfy.ldm.flux.math import apply_rope +from comfy.ldm.krea2.model import Attention +from comfy.ldm.modules.attention import optimized_attention_masked +from comfy.model_patcher import ModelPatcher +from einops import rearrange + +from ..model_patcher_mutations import ModelCallableObjectPatchMutation +from .krea2 import TRANSFORMER_MASK_KEY + +LOGGER = logging.getLogger(__name__) + + +def krea2_host_mutations( + model: ModelPatcher, +) -> tuple[ModelCallableObjectPatchMutation, ...]: + """Adapt only the known five-argument host; retain native reference-capable Krea.""" + diffusion = model.get_model_object("diffusion_model") + parameters = tuple(signature(diffusion._forward).parameters) + if "ref_latents" in parameters: + return () + if parameters != ( + "x", + "timesteps", + "context", + "attention_mask", + "transformer_options", + "kwargs", + ): + raise ValueError( + f"Krea NegPiP does not support model signature {parameters!r}." + ) + mutations: list[ModelCallableObjectPatchMutation] = [] + for index, block in enumerate(diffusion.blocks): + attention = block.attn + if not isinstance(attention, Attention): + raise TypeError(f"Krea NegPiP requires host Attention at block {index}.") + mutations.append( + ModelCallableObjectPatchMutation( + f"diffusion_model.blocks.{index}.attn.forward", + Krea2HostAttention(attention, index, len(diffusion.blocks)), + ) + ) + if not mutations: + raise ValueError("Krea NegPiP requires at least one joint attention block.") + LOGGER.info( + "Krea NegPiP installed model-local host attention hooks", + extra={"blocks": len(mutations)}, + ) + return tuple(mutations) + + +@dataclass(frozen=True) +class Krea2HostAttention: + """Preserve host attention math while exposing its missing pre-RoPE patch point. + + Comfy 0.28's public model boundary lacks attention callbacks. Object patches + scope this adapter to the derived MODEL and Comfy restores them on unload. + Text-fusion attention stays untouched; only joint text/image blocks use it. + """ + + attention: Any # Comfy Attention exposes dynamically constructed linear modules. + block_index: int + total_blocks: int + + def __call__( + self, + x: torch.Tensor, + freqs: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + transformer_options: dict[str, Any] | None = None, + ) -> torch.Tensor: + """Run host QKV projections and attention with local patch metadata.""" + options = {} if transformer_options is None else transformer_options.copy() + multiplier = options.get(TRANSFORMER_MASK_KEY) + if multiplier is None: + return cast( + torch.Tensor, + type(self.attention).forward( + self.attention, + x, + freqs, + mask, + transformer_options=options, + ), + ) + if not isinstance(multiplier, torch.Tensor) or multiplier.ndim != 3: + raise ValueError( + "Krea NegPiP host attention requires a processed token sign mask." + ) + options.update( + block_index=self.block_index, + total_blocks=self.total_blocks, + block_type="single", + img_slice=[multiplier.shape[1], x.shape[1]], + ) + attention = self.attention + q, k, v, gate = ( + attention.wq(x), + attention.wk(x), + attention.wv(x), + attention.gate(x), + ) + q = rearrange(q, "B L (H D) -> B H L D", H=attention.heads) + k = rearrange(k, "B L (H D) -> B H L D", H=attention.kvheads) + v = rearrange(v, "B L (H D) -> B H L D", H=attention.kvheads) + q, k = attention.qknorm(q, k) + for patch in options.get("patches", {}).get("attn1_patch", []): + result = patch( + q, k, v, pe=freqs, attn_mask=mask, extra_options=options.copy() + ) + q, k, v = result.get("q", q), result.get("k", k), result.get("v", v) + freqs, mask = result.get("pe", freqs), result.get("attn_mask", mask) + if freqs is not None: + q, k = apply_rope(q, k, freqs) + if attention.kvheads != attention.heads: + repeats = attention.heads // attention.kvheads + k = k.repeat_interleave(repeats, dim=1) + v = v.repeat_interleave(repeats, dim=1) + out = optimized_attention_masked( + q, + k, + v, + attention.heads, + mask=mask, + skip_reshape=True, + transformer_options=options, + ) + for patch in options.get("patches", {}).get("attn1_output_patch", []): + out = patch(out, options.copy()) + return cast(torch.Tensor, attention.wo(out * torch.nn.functional.sigmoid(gate))) diff --git a/simple_syrup/services/negpip_model_service.py b/simple_syrup/services/negpip_model_service.py index 09bd0bb..d52b3f3 100644 --- a/simple_syrup/services/negpip_model_service.py +++ b/simple_syrup/services/negpip_model_service.py @@ -48,6 +48,7 @@ from ..runtime.negpip.krea2 import ( from ..runtime.negpip.krea2 import ( WRAPPER_KEY as KREA_WRAPPER_KEY, ) +from ..runtime.negpip.krea2_host import krea2_host_mutations from ..runtime.negpip.standard import ( encode_token_weights_negpip, standard_attn2_negpip, @@ -190,6 +191,7 @@ class NegpipModelService: prepared_model = PATCHER_LIFECYCLE.derive_model( model, ( + *krea2_host_mutations(model), ModelCallableObjectPatchMutation( "extra_conds", krea2_extra_conds_negpip_wrapper(previous), diff --git a/tests/prompting/negpip/test_krea2_host_compatibility.py b/tests/prompting/negpip/test_krea2_host_compatibility.py new file mode 100644 index 0000000..565621e --- /dev/null +++ b/tests/prompting/negpip/test_krea2_host_compatibility.py @@ -0,0 +1,182 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Safeguard Krea NegPiP at both supported Comfy model call boundaries.""" + +from __future__ import annotations + +from typing import cast + +import pytest +import torch +from comfy.ldm.krea2.model import Attention, SingleStreamDiT +from comfy.model_patcher import ModelPatcher + +from simple_syrup.runtime.negpip.krea2 import ( + CONDITION_MASK_KEY, + TRANSFORMER_MASK_KEY, + krea2_attn1_negpip, + krea2_diffusion_negpip_wrapper, +) +from simple_syrup.runtime.negpip.krea2_host import ( + Krea2HostAttention, + krea2_host_mutations, +) +from simple_syrup.runtime.patcher_lifecycle import PATCHER_LIFECYCLE + + +@pytest.mark.parametrize("reference_argument", (False, True)) +def test_diffusion_wrapper_preserves_both_host_signatures( + reference_argument: bool, +) -> None: + """Inject one local mask without binding transformer options twice.""" + source_options: dict[str, object] = {"existing": True} + mask = torch.tensor([[[-1.0], [1.0]]]) + + def earlier( + x: object, + timesteps: object, + context: object, + attention_mask: object, + transformer_options: dict[str, object], + **kwargs: object, + ) -> dict[str, object]: + """Expose Comfy 0.28's exact five-positional model boundary.""" + del x, timesteps, context, attention_mask, kwargs + return transformer_options + + def current( + x: object, + timesteps: object, + context: object, + attention_mask: object, + ref_latents: object, + transformer_options: dict[str, object], + **kwargs: object, + ) -> dict[str, object]: + """Expose Comfy's reference-capable six-positional model boundary.""" + del x, timesteps, context, attention_mask, kwargs + assert ref_latents is None + return transformer_options + + args = ( + (None, None, None, None, None, source_options) + if reference_argument + else (None, None, None, None, source_options) + ) + prepared = cast( + dict[str, object], + krea2_diffusion_negpip_wrapper( + current if reference_argument else earlier, + *args, + **{CONDITION_MASK_KEY: mask}, + ), + ) + assert prepared is not source_options + assert prepared[TRANSFORMER_MASK_KEY] is mask + assert prepared["existing"] is True + assert TRANSFORMER_MASK_KEY not in source_options + + +@pytest.mark.parametrize("negative_sign", (False, True)) +def test_host_attention_matches_native_math_and_keeps_call_options_local( + negative_sign: bool, +) -> None: + """Earlier host hooks reproduce native grouped attention without source mutation.""" + torch.manual_seed(4) + attention = Attention(128, 4, kvheads=2, operations=torch.nn) + with torch.no_grad(): + attention.qknorm.qnorm.scale.zero_() + attention.qknorm.knorm.scale.zero_() + x = torch.randn(1, 7, 128) + mask = torch.tensor([[[-1.0 if negative_sign else 1.0], [1.0]]]) + options = { + TRANSFORMER_MASK_KEY: mask, + "patches": {"attn1_patch": [krea2_attn1_negpip]}, + } + native_options = {**options, "block_index": 0, "img_slice": [2, 7]} + with torch.no_grad(): + expected = attention(x, transformer_options=native_options) + actual = Krea2HostAttention(attention, 0, 28)(x, transformer_options=options) + assert torch.equal(actual, expected) + assert "block_index" not in options + assert "img_slice" not in options + + +def test_host_attention_rejects_malformed_token_mask() -> None: + """Malformed sign metadata must fail before executing attention projections.""" + with pytest.raises(ValueError, match="processed token sign mask"): + Krea2HostAttention(None, 0, 28)( + torch.zeros(1, 3, 4), transformer_options={TRANSFORMER_MASK_KEY: "invalid"} + ) + + +def test_host_attention_delegates_unsigned_calls() -> None: + """Unmarked conditioning retains the host's ordinary attention operation.""" + torch.manual_seed(4) + attention = Attention(128, 4, kvheads=2, operations=torch.nn) + with torch.no_grad(): + attention.qknorm.qnorm.scale.zero_() + attention.qknorm.knorm.scale.zero_() + x = torch.randn(1, 7, 128) + expected = attention(x) + actual = Krea2HostAttention(attention, 0, 28)(x) + assert torch.equal(actual, expected) + + +class _EarlierModel(torch.nn.Module): + """Expose the earlier host boundary around one real Comfy attention module.""" + + def __init__(self) -> None: + """Retain a genuine patchable module hierarchy without loading weights.""" + super().__init__() + block = torch.nn.Module() + block.attn = Attention(128, 4, kvheads=2, operations=torch.nn) + self.blocks = torch.nn.ModuleList([block]) + + def _forward( + self, + x: object, + timesteps: object, + context: object, + attention_mask: object = None, + transformer_options: object = None, + **kwargs: object, + ) -> object: + """Represent only the external host argument contract under test.""" + del timesteps, context, attention_mask, transformer_options, kwargs + return x + + +def test_earlier_host_object_patches_restore_on_unload() -> None: + """Patch only a derived MODEL and restore the shared host method on unload.""" + root = torch.nn.Module() + root.diffusion_model = _EarlierModel() + device = torch.device("cpu") + source = ModelPatcher(root, load_device=device, offload_device=device) + attention = root.diffusion_model.blocks[0].attn + assert isinstance(attention, torch.nn.Module) + original = attention.forward + derived = PATCHER_LIFECYCLE.derive_model( + source, krea2_host_mutations(source), operation="test Krea host adaptation" + ) + assert source.object_patches == {} + assert attention.forward == original + derived.patch_model(load_weights=False) + try: + assert isinstance(attention.forward, Krea2HostAttention) + finally: + derived.unpatch_model(unpatch_weights=False) + assert attention.forward == original + + +def test_reference_capable_host_needs_no_object_replacements() -> None: + """Leave native Krea attention and its reference-image path entirely untouched.""" + root = torch.nn.Module() + diffusion = object.__new__(SingleStreamDiT) + torch.nn.Module.__init__(diffusion) + root.diffusion_model = diffusion + device = torch.device("cpu") + model = ModelPatcher(root, load_device=device, offload_device=device) + assert krea2_host_mutations(model) == () diff --git a/tests/prompting/negpip/test_negpip_model_service.py b/tests/prompting/negpip/test_negpip_model_service.py index 58e4e3c..8154242 100644 --- a/tests/prompting/negpip/test_negpip_model_service.py +++ b/tests/prompting/negpip/test_negpip_model_service.py @@ -11,6 +11,7 @@ from typing import Any, cast import pytest import torch +from comfy.ldm.krea2.model import SingleStreamDiT from comfy.model_base import SDXL, Anima, BaseModel, Krea2, SDXLRefiner from comfy.model_patcher import ModelPatcher from comfy.patcher_extension import WrappersMP @@ -161,6 +162,10 @@ def _model_patcher(model_class: type[BaseModel]) -> ModelPatcher: model = object.__new__(model_class) torch.nn.Module.__init__(model) + if model_class is Krea2: + diffusion = object.__new__(SingleStreamDiT) + torch.nn.Module.__init__(diffusion) + model.diffusion_model = diffusion device = torch.device("cpu") return ModelPatcher(model, load_device=device, offload_device=device)