nodes/capitan01R-ComfyUI-Flux2Klein-Enhancer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
capitan01R
2026-04-18 01:15:03 -04:00
committed by GitHub
parent 86394cac48
commit 3bc3dfc683
2 changed files with 261 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
identity_feature_transfer.py
+146
View File
@@ -0,0 +1,146 @@
"""Identity Feature Transfer — Attention Output Steering."""
import torch
import torch.nn.functional as F
class IdentityFeatureTransfer:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL", {
"tooltip": "Requires ReferenceLatent connected. The reference must be in the image stream.",
}),
"strength": ("FLOAT", {
"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "Per-block blend factor. Fires at every active block so the effect is cumulative. Start at 0.10 to 0.20.",
}),
"start_block": ("INT", {
"default": 0, "min": 0, "max": 23,
"tooltip": "First block index to apply. 0 = earliest. Index is shared across double and single blocks (resets when single blocks begin).",
}),
"end_block": ("INT", {
"default": 23, "min": 0, "max": 23,
"tooltip": "Last block index to apply. Covers 8 double blocks (0-7) then 24 single blocks (index resets 0-23). Higher values extend coverage into later single blocks.",
}),
"mode": (["cosine_pull", "topk_replace", "mean_transfer"], {
"default": "cosine_pull",
"tooltip": "cosine_pull: pulls each gen token toward its best-matching ref token. topk_replace: only affects the top K%% most similar tokens. mean_transfer: shifts overall feature distribution toward the reference.",
}),
"top_k_percent": ("FLOAT", {
"default": 0.25, "min": 0.01, "max": 1.0, "step": 0.01,
"tooltip": "topk_replace mode only. Fraction of generation tokens to affect. 0.25 = top 25%% most similar.",
}),
},
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "conditioning/flux2klein"
def apply(self, model, strength=0.15, start_block=0, end_block=23,
mode="cosine_pull", top_k_percent=0.25):
m = model.clone()
_strength = strength
_start = start_block
_end = end_block
_mode = mode
_topk_pct = top_k_percent
def output_patch(attn, extra_options):
ref_tokens_list = extra_options.get("reference_image_num_tokens", [])
if not ref_tokens_list:
return attn
block_idx = extra_options.get("block_index", 0)
if block_idx < _start or block_idx > _end:
return attn
img_slice = extra_options.get("img_slice", None)
if img_slice is None:
return attn
txt_end = img_slice[0]
total_seq = img_slice[1]
total_ref = sum(ref_tokens_list)
if total_ref <= 0:
return attn
gen_start = txt_end
gen_end = total_seq - total_ref
ref_start = total_seq - total_ref
ref_end = total_seq
if gen_end <= gen_start or ref_end <= ref_start:
return attn
gen_features = attn[:, gen_start:gen_end]
ref_features = attn[:, ref_start:ref_end]
if _mode == "cosine_pull":
gen_norm = F.normalize(gen_features.float(), dim=-1)
ref_norm = F.normalize(ref_features.float(), dim=-1)
sim = torch.bmm(gen_norm, ref_norm.transpose(1, 2))
max_sim, max_idx = sim.max(dim=-1)
weight = max_sim.clamp(0.0, 1.0) * _strength
weight = weight.unsqueeze(-1).to(attn.dtype)
max_idx_expanded = max_idx.unsqueeze(-1).expand(-1, -1, ref_features.shape[-1])
best_ref = torch.gather(ref_features, 1, max_idx_expanded)
new_gen = gen_features + (best_ref - gen_features) * weight
attn = attn.clone()
attn[:, gen_start:gen_end] = new_gen
elif _mode == "topk_replace":
gen_norm = F.normalize(gen_features.float(), dim=-1)
ref_norm = F.normalize(ref_features.float(), dim=-1)
sim = torch.bmm(gen_norm, ref_norm.transpose(1, 2))
max_sim, max_idx = sim.max(dim=-1)
k = max(1, int(gen_features.shape[1] * _topk_pct))
topk_vals, topk_indices = max_sim.topk(k, dim=-1)
max_idx_expanded = max_idx.unsqueeze(-1).expand(-1, -1, ref_features.shape[-1])
best_ref = torch.gather(ref_features, 1, max_idx_expanded)
attn = attn.clone()
for b in range(attn.shape[0]):
for i in range(k):
idx = topk_indices[b, i].item()
sim_val = topk_vals[b, i].item()
if sim_val > 0:
w = min(sim_val * _strength, 1.0)
pos = gen_start + idx
attn[b, pos] = (1.0 - w) * attn[b, pos] + w * best_ref[b, idx]
elif _mode == "mean_transfer":
gen_mean = gen_features.mean(dim=1, keepdim=True)
ref_mean = ref_features.mean(dim=1, keepdim=True)
delta = (ref_mean - gen_mean) * _strength
new_gen = gen_features + delta
attn = attn.clone()
attn[:, gen_start:gen_end] = new_gen
return attn
m.set_model_attn1_output_patch(output_patch)
return (m,)
NODE_CLASS_MAPPINGS = {
"IdentityFeatureTransfer": IdentityFeatureTransfer,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"IdentityFeatureTransfer": "FLUX.2 Klein Identity Feature Transfer",
}
identity_guidance.py
+115
View File
@@ -0,0 +1,115 @@
"""Identity Guidance — Sampling Loop Correction."""
import torch
import torch.nn.functional as F
class IdentityGuidance:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"identity_latent": ("LATENT", {
"tooltip": "VAE-encoded reference image at full resolution.",
}),
"strength": ("FLOAT", {
"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "How hard to pull toward the reference each step. 0.3 = move 30% of the distance.",
}),
"start_percent": ("FLOAT", {
"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "When to start correcting. 0.0 = beginning of denoising.",
}),
"end_percent": ("FLOAT", {
"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "When to stop correcting. 0.8 = last 20% runs freely for texture refinement.",
}),
"mode": (["adaptive", "direct", "channel_match"], {
"default": "adaptive",
"tooltip": "adaptive: pulls only where prediction resembles reference. direct: pulls everywhere equally. channel_match: matches color/feature statistics without copying spatial content.",
}),
},
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "conditioning/flux2klein"
def apply(self, model, identity_latent, strength=0.3,
start_percent=0.0, end_percent=0.8, mode="adaptive"):
m = model.clone()
ref = identity_latent["samples"]
_ref = ref
_strength = strength
_start = start_percent
_end = end_percent
_mode = mode
def post_cfg_fn(args):
denoised = args["denoised"]
sigma = args["sigma"]
s_now = float(sigma.flatten()[0])
progress = max(0.0, min(1.0, 1.0 - s_now))
if progress < _start or progress > _end:
return denoised
ref_resized = _ref.to(device=denoised.device, dtype=denoised.dtype)
if ref_resized.shape[0] != denoised.shape[0]:
ref_resized = ref_resized[:1].expand(denoised.shape[0], -1, -1, -1)
if ref_resized.shape[2:] != denoised.shape[2:]:
ref_resized = F.interpolate(
ref_resized, size=denoised.shape[2:],
mode="bilinear", align_corners=False,
)
if ref_resized.shape[1] != denoised.shape[1]:
if ref_resized.shape[1] > denoised.shape[1]:
ref_resized = ref_resized[:, :denoised.shape[1]]
else:
ref_resized = F.pad(ref_resized, (0, 0, 0, 0,
0, denoised.shape[1] - ref_resized.shape[1]))
if _mode == "direct":
correction = ref_resized - denoised
denoised = denoised + correction * _strength
elif _mode == "adaptive":
d_flat = denoised.flatten(2)
r_flat = ref_resized.flatten(2)
cos_sim = F.cosine_similarity(d_flat, r_flat, dim=1)
weight = cos_sim.clamp(0.0, 1.0)
weight = weight.unsqueeze(1)
weight = weight.view(denoised.shape[0], 1,
denoised.shape[2], denoised.shape[3])
correction = ref_resized - denoised
denoised = denoised + correction * weight * _strength
elif _mode == "channel_match":
ref_mean = ref_resized.mean(dim=(2, 3), keepdim=True)
ref_std = ref_resized.std(dim=(2, 3), keepdim=True).clamp(min=1e-5)
den_mean = denoised.mean(dim=(2, 3), keepdim=True)
den_std = denoised.std(dim=(2, 3), keepdim=True).clamp(min=1e-5)
normed = (denoised - den_mean) / den_std
matched = normed * ref_std + ref_mean
denoised = denoised + (matched - denoised) * _strength
return denoised
m.model_options.setdefault("sampler_post_cfg_function", []).append(post_cfg_fn)
return (m,)
NODE_CLASS_MAPPINGS = {
"IdentityGuidance": IdentityGuidance,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"IdentityGuidance": "FLUX.2 Klein Identity Guidance",
}