diff --git a/CHANGELOG.md b/CHANGELOG.md index 5dc814b..d463bc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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. diff --git a/README.md b/README.md index fbbb43b..3a689aa 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/krea2_regional.py b/krea2_regional.py index fa62722..af47f91 100644 --- a/krea2_regional.py +++ b/krea2_regional.py @@ -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": { diff --git a/test_nodes.py b/test_nodes.py index 79c52f0..724b502 100644 --- a/test_nodes.py +++ b/test_nodes.py @@ -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")