From 678891c1cdec6ebad4ef413381a183a64f03bcd4 Mon Sep 17 00:00:00 2001 From: capitan01R Date: Mon, 26 Jan 2026 09:21:25 -0500 Subject: [PATCH] Update flux2_klein_enhancer.py --- flux2_klein_enhancer.py | 152 +++++++++++++++++++++++++++++++++------- 1 file changed, 126 insertions(+), 26 deletions(-) diff --git a/flux2_klein_enhancer.py b/flux2_klein_enhancer.py index cc34a1d..7901c26 100644 --- a/flux2_klein_enhancer.py +++ b/flux2_klein_enhancer.py @@ -1,15 +1,14 @@ """ -FLUX.2 Klein Conditioning Enhancer - Fixed Version +FLUX.2 Klein Conditioning Enhancer - v1.1 + Multiple Preserve Modes -Built from empirical diagnostic data: -- Conditioning shape: [batch, 512, 12288] -- Active text region: Dynamic (check attention_mask), typically ~77 tokens -- Padding region: Remaining positions (std ~2.3) -- Metadata: pooled_output, attention_mask, reference_latents (for edit) -- Text encoder: Qwen3 8B (4096 hidden → 12288 projected) +Added features: +- preserve_mode: Choose between different preservation behaviors + * linear: Original blend-back method (default, backward compatible) + * dampen: Reduce modification strength before applying + * blend_after: Same as linear (legacy naming) + * hybrid: Dampen first, then blend -CRITICAL FIX: Previous version had mean-recentering that undid the scaling. -This version uses direct operations with no reversion loops. +Main v1 logic: COMPLETELY UNCHANGED """ import torch @@ -30,6 +29,7 @@ class Flux2KleinEnhancer: - magnitude: Direct scaling of active region embeddings - contrast: Amplify differences between tokens - normalize: Equalize token magnitudes + - preserve_original: Blend back original embeddings with multiple modes All operations modify the active text region only (padding untouched). """ @@ -47,14 +47,14 @@ class Flux2KleinEnhancer: "magnitude": ("FLOAT", { "default": 1.0, "min": 0.0, - "max": 3.0, + "max": 10.0, "step": 0.05, "tooltip": "Scale active region embeddings. >1 = stronger prompt, <1 = weaker." }), "contrast": ("FLOAT", { "default": 0.0, "min": -1.0, - "max": 2.0, + "max": 10.0, "step": 0.05, "tooltip": "Amplify differences between tokens. >0 = sharper concepts." }), @@ -67,10 +67,21 @@ class Flux2KleinEnhancer: }), }, "optional": { + "preserve_original": ("FLOAT", { + "default": 0.0, + "min": -2.0, + "max": 10.0, + "step": 0.05, + "tooltip": "Blend back original embeddings. 0=full enhancement, 1=no change. For image edit: higher=keep original structure." + }), + "preserve_mode": (["linear", "dampen", "blend_after", "hybrid"], { + "default": "linear", + "tooltip": "linear: direct blend (original) | dampen: reduce modifications first | blend_after: same as linear | hybrid: dampen then blend" + }), "edit_text_weight": ("FLOAT", { "default": 1.0, "min": 0.0, - "max": 3.0, + "max": 10.0, "step": 0.05, "tooltip": "[IMAGE EDIT] Additional scaling for edit mode. <1 = preserve original, >1 = follow prompt." }), @@ -109,8 +120,8 @@ class Flux2KleinEnhancer: return min(77, cond_shape[1]) def enhance(self, conditioning, magnitude=1.0, contrast=0.0, normalize_strength=0.0, - edit_text_weight=1.0, active_end_override=0, low_vram=False, - device="auto", debug=False): + preserve_original=0.0, preserve_mode="linear", edit_text_weight=1.0, + active_end_override=0, low_vram=False, device="auto", debug=False): if not conditioning: return (conditioning,) @@ -120,6 +131,7 @@ class Flux2KleinEnhancer: magnitude == 1.0 and contrast == 0.0 and normalize_strength == 0.0 and + preserve_original == 0.0 and edit_text_weight == 1.0 ) if no_op: @@ -163,18 +175,46 @@ class Flux2KleinEnhancer: print(f"Shape: {cond.shape}") print(f"Active region: 0 to {active_end}") print(f"Edit mode: {is_edit_mode}") + print(f"Preserve mode: {preserve_mode}") if active_end < seq_len: print(f"Active std: {cond[:, :active_end, :].std().item():.4f}") print(f"Padding std: {cond[:, active_end:, :].std().item():.4f}") # Extract active region only active = cond[:, :active_end, :].clone() + original_active = active.clone() # Store for preserve_original if debug: orig_norm = active.norm(dim=-1).mean().item() print(f"\nBefore modifications:") print(f" Active region mean norm: {orig_norm:.4f}") + # ============================================================ + # DAMPEN MODE: Apply preservation BEFORE modifications + # ============================================================ + if preserve_mode == "dampen" and preserve_original != 0.0: + # Calculate dampening factors for all operations + damping = 1.0 - preserve_original + + # Dampen the modification parameters + dampened_magnitude = 1.0 + (magnitude - 1.0) * damping + dampened_contrast = contrast * damping + dampened_normalize = normalize_strength * damping + dampened_edit_weight = 1.0 + (edit_text_weight - 1.0) * damping + + if debug: + print(f"\nDampen mode ({preserve_original:.2f}):") + print(f" magnitude: {magnitude:.2f} -> {dampened_magnitude:.3f}") + print(f" contrast: {contrast:.2f} -> {dampened_contrast:.3f}") + print(f" normalize: {normalize_strength:.2f} -> {dampened_normalize:.3f}") + print(f" edit_weight: {edit_text_weight:.2f} -> {dampened_edit_weight:.3f}") + + # Use dampened values + magnitude = dampened_magnitude + contrast = dampened_contrast + normalize_strength = dampened_normalize + edit_text_weight = dampened_edit_weight + # ============================================================ # OPERATION 1: Contrast (amplify token differences) # ============================================================ @@ -216,6 +256,21 @@ class Flux2KleinEnhancer: if debug: print(f"\nEdit text weight ({edit_text_weight:.2f}): applied for image edit mode") + # ============================================================ + # OPERATION 5: Preserve Original (blend back) + # ============================================================ + if preserve_original != 0.0 and preserve_mode in ["linear", "blend_after", "hybrid"]: + # Linear interpolation: result = enhanced * (1-p) + original * p + # p=0: fully enhanced + # p=1: original unchanged + # p>1: over-preserve (can create interesting effects) + active = active * (1.0 - preserve_original) + original_active * preserve_original + + if debug: + blend_pct = min(preserve_original * 100, 100) + mode_name = "Hybrid blend" if preserve_mode == "hybrid" else "Linear blend" + print(f"\n{mode_name} ({preserve_original:.2f}): {blend_pct:.0f}% original blended back") + # Write back to full tensor result = cond.clone() result[:, :active_end, :] = active @@ -229,7 +284,7 @@ class Flux2KleinEnhancer: output.append((result.to("cpu", dtype=original_dtype), meta)) - del cond, active, result + del cond, active, original_active, result if device.type == "cuda": torch.cuda.empty_cache() @@ -263,21 +318,21 @@ class Flux2KleinDetailController: "front_mult": ("FLOAT", { "default": 1.0, "min": 0.0, - "max": 3.0, + "max": 10.0, "step": 0.05, "tooltip": "First 25% of active tokens" }), "mid_mult": ("FLOAT", { "default": 1.0, "min": 0.0, - "max": 3.0, + "max": 10.0, "step": 0.05, "tooltip": "Middle 50% of active tokens" }), "end_mult": ("FLOAT", { "default": 1.0, "min": 0.0, - "max": 3.0, + "max": 10.0, "step": 0.05, "tooltip": "Last 25% of active tokens" }), @@ -298,10 +353,21 @@ class Flux2KleinDetailController: "emphasis_mult": ("FLOAT", { "default": 1.0, "min": 0.0, - "max": 3.0, + "max": 10.0, "step": 0.1, "tooltip": "Multiplier for emphasis region" }), + "preserve_original": ("FLOAT", { + "default": 0.0, + "min": -2.0, + "max": 10.0, + "step": 0.05, + "tooltip": "Blend back original embeddings. 0=full enhancement, 1=no change. For image edit: higher=keep original structure." + }), + "preserve_mode": (["linear", "dampen", "blend_after", "hybrid"], { + "default": "linear", + "tooltip": "linear: direct blend (original) | dampen: reduce modifications first | blend_after: same as linear | hybrid: dampen then blend" + }), "low_vram": ("BOOLEAN", {"default": False}), "device": (devices, {"default": "auto"}), "debug": ("BOOLEAN", {"default": False}), @@ -314,7 +380,8 @@ class Flux2KleinDetailController: def control(self, conditioning, front_mult=1.0, mid_mult=1.0, end_mult=1.0, emphasis_start=0, emphasis_end=0, emphasis_mult=1.0, - low_vram=False, device="auto", debug=False): + preserve_original=0.0, preserve_mode="linear", low_vram=False, + device="auto", debug=False): if not conditioning: return (conditioning,) @@ -324,7 +391,8 @@ class Flux2KleinDetailController: front_mult == 1.0 and mid_mult == 1.0 and end_mult == 1.0 and - (emphasis_end == 0 or emphasis_mult == 1.0) + (emphasis_end == 0 or emphasis_mult == 1.0) and + preserve_original == 0.0 ) if no_op: if debug: @@ -363,6 +431,7 @@ class Flux2KleinDetailController: active_end = min(77, seq_len) active = cond[:, :active_end, :].clone() + original_active = active.clone() # Store for preserve_original num_tokens = active.shape[1] if debug: @@ -370,6 +439,26 @@ class Flux2KleinDetailController: print(f"Flux2KleinDetailController Item {idx}") print(f"{'='*50}") print(f"Active tokens: {num_tokens}") + print(f"Preserve mode: {preserve_mode}") + + # ============================================================ + # DAMPEN MODE: Apply preservation BEFORE regional operations + # ============================================================ + if preserve_mode == "dampen" and preserve_original != 0.0: + # Dampen the regional multipliers based on preserve_original + damping = 1.0 - preserve_original + + front_mult = 1.0 + (front_mult - 1.0) * damping + mid_mult = 1.0 + (mid_mult - 1.0) * damping + end_mult = 1.0 + (end_mult - 1.0) * damping + emphasis_mult = 1.0 + (emphasis_mult - 1.0) * damping + + if debug: + print(f"\nDampen mode ({preserve_original:.2f}):") + print(f" front_mult: -> {front_mult:.3f}") + print(f" mid_mult: -> {mid_mult:.3f}") + print(f" end_mult: -> {end_mult:.3f}") + print(f" emphasis_mult: -> {emphasis_mult:.3f}") # Calculate region boundaries front_end = int(num_tokens * 0.25) @@ -382,17 +471,17 @@ class Flux2KleinDetailController: if front_mult != 1.0 and front_end > 0: active[:, :front_end, :] *= front_mult if debug: - print(f" Front: x{front_mult}") + print(f" Front: x{front_mult:.3f}") if mid_mult != 1.0 and mid_end > front_end: active[:, front_end:mid_end, :] *= mid_mult if debug: - print(f" Mid: x{mid_mult}") + print(f" Mid: x{mid_mult:.3f}") if end_mult != 1.0 and num_tokens > mid_end: active[:, mid_end:, :] *= end_mult if debug: - print(f" End: x{end_mult}") + print(f" End: x{end_mult:.3f}") # Custom emphasis region if emphasis_end > 0 and emphasis_mult != 1.0: @@ -400,7 +489,18 @@ class Flux2KleinDetailController: emp_end = max(emp_start + 1, min(emphasis_end, num_tokens)) active[:, emp_start:emp_end, :] *= emphasis_mult if debug: - print(f" Emphasis [{emp_start}:{emp_end}]: x{emphasis_mult}") + print(f" Emphasis [{emp_start}:{emp_end}]: x{emphasis_mult:.3f}") + + # ============================================================ + # Preserve Original (blend back) + # ============================================================ + if preserve_original != 0.0 and preserve_mode in ["linear", "blend_after", "hybrid"]: + active = active * (1.0 - preserve_original) + original_active * preserve_original + + if debug: + blend_pct = min(preserve_original * 100, 100) + mode_name = "Hybrid blend" if preserve_mode == "hybrid" else "Linear blend" + print(f"\n{mode_name} ({preserve_original:.2f}): {blend_pct:.0f}% original blended back") # Write back result = cond.clone() @@ -412,7 +512,7 @@ class Flux2KleinDetailController: output.append((result.to("cpu", dtype=original_dtype), meta)) - del cond, active, result + del cond, active, original_active, result if device.type == "cuda": torch.cuda.empty_cache()