ComfyUI-Krea2-strip-layers

ComfyUI-Krea2-strip-layers

ComfyUI nodes that load a Krea2 LoRA, strip selected layer groups from it in memory, and apply the remaining weights to the diffusion model. A simple all-or-nothing node and an advanced node with per-block selection / weighting.

In practice it is useful to be able to use only some layers from a NSFW LoRA for it to combine better with a character LoRA.

Installation

Clone this repository to your ComfyUI custom_nodes directory.

Example workflow

An example workflow lives in workflow/workflow.webp. You can download and drag and drop it directly onto the ComfyUI canvas to load the graph.

Nodes

Both nodes are in category loaders/krea2 and output a MODEL.

Load LoRA (Krea2 Strip Layers)

The simple, all-or-nothing node.

| Input | Type | Default | Description | |-------|------|---------|-------------| | model | MODEL | — | The diffusion model the stripped LoRA is applied to. | | lora_name | dropdown | — | The Krea2 LoRA to strip and apply. | | strength_model | FLOAT | 1.0 | How strongly to apply the stripped LoRA. | | strip_dit_layers | BOOLEAN | True | Strip the DiT: all transformer blocks + first/last layers. | | strip_time_mlp | BOOLEAN | True | Strip time MLP / timestep projection: tmlp_, tproj_1 | | strip_txtfusion_layerwise | BOOLEAN | False | Strip txtfusion layerwise blocks (full censor) | | strip_txtfusion_projector | BOOLEAN | False | Strip txtfusion projector (zero impact) | | strip_txtfusion_refiner | BOOLEAN | False | Strip txtfusion refiner (tiny impact / full censor) |

Load LoRA Advanced (Krea2 Strip Layers)

Adds per-block selection and per-block weighting.

| Input | Type | Default | Description | |-------|------|---------|-------------| | model | MODEL | — | The diffusion model the LoRA is applied to. | | lora_name | dropdown | — | The Krea2 LoRA to strip and apply. | | strength_model | FLOAT | 1.0 | Global LoRA strength. Per-block weights (if enabled) multiply this. | | strip_all_blocks | BOOLEAN | True | Strip all transformer blocks. Ignored when use_block_selection or use_block_weights is ON. | | use_block_selection | BOOLEAN | False | Strip only the blocks in block_selection (ignores strip_all_blocks). Ignored when use_block_weights is ON. | | block_selection | STRING | 0-27 | Blocks to strip when use_block_selection is ON. Range/list like 0-27, 10-27, 5,8,20-27. Empty = strip none. | | use_block_weights | BOOLEAN | False | Per-block weight mode: scale each block by block_weights instead of stripping (weight 0 = block off). Overrides the block strip toggles above. | | block_weights | STRING | 1 | Per-block weights (see below). | | strip_first_last | BOOLEAN | True | Strip DiT input/output layers: first, last_linear, last.linear | | strip_time_mlp | BOOLEAN | True | Strip time MLP / timestep projection: tmlp_, tproj_1 | | strip_txtfusion_layerwise | BOOLEAN | False | Strip txtfusion layerwise blocks (full censor) | | strip_txtfusion_projector | BOOLEAN | False | Strip txtfusion projector (zero impact) | | strip_txtfusion_refiner | BOOLEAN | False | Strip txtfusion refiner (tiny impact / full censor) |

How it works

