Update flux2_klein_enhancer.py
This commit is contained in:
capitan01R
+126
-26
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user