NSFW Krea2 (Low VRAM)

Based on good NSFW Lora/LoKr, I've experimented with reduced size/weight: dropping the DIT/UNET layer part because the model already know how to do NSFW. keeping only the text encoder/fusion layers from the NSFW lora. This technique can reduce Krea2 lora weight (reduce ANY lora/lokr weight) exemple with this lora 218mo => 13.5mo exemple with LoKr 1.46go => 40mo or even 20mo in some case I simply use the safetensors python lib and keep only the txtfusion.layerwise_blocks and txtfusion.refiner layers from NSFW lora. Dropping the other layers as the model know how to draw nsfw stuff ;). I wanted to experiment with the Krea2 model and understand two things: - If the model was trainned on NSFW (I think it was or it couldn't draw NSFW without UNET/DIT layer part) - Where is the NSFW influence in the model: txtfusion.layerwise_blocks seems to have 80% of the NSFW decision making and txtfusion.refiner 20% ish. txtfusion.projector is a small vector and I think it should not be touched too much (it merge the 12 big vector from qwen3vl) but doesn't directly influence NSFW. Detail and code at the end. detail of the token handling before DIT as I think I understand it (maybe wrong): ''' Frozen Qwen3-VL-4B (Text Encoder) (hidden dim ~2560) | | Select 12 specific layers | [2,5,8,11,14,17,20,23,26,29,32,35] v 12 x (Seq_len, 2560) hidden states | | Stack / Concat along layer dim v [12, Seq_len, 2560] ← multilayer input | +---------------------------------------------------+ | txtfusion.layerwise_blocks | | (2 blocks, cross-layer attn) | +---------------------------------------------------+ | Block 0: - prenorm (RMSNorm scale [2560]) - attn: wq/wk/wv/wo/gate [2560,2560] + qknorm - mlp: down[2560→6912] / gate[6912,2560] / up[6912,2560] - postnorm (RMSNorm scale [2560]) | Block 1: (same structure as Block 0) | v Fused layer-wise features [Seq_len, 2560] | +---------------------------------------------------+ | txtfusion.projector | | Linear(12 → 1) weight [1, 12] | +---------------------------------------------------+ | v Single fused representation per token (or pooled/processed) [Seq_len, 2560] | +---------------------------------------------------+ | txtfusion.refiner_blocks | | (2 bidirectional transformer blocks)| +---------------------------------------------------+ | Refiner Block 0: - prenorm (RMSNorm [2560]) - attn: wq/wk/wv/wo/gate [2560,2560] + qknorm - mlp: down/gate/up (same as above) - postnorm | Refiner Block 1: (identical structure) | v Refined bidirectional text features [Seq_len, 2560] | +---------------------------------------------------+ | txtmlp | | (final text MLP / projection) | +---------------------------------------------------+ | txtmlp.0.scale [2560] ← residual / scaling | txtmlp.1: Linear [2560 → 6144] + bias | txtmlp.3: Linear [6144 → 6144] + bias | v Final conditioned text embeddings (ready for cross-attention into DiT) | ↓ Injected into first DiT / unet_blocks (via cross-attention in the main transformer) Code for the layer stripping: from safetensors import safe_open from safetensors.torch import save_file from pathlib import Path def strip_layers(input_file, output_file, remove_prefixes): tensors = {} with safe_open(input_file, framework="pt", device="cpu") as f: metadata = f.metadata() for key in f.keys(): if any(key.startswith(prefix) for prefix in remove_prefixes): print(f"Removing {key}") continue tensors[key] = f.get_tensor(key) save_file(tensors, output_file, metadata=metadata) inputFile = "i:/ComfyUI_windows_portable/ComfyUI/models/loras/krea2/KNPV3_1.safetensors" stripLayer = [ # linked to the DIT/UNET 'blocks', 'first', 'last_linear', 'last.linear' # liked to unet/DIT time MLP and timestep projection 'tmlp_', 'tproj_1', # keep commented for NSFW #'txtfusion.layerwise_blocks', 'txtfusion_layerwise_blocks', # full censor ] prefix = 'diffusion_model.' stripLayer = [prefix+x for x in stripLayer] print('layer to strip: ', stripLayer) outPutFile = Path(inputFile).name strip_layers( inputFile, outPutFile, stripLayer )