Each enabled toggle maps to a set of key prefixes; any state-dict entry whose key starts with one of them is dropped before the LoRA is applied (via ComfyUI's comfy.sd.load_lora_for_models, model-only path). The block selection / weight controls below belong to the Advanced node.

The console logs how many keys were removed vs kept. If 0 keys are removed it warns and prints the LoRA's top-level prefixes, so you can tell it wasn't matched.

Supported key conventions

The node handles several key conventions for blocks and txtfusion:

| | ComfyUI-native | diffusers | |---|---|---| | top-level | diffusion_model. | transformer. | | blocks | blocks.N | transformer_blocks.N | | txtfusion | txtfusion. | text_fusion. | | attn / mlp | attn.wq…, mlp.… | attn.to_q…, ff.… |

Use tools/inspect_lora.py to see which convention (and math format — LoRA / LoKr / LoHa) a file uses:

python tools/inspect_lora.py path/to/lora.safetensors

Note: strip_first_last and strip_time_mlp currently match only the diffusion_model. names; on a transformer.-format LoRA they're no-ops. Block stripping/weighting and the txtfusion toggles work on both.

Selecting blocks

By default strip_all_blocks removes every transformer block (all-or-nothing). Turn on use_block_selection to instead strip only the blocks listed in block_selection — a Krea2 model has 28 blocks (0–27). Since the field lists blocks to remove, to keep only a subset strip its complement: e.g. to keep only the first 8 blocks of an NSFW LoRA so it blends with a character LoRA, set block_selection to 8-27. Empty means strip no blocks.

This is useful to try to keep some concepts from a LoRA, while trying to keep some consistency from the original model or another LoRA.

Per-block weights

Turn on use_block_weights for continuous control: each block's LoRA contribution is scaled by strength_model × block_weight, so 1.0 = full, 0.5 = half, 0.0 = off (equivalent to stripping). This overrides the block strip toggles. block_weights accepts three forms:

| Form | Example | Meaning | |------|---------|---------| | single value | 0.5 | every block at ×0.5 | | positional list | 1,1,0.5,0 | block 0=1, 1=1, 2=0.5, 3=0; remaining blocks default to 1.0 | | labeled overrides | 12:0.5, 20:0 | only blocks 12 and 20 changed; everything else 1.0 |

Weights may be negative and are applied via ComfyUI's own LoRA patch mechanism (one add_patches per distinct weight), so any format (LoRA / LoKr / LoHa) is reconstructed correctly. Non-block groups (txtfusion, first/last, time-MLP) are still governed by their strip toggles.

txtfusion / NSFW notes

The txtfusion groups act as content censors in the Krea 2 model. The default node settings are to apply the LoRA's version of the groups by default. Here are what each of them do from empirical testing:

• layerwise — full censor
• projector — zero impact
• refiner — tiny impact / full censor

In other words for layerwise and refiner:

• Toggle off has an uncensor effect if the LoRA you're applying has NSFW capabilities.
• Toggle on removes the groups, so you're going back to the diffusion model's original censoring.

Standalone script (tools/strip_lora.py)

If you'd rather strip a LoRA to a new file on disk (instead of in memory at apply time), use the standalone CLI. It reuses the same key conventions and block logic as the nodes, so the result is identical — just written out.

# Default: strip the whole DiT (all blocks + first/last) and the time MLP,
# keep the txtfusion groups. Writes <name>.stripped.safetensors next to it.
python tools/strip_lora.py input.safetensors

# Explicit output path.
python tools/strip_lora.py input.safetensors -o out.safetensors

# Report what would be stripped without writing anything.
python tools/strip_lora.py input.safetensors --dry-run

# Keep the DiT, only strip the txtfusion censor groups.
python tools/strip_lora.py input.safetensors --no-dit --no-time-mlp \
    --txtfusion-layerwise --txtfusion-refiner

# Strip only a subset of transformer blocks (keep blocks 0-7 -> strip 8-27).
python tools/strip_lora.py input.safetensors --no-dit --blocks 8-27

| Flag | Default | Description | |------|---------|-------------| | input | — | Input .safetensors LoRA (positional). | | -o, --output | <input>.stripped.safetensors | Output path. | | --dry-run | off | Report removed/kept counts, write nothing. | | --dit / --no-dit | on | Strip / keep all transformer blocks. | | --blocks SPEC | — | Strip only these block indices (e.g. 8-27, 5,8,20-27); overrides --dit. | | --first-last / --no-first-last | follows --dit | Strip / keep the first/last I/O layers. | | --time-mlp / --no-time-mlp | on | Strip / keep the time MLP (tmlp_, tproj_1). | | --txtfusion-layerwise | off | Strip txtfusion layerwise blocks (full censor). | | --txtfusion-projector | off | Strip txtfusion projector (zero impact). | | --txtfusion-refiner | off | Strip txtfusion refiner (tiny impact / full censor). |

Requires the safetensors package (already present in a ComfyUI environment): pip install safetensors. LoRA metadata is preserved. Like the nodes, it warns if 0 keys matched (wrong key convention) and prints the top-level prefixes.

Credits

The layer stripping code's original author is Puppet_Master as seen on https://civitai.red/models/2742336/nsfw-krea2-low-vram.

License

MIT — see LICENSE.