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.