Conditioning enhancement and reference latent control for FLUX.2 Klein 9B in ComfyUI. Built from empirical analysis and forward-pass tracing of the model's dual-stream architecture.

What This Does (( NEW DISCOVERY!! ))

FLUX.2 Klein uses a Qwen3 8B text encoder that outputs conditioning tensors of shape [batch, 512, 12288]. Through diagnostic analysis and model hook tracing, I verified:

Text Conditioning: [1, 512, 12288] tensor with ~67 active tokens (auto-detected from attention mask)
Reference Latent: Stored separately in metadata as [1, 128, H, W] - NOT merged into text conditioning
Dual-Stream Architecture: Text and image streams are processed separately through double_blocks, then concatenated for single_blocks

Verified Architecture (from forward-pass hooks)

┌─────────────────────────────────────────────────────────────────┐
│                    FLUX.2 Klein Forward Pass                    │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  Reference Latent [1,128,55,74]  →  patchify  →  [1, 4070, 128] │
│  Noisy Latent [1,128,55,74]      →  patchify  →  [1, 4070, 128] │
│                                         │                       │
│                                    CONCATENATE                  │
│                                         ↓                       │
│                                  [1, 8140, 128]                 │
│                                         │                       │
│                                      img_in                     │
│                                         ↓                       │
│                                  [1, 8140, 4096]                │
│                                                                 │
│  Text Conditioning [1,512,12288] → txt_in → [1, 512, 4096]      │
│                                                                 │
├─────────────────────────────────────────────────────────────────┤
│  DOUBLE BLOCKS (×8): Separate streams                           │
│    img_stream: [1, 8140, 4096]                                  │
│    txt_stream: [1, 512, 4096]                                   │
├─────────────────────────────────────────────────────────────────┤
│  SINGLE BLOCKS (×24): Concatenated                              │
│    combined: [1, 8652, 4096]  (8140 + 512)                      │
└─────────────────────────────────────────────────────────────────┘

Without reference latent: img_in receives [1, 4070, 128], single_blocks process [1, 4582, 4096]

This means:

Text enhancement modifies the txt_stream input
Reference control modifies half of the img_stream input (the reference portion)
These are independent controls that can be combined

Installation

Navigate to your ComfyUI custom nodes folder:
```
cd ComfyUI/custom_nodes/
```

Clone this repository:

git clone https://github.com/capitan01R/ComfyUI-Flux2Klein-Enhancer.git

Restart ComfyUI

Nodes

FLUX.2 Klein Enhancer

General-purpose text conditioning enhancement for both text-to-image and image editing workflows.

Parameters

Parameter	Default	Range	Description
`magnitude`	1.0	0.0 to 3.0	Direct scaling of active region embeddings. Values above 1.0 increase prompt influence, below 1.0 decrease it.
`contrast`	0.0	-1.0 to 2.0	Amplifies differences between tokens. Positive values sharpen concept separation, negative values blend them.
`normalize_strength`	0.0	0.0 to 1.0	Equalizes token magnitudes. Higher values balance emphasis across all tokens in the prompt.
`edit_text_weight`	1.0	0.0 to 3.0	Image edit mode only. Values below 1.0 preserve more of the original image, above 1.0 follows the prompt more strongly.
`active_end_override`	0	0 to 512	Manual override for active region end. 0 = auto-detect from attention mask.
`low_vram`	False	True/False	Use float16 computation on CUDA devices.
`device`	auto	auto/cpu/cuda:N	Compute device selection.
`debug`	False	True/False	Prints tensor statistics and modification details to console.

FLUX.2 Klein Detail Controller

Regional control over prompt conditioning. Divides active tokens into front/mid/end sections.

Parameters

Parameter	Default	Range	Description
`front_mult`	1.0	0.0 to 3.0	Multiplier for first 25% of active tokens (typically subject/main concept).
`mid_mult`	1.0	0.0 to 3.0	Multiplier for middle 50% of active tokens (typically details/modifiers).
`end_mult`	1.0	0.0 to 3.0	Multiplier for last 25% of active tokens (typically style/quality terms).
`emphasis_start`	0	0 to 200	Start position of custom emphasis region.
`emphasis_end`	0	0 to 200	End position of custom emphasis region. 0 = disabled.
`emphasis_mult`	1.0	0.0 to 3.0	Multiplier for the custom emphasis region.
`low_vram`	False	True/False	Use float16 computation on CUDA devices.
`device`	auto	auto/cpu/cuda:N	Compute device selection.
`debug`	False	True/False	Prints debug information to console.

