Files
kongo_junandClaude Fable 5 f25e3d6fa4 Add BlockSwapKrea2: hook-triggered swap for Krea 2 (Ostris Edit compatible)
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
2026-07-13 20:30:48 +09:00

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"]