OzzyGT/krea2_lora_mask_blocks

Krea 2 masked-LoRA modular blocks

Custom Modular Diffusers blocks that give each loaded LoRA its own region of the image. Two LoRAs that address the same thing -- two character LoRAs, most often -- blend wherever both apply, so each subject ends up carrying traces of the other. lora_masks maps an adapter name to the region it may write in, and everything outside that region stays on the base model.

Mixing kinds is not the problem this solves: a character LoRA and a style LoRA usually compose fine on their own. It is same-kind pairs that need separating.

lora_masks={}                                       -> every adapter applies everywhere, as normal
lora_masks={"alice": alice_mask}                    -> "alice" only where its mask is white
lora_masks={"alice": alice_mask, "bob": bob_mask}   -> one region per character

The keys are adapter names -- whatever you passed as adapter_name to load_lora_weights -- and the values are mask images. Only the image tokens are masked; the text tokens always see every adapter unmasked.

Both blocksets are the stock Krea 2 ones with the denoise loop swapped, so with no lora_masks they generate exactly what the stock pipeline does:

Loading & running

import sdnq  # needed to load the quantized text encoder
import torch

from diffusers import ModularPipeline
from diffusers.utils import load_image


pipe = ModularPipeline.from_pretrained(
    "OzzyGT/krea2_lora_mask_blocks", trust_remote_code=True
)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")

pipe.load_lora_weights("path/to/alice.safetensors", adapter_name="alice")
pipe.load_lora_weights("path/to/bob.safetensors", adapter_name="bob")
pipe.set_adapters(["alice", "bob"], [1.0, 1.0])

# white = where that LoRA applies
alice_mask = load_image("alice_region.png")
bob_mask = load_image("bob_region.png")

image = pipe(
    prompt="alice and bob standing side by side in a ruined cathedral",
    height=1024,
    width=1024,
    lora_masks={"alice": alice_mask, "bob": bob_mask},
    output="images",
)[0]
image.save("masked_lora.png")

That runs Krea 2 Turbo. The blocks also work with the raw model and CFG, as Krea2LoraMaskAutoBlocks.

For loading a different checkpoint or swapping components, see Modular pipeline in the diffusers docs.

Masks

White is where the LoRA applies, black leaves the base model alone, greys fade it -- the same reading as a diffusers inpainting mask_image, where white marks the region being acted on. Values stay continuous, so a feathered brush gives a soft edge.

One image per adapter, grayscale or RGB at any resolution at the same aspect ratio of the output.

Reusing the mechanism

lora_spatial_mask.py knows nothing about Krea. A family plugs in by mixing LoraSpatialMaskDenoiseMixin into its denoise loop wrapper and declaring where the image tokens sit:

class Flux2MaskedDenoiseStep(LoraSpatialMaskDenoiseMixin, Flux2DenoiseLoopWrapper):
    block_classes = [...]
    block_names = [...]
    image_tokens_last = False  # Flux2 packs [image | text]
    exclude_name_substrings = ()  # projections that never see image tokens

The apply_lora_spatial_masks / remove_lora_spatial_masks pair is also usable outside modular pipelines, on any transformer with PEFT adapters loaded.

Not covered: LoKr adapters, and models whose LoRA layers see a (sequence, batch, feature) layout rather than batch-first.