FLUX.2 Klein Ref Latent Controller

Controls how strongly a specific reference image influences the generation. Requires a MODEL input and returns an updated MODEL. Chain multiple nodes to control each reference independently.

Parameters

Parameter	Default	Range	Description
`strength`	1.0	0.0 to 1000.0	Reference attention strength. 0 = reference ignored, 1 = normal, >1 = stronger structure.
`reference_index`	0	0 to 7	Which reference image to control (0 = first).
`spatial_fade`	none	none/center_out/edges_out/top_down/left_right	Per-token spatial gradient applied to the strength.
`spatial_fade_strength`	0.5	0.0 to 1.0	Intensity of the spatial fade.
`debug`	False	True/False	Prints block index, token range, and strength to console.

FLUX.2 Klein Text/Ref Balance

Single slider to balance text conditioning vs. all reference images. Requires a MODEL input and returns an updated MODEL.

Parameters

Parameter	Default	Range	Description
`balance`	0.005	0.000 to 1.000	0 = reference only, 0.5 = balanced, 1 = text only.
`debug`	False	True/False	Prints text and ref scale factors per block to console.

Preserve Original - Solving FLUX Klein's Preservation Problem

FLUX Klein has a consistency problem. Sometimes it nails the preservation of subjects and objects. Sometimes it completely ignores what you're trying to keep and generates something else entirely. There was no native way to control this.

This node exposes preservation control that FLUX Klein doesn't provide. You can now control exactly how much original structure is maintained versus how much the prompt can modify the generation.

The Modes

dampen (Recommended)

Reduces modification strength before applying changes. This is the most reliable mode for precise preservation.

For consistent identity/object preservation: 1.20 to 1.30

linear

Applies full modifications, then blends the result back with the original.

hybrid

Dampens parameters first, then blends the result.

blend_after

Same as linear, just a different name.

Usage

1.20-1.30 (dampen): Recommended starting point for solid preservation
1.40-1.50: Tighter control when needed, very prompt-dependent
0.0-1.0: Standard range from full enhancement to balanced preservation

How It Works

Text Conditioning Enhancement

Magnitude

Direct scaling of all embedding vectors in the active region:

active = active * magnitude

Contrast (Safe Implementation)

Computes the mean embedding across the sequence, then amplifies deviations:

seq_mean = active.mean(dim=1, keepdim=True)
deviation = active - seq_mean
if contrast >= 0:
    scale = 1.0 + contrast
else:
    scale = math.exp(contrast)  # Never inverts, asymptotes to 0
active = seq_mean + deviation * scale

Note: Negative contrast uses exponential scaling to prevent semantic inversion that occurs with linear scaling below -1.

Normalize Strength

Equalizes token magnitudes toward a uniform value:

token_norms = active.norm(dim=-1, keepdim=True)
mean_norm = token_norms.mean()
normalized = active / token_norms * mean_norm
active = active * (1.0 - normalize_strength) + normalized * normalize_strength

Reference Latent Control

Strength is applied to k and v inside every attention block via attn1_patch, after all normalisation has completed:

# Token layout: [ txt_tokens | main_img_tokens | ref_0_tokens | ... | ref_N_tokens ]
k[:, :, seq_start:seq_end, :] *= strength
v[:, :, seq_start:seq_end, :] *= strength

reference_image_num_tokens from extra_options gives the exact token count per reference, allowing precise per-reference indexing without affecting other tokens.

Active Region Detection

Auto-detected from attention mask:

attn_mask = meta.get("attention_mask", None)
nonzero = attn_mask[0].nonzero()
active_end = int(nonzero[-1].item()) + 1

Presets

Text-to-Image

              BASE   GENTLE   MOD   STRONG   AGG     MAX    CRAZY
              ----    ----    ----    ----    ----    ----    ----
magnitude:    1.20    1.15    1.25    1.35    1.50    1.75    2.50
contrast:     0.00    0.10    0.20    0.30    0.40    0.60    1.20
normalize:    0.00    0.00    0.00    0.15    0.25    0.35    0.60
edit_weight:  1.00    1.00    1.00    1.00    1.00    1.00    1.00

Image Edit

              PRESERVE   SUBTLE   BALANCED   FOLLOW   FORCE
              --------   ------   --------   ------   -----
magnitude:       0.85     1.00       1.10     1.20    1.35
contrast:        0.00     0.05       0.10     0.15    0.25
normalize:       0.00     0.00       0.10     0.10    0.15
edit_weight:     0.70     0.85       1.00     1.25    1.50
ref_strength:    1.50     1.20       1.00     0.70    0.30

