diff --git a/identity_feature_transfer.py b/identity_feature_transfer.py new file mode 100644 index 0000000..4d2b182 --- /dev/null +++ b/identity_feature_transfer.py @@ -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", +} diff --git a/identity_guidance.py b/identity_guidance.py new file mode 100644 index 0000000..5f253e3 --- /dev/null +++ b/identity_guidance.py @@ -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", +}