nodes/januspluto-ComfyUI-Krea2-Regional
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Scheduled restrict window + proper handling of timestep-embedding LoRA keys

Browse Source
This commit is contained in:
januspluto
2026-07-27 23:11:33 +01:00
parent 3f6eaad5d2
commit 65efe53ead
4 changed files with 178 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+15
View File
@@ -2,6 +2,21 @@
## 1.1.0 (unreleased)
- **Scheduled restrict** (`restrict_end_percent` on Apply Regional): keep
`restrict_img_attn` isolation during the early identity-forming steps,
then release it so late steps integrate seams/lighting. Sigma-space
window; token masks and adaptive state survive the mid-run switch.
- Removed the `unmaskable_layers` widget (the capability remains in the
loader API; the node always skips timestep-embedding layers, with a clear
INFO log). If you saved a workflow with yesterday's dev build, re-check
the last widget on Apply Regional after updating.
- **`unmaskable_layers`** (Apply Regional): newer trainers also train Krea2's
timestep-embedding layers (`tmlp`/`tproj`); these modulate the whole image
and can't be region-masked. They're now classified properly (INFO, not a
scary warning) and can optionally be merged globally via ComfyUI's own
weight patching for extra likeness.
- Fix: compatibility with current ComfyUI master, which added a
`ref_latents` positional parameter to Krea2's forward — the wrapper is now
signature-agnostic, so future parameter insertions can't break it.
README.md
+5 -2
View File
@@ -29,7 +29,7 @@ Clone into `ComfyUI/custom_nodes/` and restart:
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/januspluto/ComfyUI-Krea2-Regional.git
git clone https://github.com/YOUR_USERNAME/ComfyUI-Krea2-Regional.git
```
Requires a ComfyUI build with native Krea 2 support (v0.26+). No extra Python
@@ -119,7 +119,10 @@ In order of impact:
3. **`exclusive_masks`** (Apply Regional, default on) — where grown/feathered
masks overlap, each token keeps only its strongest region.
4. **`restrict_img_attn`** (Apply Regional) — blocks image-to-image attention
across regions. Strongest lever; can look collaged at hard seams.
across regions. Strongest lever; can look collaged at hard seams. Pair it
with **`restrict_end_percent`** (~0.4-0.6) to get both: hard isolation
while each subject's identity forms, then an open second half where the
model integrates seams and lighting into one cohesive image.
5. **`base_loras_exclude_regions`** (Apply Regional) — style/base LoRAs
apply everywhere *except* inside subject regions, so a style LoRA can
skin the scene without contaminating character LoRAs.
krea2_regional.py
+87 -11
View File
@@ -98,6 +98,7 @@ def _normalize_lora_sd(sd: dict) -> dict[str, dict]:
continue
out = {}
for key, g in groups.items():
orig_key = key
entry = None
if "down" in g and "up" in g:
entry = {"type": "lora", "down": g["down"], "up": g["up"]}
@@ -131,6 +132,7 @@ def _normalize_lora_sd(sd: dict) -> dict[str, dict]:
"alpha_scale": alpha_scale}
if entry is None:
continue
entry["orig"] = orig_key
for p in _PREFIXES:
if key.startswith(p):
key = key[len(p):]
@@ -257,25 +259,46 @@ def _patched_linear(base: torch.nn.Module) -> torch.nn.Module:
return patched
def _inject_lora(model_patcher, lora_sd: dict, uid: str, strength: float) -> int:
"""Attach one LoRA (as adapter `uid`) onto the DiT block Linears."""
def _inject_lora(model_patcher, lora_sd: dict, uid: str, strength: float,
unmaskable: str = "skip") -> int:
"""Attach one LoRA (as adapter `uid`) onto the DiT block Linears.
Keys that hit layers OUTSIDE the token sequence (tmlp/tproj timestep
embedders, etc. — trained by newer ai-toolkit configs) can't be
region-masked: the timestep vector modulates the whole image through
AdaLN, so there is no per-token axis. `unmaskable`:
"skip" -> leave them out (default; logged as INFO)
"apply globally" -> merge them as normal weight patches via ComfyUI's
own LoRA loader (affects the WHOLE image)
"""
groups = _normalize_lora_sd(lora_sd)
diffusion_model = model_patcher.get_model_object("diffusion_model")
lookup: dict[str, str] = {}
full_lookup: dict[str, str] = {}
for name, mod in diffusion_model.named_modules():
if (name.startswith(("blocks.", "txtfusion.", "txtmlp", "first",
"last")) and hasattr(mod, "weight")
and callable(mod)):
if mod.weight is not None and mod.weight.ndim == 2:
if (hasattr(mod, "weight") and callable(mod)
and getattr(mod, "weight", None) is not None
and mod.weight.ndim == 2):
full_lookup[name] = name
full_lookup[name.replace(".", "_")] = name
if name.startswith(("blocks.", "txtfusion.", "txtmlp", "first",
"last")):
lookup[name] = name
lookup[name.replace(".", "_")] = name
patched, skipped = 0, []
global_keys: dict[str, str] = {}
for key, g in groups.items():
path = lookup.get(key) or lookup.get(key.replace(".", "_"))
if path is None:
skipped.append(key)
gpath = (full_lookup.get(key)
or full_lookup.get(key.replace(".", "_")))
if gpath is not None:
# real layer, just not a token-sequence one
global_keys[g["orig"]] = "diffusion_model." + gpath + ".weight"
else:
skipped.append(key)
continue
full = "diffusion_model." + path
target = _patched_linear(model_patcher.get_model_object(full))
@@ -294,10 +317,29 @@ def _inject_lora(model_patcher, lora_sd: dict, uid: str, strength: float) -> int
model_patcher.add_object_patch(full, target)
patched += 1
if global_keys:
if unmaskable == "apply globally":
import comfy.lora
patches = comfy.lora.load_lora(lora_sd, global_keys,
log_missing=False)
if patches:
model_patcher.add_patches(patches, float(strength))
logging.info(
"[Krea2Regional] lora '%s': %d timestep/embedding layers "
"applied GLOBALLY (whole image) — they can't be "
"region-masked.", uid, len(global_keys))
else:
logging.info(
"[Krea2Regional] lora '%s': %d timestep/embedding layers "
"skipped (e.g. %s) — they modulate the whole image and "
"can't be region-masked. Set unmaskable_layers to 'apply "
"globally' on Apply Regional to merge them.",
uid, len(global_keys), list(global_keys)[:3])
if skipped:
logging.warning(
"[Krea2Regional] lora '%s': %d keys didn't match any DiT block layer "
"(e.g. %s) — is this a Krea 2 LoRA?", uid, len(skipped), skipped[:3]
"[Krea2Regional] lora '%s': %d keys didn't match any layer "
"(e.g. %s) — is this a Krea 2 LoRA?", uid, len(skipped),
skipped[:3]
)
return patched
@@ -538,6 +580,17 @@ def _diffusion_wrapper(executor, x, timesteps, context, *args, **kwargs):
rt["call"] += 1
rt["blk"] = 0 # per-call joint-attention block counter (capture below)
# ---- scheduled restrict: hard image<->image isolation while identity
# forms, then release so the late steps can integrate seams/lighting.
# restrict_end_sigma < 0 means "never release" (classic behavior).
restrict_now = bool(bundle["restrict_img_attn"]) and (
t0 >= bundle.get("restrict_end_sigma", -1.0) - 1e-9)
if rt.get("restrict_now") != restrict_now:
# only the attention masks change: keep token masks + adaptive state
rt["restrict_now"] = restrict_now
rt["masks_built"] = False
rt["additive"] = {}
# ---- adaptive finalize: capture window over -> snap masks to the
# discovered subject silhouettes and rebuild everything once
capturing = (ad_mode != "off" and rt["refined"] is None
@@ -568,7 +621,7 @@ def _diffusion_wrapper(executor, x, timesteps, context, *args, **kwargs):
hard = soft > 0.5
allow, txt_allow = _build_allow(
segments, hard, imglen, bundle["restrict_img_attn"], device
segments, hard, imglen, rt.get("restrict_now", False), device
)
# When every row is regional (the usual turbo cfg=1 case), keep batch
# dim 1 so the mask broadcasts against ANY caller batch — including
@@ -891,6 +944,15 @@ class Krea2ApplyRegional:
"max": 1.0, "step": 0.05,
"tooltip": "When to release the lock so the model can "
"integrate seams and lighting."}),
"restrict_end_percent": ("FLOAT", {"default": 1.0,
"min": 0.05, "max": 1.0, "step": 0.05,
"tooltip": "When restrict_img_attn is on: release the "
"image<->image isolation at this fraction of "
"the schedule. Identity locks in during the "
"restricted window; the open steps afterwards "
"integrate seams and lighting. 1.0 = "
"restricted for the whole run (classic). "
"0.4-0.6 is a good cohesion sweet spot."}),
},
"optional": {"base_loras": ("KREA2_LORAS",)},
}
@@ -903,7 +965,8 @@ class Krea2ApplyRegional:
exclusive_masks=True, adaptive_masks="off", adaptive_steps=2,
adaptive_threshold=0.45, base_loras_exclude_regions=False,
region_lock_strength=0.0, region_lock_start=0.35,
region_lock_end=0.85, base_loras=None):
region_lock_end=0.85, restrict_end_percent=1.0,
unmaskable_layers="skip", base_loras=None):
if len(conditioning) != 1:
logging.warning("[Krea2Regional] base conditioning has %d entries; "
"using the first.", len(conditioning))
@@ -937,11 +1000,24 @@ class Krea2ApplyRegional:
lora_specs.append({"uid": uid, "weight": 1.0, "seg": ridx + 1,
"mask_idx": ridx})
# scheduled restrict: convert the percent to sigma space so it
# tracks the real noise schedule at any step count. >=0.999 means
# restricted for the whole run (sentinel -1: never release).
restrict_end_sigma = -1.0
if restrict_img_attn and restrict_end_percent < 0.999:
try:
ms = model.model.model_sampling
restrict_end_sigma = float(
ms.percent_to_sigma(restrict_end_percent))
except Exception:
restrict_end_sigma = -1.0 # unknown sampling: never release
bundle = {
"segments": segments,
"txt_total": off,
"masks": [r["mask"] for r in regions],
"restrict_img_attn": restrict_img_attn,
"restrict_end_sigma": restrict_end_sigma,
"exclusive_masks": exclusive_masks,
"base_loras_exclude_regions": base_loras_exclude_regions,
"adaptive": {
test_nodes.py
+71
View File
@@ -507,3 +507,74 @@ err21 = (y21 - exp21).abs().max().item()
assert err21 < 1e-4, f"multi-lokr stacking broken: {err21}"
patcher21.unpatch_model()
print(f"21) two LoKr + standard stacking (issue #1): ok (err {err21:.1e})")
# ---- 22. unmaskable layers (tmlp/tproj): classified, optional global apply
import comfy.model_patcher as _mp22
sd22 = {
"lora_unet_blocks_0_attn_wq.lora_down.weight": torch.randn(4, CFG["features"]) * 0.1,
"lora_unet_blocks_0_attn_wq.lora_up.weight": torch.randn(CFG["features"], 4) * 0.1,
# timestep-embedding layer: real module, NOT token-sequence
"lora_unet_tmlp_0.lora_down.weight": torch.randn(4, CFG["tdim"]) * 0.1,
"lora_unet_tmlp_0.lora_up.weight": None, # placeholder, set below
# genuinely unknown key: must still hit the warning path
"lora_unet_bogus_layer.lora_down.weight": torch.randn(4, 8),
"lora_unet_bogus_layer.lora_up.weight": torch.randn(8, 4),
}
tmlp0 = container.diffusion_model.tmlp[0]
sd22["lora_unet_tmlp_0.lora_up.weight"] = \
torch.randn(tmlp0.weight.shape[0], 4) * 0.1
sd22["lora_unet_tmlp_0.lora_down.weight"] = \
torch.randn(4, tmlp0.weight.shape[1]) * 0.1
# skip mode: 1 token-sequence layer patched, no weight patches added
p22a = _mp22.ModelPatcher(container, load_device="cpu", offload_device="cpu")
n22a = _inj(p22a, sd22, "unm_a", 1.0)
assert n22a == 1, n22a
assert not any("tmlp" in k for k in p22a.patches), "skip mode must not patch"
print("22a) unmaskable skip mode: ok")
# apply-globally mode: tmlp weight patch registered and changes the weight
p22b = _mp22.ModelPatcher(container, load_device="cpu", offload_device="cpu")
n22b = _inj(p22b, sd22, "unm_b", 1.0, unmaskable="apply globally")
key22 = "diffusion_model.tmlp.0.weight"
assert key22 in p22b.patches, list(p22b.patches)[:5]
w_before = tmlp0.weight.detach().clone()
p22b.patch_model(device_to="cpu")
w_after = container.diffusion_model.tmlp[0].weight.detach().clone()
delta22 = (w_after - w_before).abs().max().item()
assert delta22 > 1e-6, "global patch must change the tmlp weight"
p22b.unpatch_model()
assert torch.allclose(container.diffusion_model.tmlp[0].weight, w_before)
print(f"22b) unmaskable global mode merges weights: ok (|d|={delta22:.2e})")
# ---- 23. scheduled restrict: isolated early, released late, txt intact
b23 = dict(bundle(loras=lora_specs, restrict=True))
b23["restrict_end_sigma"] = 0.5 # release at sigma < 0.5
to23 = wrapped_opts(b23)
img_sl = slice(txtlen, txtlen + imglen)
with torch.no_grad():
run(container.diffusion_model, x, torch.full((B,), 0.9), context, to23)
j_early = b23["_rt"]["joint"][0, 0]
tm_early = b23["_rt"]["token_masks"]
assert not bool(j_early[img_sl, img_sl].all()), \
"early steps must be image-restricted"
with torch.no_grad():
run(container.diffusion_model, x, torch.full((B,), 0.4), context, to23)
j_late = b23["_rt"]["joint"][0, 0]
assert bool(j_late[img_sl, img_sl].all()), \
"after the window, image<->image opens up"
s0, e0 = b23["segments"][0]
s1, e1 = b23["segments"][1]
assert not bool(j_late[s0:e0, s1:e1].any()), \
"prompt isolation must persist after release"
assert b23["_rt"]["token_masks"] is tm_early, \
"token masks (lora gating) must survive the flip untouched"
# sentinel -1: never releases
b23b = dict(bundle(loras=lora_specs, restrict=True))
to23b = wrapped_opts(b23b)
with torch.no_grad():
run(container.diffusion_model, x, torch.full((B,), 0.9), context, to23b)
run(container.diffusion_model, x, torch.full((B,), 0.05), context, to23b)
assert not bool(b23b["_rt"]["joint"][0, 0][img_sl, img_sl].all()), \
"default (no end sigma) stays restricted for the whole run"
print("23) scheduled restrict window: ok")