BETA: Mask-Guided Reference Latent Controller

Experimental — results are promising but behavior may vary depending on prompt and image complexity.

FLUX.2 Klein Mask Ref Controller

Spatially controls the reference latent using a mask. Masked area gets targeted by the prompt while staying true to its original structure. Unmasked area gets fully freed up for the prompt to take over.

Not inpainting — works entirely at the conditioning level through the reference latent stream.

Parameters

Parameter	Default	Range	Description
`mask`	—	MASK	Defines targeted vs free regions. Connect from any ComfyUI mask node.
`strength`	1.0	0.0 to 1.0	How free the unmasked area is. 1.0 = reference fully removed there. 0.5 = half reference kept. Lower values also bleed influence into surrounding areas.
`invert_mask`	False	True/False	Flip targeted and free regions.
`feather`	0	0 to 64	Gaussian blur on mask edges in latent space. Reduces hard seams at boundaries.
`channel_mode`	all	all/low/high	Which latent channels the mask affects. `low` = structure/layout (ch 0-63), `high` = texture/detail (ch 64-127).
`debug`	False	True/False	Print spatial stats and attenuation to console.

Notes

Requires image edit mode — reference latents must be present in conditioning metadata
strength also controls boundary bleed — 1.0 is a tight boundary, lower values spread influence into neighboring regions
Useful for targeting a specific subject within a scene while leaving the rest open to the prompt

Identity Preservation Nodes

Two nodes that approach identity preservation from outside the model's conditioning pipeline. They bypass text streams, attention masks, and token injection entirely.

FLUX.2 Klein Identity Guidance

Operates in the sampling loop. After the model predicts the denoised image at each step, compares it to the reference latent and pulls it back. The model runs freely, then gets corrected.

Takes a VAE-encoded reference image directly. ReferenceLatent should still be connected on the conditioning path so the model has reference context.

Wiring:

[Checkpoint] → MODEL → [Identity Guidance] → MODEL → [KSampler]
                              ↑
                    [VAE Encode of reference]

Parameters

Parameter	Default	Range	Description
`strength`	0.3	0.0 to 1.0	How hard to pull toward the reference each step. 0.3 = move 30% of the distance.
`start_percent`	0.0	0.0 to 1.0	When to start correcting. 0.0 = beginning of denoising.
`end_percent`	0.8	0.0 to 1.0	When to stop correcting. 0.8 = last 20% runs freely for texture refinement.
`mode`	adaptive	adaptive/direct/channel_match	How correction is applied (see below).

Modes

adaptive: Pulls only where the prediction already resembles the reference. Preserves prompt-driven changes like new backgrounds or poses.
direct: Pulls everywhere equally. Strongest identity lock, least prompt freedom.
channel_match: Forces the generation's color and feature statistics to match the reference without copying spatial content.

FLUX.2 Klein Identity Feature Transfer

Operates inside the model's attention layers. After each attention block computes, finds where the generation's features are similar to the reference's features and pushes them closer.

Requires ReferenceLatent connected. The reference must already be in the image stream.

Wiring:

[Checkpoint] → MODEL → [Identity Feature Transfer] → MODEL → [KSampler]
                                                        ↑
                        [ReferenceLatent] → CONDITIONING → [KSampler]

Parameters

Parameter	Default	Range	Description
`strength`	0.15	0.0 to 1.0	Per-block blend factor. Fires at every active block (cumulative). Start at 0.10 to 0.20.
`start_block`	0	0 to 23	First block index to apply. Index is shared across double and single blocks (resets when single blocks begin).
`end_block`	23	0 to 23	Last block index to apply. Covers 8 double blocks (0-7) then 24 single blocks (index resets 0-23).
`mode`	cosine_pull	cosine_pull/topk_replace/mean_transfer	How features are transferred (see below).
`top_k_percent`	0.25	0.01 to 1.0	topk_replace mode only. Fraction of tokens to affect.

Modes

cosine_pull: Each generation token finds its most similar reference token and gets pulled toward it. Strength scales with similarity.
topk_replace: Only the top K% most similar tokens are affected. Everything else stays untouched.
mean_transfer: Shifts the overall feature distribution toward the reference without spatial matching.

FLUX.2 Klein Identity Feature Transfer Advanced

Same approach as the standard transfer node, with finer control: per-band strength on double vs single blocks, a similarity floor that gates how tight or loose the pull is, a block schedule, and an optional subject mask that restricts the pull to the masked area of the reference.

