Krea2OstrisEditModelPatch object-patches diffusion_model.forward and drives block submodules directly (block.mod / block.attn / block.mlp), so the forward-wrapper swap never fires on the reference-latent paths. Every Krea2 block path begins with block.mod and ends with block.mlp, so this node triggers the swap with a forward_pre_hook on mod and a forward_hook on mlp. Verified on RTX 4090 / ComfyUI 0.26.2 with krea2_turbo_fp8 (no LoRA), blocks_to_swap=14: native, multi-reference and kv_cache paths all bit-exact vs no-swap, 6.4 GB resident instead of 12.9 GB. Runtime LoRA and bf16-scale checkpoints are documented as known-broken (CUDA invalid argument; shared bake-path suspect, see README). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TCjmaeny4KRS2TzXyYEgZS
9 lines
506 B
Python
9 lines
506 B
Python
from .blockswap import NODE_CLASS_MAPPINGS as _WAN_NODES, NODE_DISPLAY_NAME_MAPPINGS as _WAN_NAMES
|
|
from .blockswap_ltx import NODE_CLASS_MAPPINGS as _LTX_NODES, NODE_DISPLAY_NAME_MAPPINGS as _LTX_NAMES
|
|
from .blockswap_krea2 import NODE_CLASS_MAPPINGS as _K2_NODES, NODE_DISPLAY_NAME_MAPPINGS as _K2_NAMES
|
|
|
|
NODE_CLASS_MAPPINGS = {**_WAN_NODES, **_LTX_NODES, **_K2_NODES}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {**_WAN_NAMES, **_LTX_NAMES, **_K2_NAMES}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|