Requires ReferenceLatent connected. The reference must already be in the image stream.

Wiring:

[Checkpoint] → MODEL → [Identity Feature Transfer Advanced] → MODEL → [KSampler]
                                              ↑                          ↑
                                       [MASK (optional)]    [ReferenceLatent] → CONDITIONING

Parameters

Parameter	Default	Range	Description
`reference_index`	0	0 to 15	Which reference image to draw features from when multiple are connected (0 = first).
`mode`	cosine_pull	cosine_pull/topk_replace/mean_transfer	Same modes as the standard node. cosine_pull is the only mode that uses the subject mask.
`top_k_percent`	0.25	0.01 to 1.0	`topk_replace` mode only. Fraction of generation tokens to affect.
`double_enable`	True	True/False	Apply transfer on double blocks (0-7). These shape pose, color, and identity early.
`double_strength`	0.15	0.0 to 1.0	Per-block blend factor for double blocks. Cumulative. Raise (0.4-0.6) for stronger guidance, especially with multi-subject references.
`double_start` / `double_end`	0 / 7	0 to 7	Range of double blocks to apply on.
`single_enable`	True	True/False	Apply transfer on single blocks (0-23). These refine style and texture.
`single_strength`	0.15	0.0 to 1.0	Per-block blend factor for single blocks. Cumulative.
`single_start` / `single_end`	0 / 23	0 to 23	Range of single blocks to apply on.
`block_schedule`	flat	flat / ramp_down / ramp_up / peak_mid	Strength curve across the active block range.
`sim_floor`	0.20	0.0 to 0.95	Cosine threshold gating which matches contribute. Low (~0.05) = wide pull, tight identity lock, suited to subtle edits. High = sparse pull, more freedom for broader edits.
`mask_threshold`	0.5	0.0 to 1.0	Used only when `subject_mask` is connected. Reference tokens whose pooled mask value falls below this are excluded from the pull. 0.5 keeps boundary tokens; raise toward 1.0 to shrink the effective mask inward.
`subject_mask` (optional)	—	MASK	When connected, the cosine pull samples only from masked-in reference tokens. The conditioning latent is not modified. Mask aspect must match the encoded reference aspect.

Tuning notes

Subtle edit (outfit swap, same character): lower sim_floor to ~0.05, keep strengths around 0.15-0.20. The pull is wide and locks identity tightly so only the prompted change lands.
Broader edit (new scene, same person): raise sim_floor toward 0.4-0.6. The pull becomes sparse and gives the model freedom to drift.
Reference contains multiple subjects: paint the desired one as subject_mask and raise both block strengths to 0.4-0.6. The default 0.15 is too soft to compete with the model's natural attention to the unmasked subject.
Subject mask source: ComfyUI's built-in MaskEditor on the reference image works fine. Any MASK output is accepted.

FLUX.2 Klein Identity Feature Transfer V3

Shipment-ready identity preservation node for FLUX.2 Klein. This is the easier version to use.

V3 uses a commit system instead of constantly averaging the generation toward the whole reference. Each generation token looks for the best matching reference token. If the same match stays stable, it locks to that match and then keeps a lighter anchor instead of continuing to pull hard every block.

This gives cleaner preservation and less feature mush.

Requires ReferenceLatent connected. The reference must already be in the image stream.

Wiring:

[Checkpoint] -> MODEL -> [Identity Feature Transfer V3] -> MODEL -> [KSampler]
                                             ↑                         ↑
                                      [MASK optional]    [ReferenceLatent] -> CONDITIONING

Presets

Preset	Use
`MIDUM_LOCK`	Recommended starting point. Uses the tested middle preset.
`HARD_LOCK`	Strongest preset. Use when the reference keeps drifting.
`SOFT_LOCK`	Softer preset. Use when the lock is too aggressive.
`custom`	Uses your manual schedules and settings.

Any preset except custom overrides the manual settings below. This is intentional so the node is easier to use.

Parameters

Parameter	Default	Description
`preset`	MIDUM_LOCK	Main control. Pick how strongly the reference should hold.
`reference_index`	0	Which reference image to use when multiple references are connected.
`double_schedule`	0-3:mid=0.25; 4:mid=0.35; 5:mid=0.65; 6-7:mid=0.45	Custom preset only. Controls double blocks.
`single_schedule`	0:mid=0.35; 1:mid=0.25; 2-10:mid=0.30; 11-19:mid=0.25; 20:mid=0.08; 21:mid=0.10; 22:mid=0.15; 23:mid=0.20	Custom preset only. Controls single blocks.
`double_sim`	0.020	Custom preset only. Higher means fewer double-block matches are allowed.
`single_sim`	0.020	Custom preset only. Higher means fewer single-block matches are allowed.
`commit_margin`	0.035	Custom preset only. Higher means the match has to be more obvious before it locks.
`commit_confirm`	2	Custom preset only. How many times the same match must repeat before it locks.
`commit_anchor`	0.35	Custom preset only. How much pull remains after a token has locked.
`mask_threshold`	0.25	Used only when `subject_mask` is connected. Lower keeps more edge tokens.
`subject_mask`	optional	Use this when the reference has more than one subject.

Schedule format

4-7:0.20
2-10:0.10; 14-20:0.07

Left side is the block range. Right side is the strength.

Double blocks are 0-7. Single blocks are 0-23.

Starting points

Use MIDUM_LOCK first.

If identity is weak, switch to HARD_LOCK.

If the reference is copying too much or distorting the result, switch to SOFT_LOCK.

Use custom only after you know what you want to change.

Combining Both Nodes

Both nodes can be stacked. Identity Guidance handles macro-level correction in latent space. Feature Transfer handles micro-level feature alignment inside attention. They operate at different stages and don't interfere.

[Checkpoint] → MODEL → [Identity Feature Transfer] → MODEL → [Identity Guidance] → MODEL → [KSampler]

Technical Details

Model: FLUX.2 Klein 9B
Text Encoder: Qwen3 8B (4096 hidden dim, 36 layers)
Conditioning Shape: [batch, 512, 12288]
Reference Latent Shape: [batch, 128, H, W] (stored in metadata)
Joint Attention Dim: 12288
Architecture: 8 double_blocks (separate streams) + 24 single_blocks (concatenated)
img_in projection: [128] → [4096]
txt_in projection: [12288] → [4096]
Guidance Embeds: False (step-distilled model, no CFG)

Methodology

All findings were verified through:

Conditioning Diagnostic: Tensor structure analysis
Forward Hook Tracing: Attached hooks to img_in, txt_in, double_blocks, single_blocks
Comparative Analysis: With vs without reference latent runs
Model Introspection: forward() signature analysis

The reference latent concatenation was confirmed by observing:

With reference: img_in receives [1, 8140, 128] (4070 × 2 patches)
Without reference: img_in receives [1, 4070, 128]

Acknowledgments

Built through empirical analysis and forward-pass hook tracing of FLUX.2 Klein's conditioning structure and dual-stream architecture.

License

This project is licensed under the PolyForm Noncommercial License 1.0.0.

✅ Free for noncommercial use — personal projects, research, hobby work, education, charitable/public-interest organizations, and other noncommercial purposes.
✅ You may modify, distribute, and build on it for any noncommercial purpose.
❌ Commercial use requires a separate license. If you (or your company) want to use this in a commercial product, paid service, or any revenue-generating context, please contact me first to arrange commercial licensing.

For commercial licensing inquiries, open an issue on the GitHub repository or reach out via the contact info on my GitHub profile.

README.md Unescape Escape

ComfyUI-Flux2Klein-Enhancer

What This Does (( NEW DISCOVERY!! ))

Verified Architecture (from forward-pass hooks)

Installation

Nodes

FLUX.2 Klein Enhancer

Parameters

FLUX.2 Klein Detail Controller

Parameters

FLUX.2 Klein Ref Latent Controller

Parameters

FLUX.2 Klein Text/Ref Balance

Parameters

Preserve Original - Solving FLUX Klein's Preservation Problem

The Modes

dampen (Recommended)

linear

hybrid

blend_after

Usage

How It Works

Text Conditioning Enhancement

Magnitude

Contrast (Safe Implementation)

Normalize Strength

Reference Latent Control

Active Region Detection

Presets

Text-to-Image

Image Edit

BETA: Mask-Guided Reference Latent Controller

FLUX.2 Klein Mask Ref Controller

Parameters

Notes

Identity Preservation Nodes

FLUX.2 Klein Identity Guidance

Parameters

Modes

FLUX.2 Klein Identity Feature Transfer

Parameters

Modes

FLUX.2 Klein Identity Feature Transfer Advanced

Parameters

Tuning notes

FLUX.2 Klein Identity Feature Transfer V3

Presets

Parameters

Schedule format

Starting points

Combining Both Nodes

Technical Details

Methodology

Acknowledgments

License

README.md