VAE Inspector (Flux 2 Klein — 125 Tensors)

What it is

Klein has a documented, reproducible quirk: edits come back with a warm color shift, and it shows up regardless of quantization - reported on fp8, confirmed on fp16, so it's not a precision artifact you can quantize your way out of. The community's fix so far has been external - histogram or color matching against the source after the fact, which a tile-upscaler node in the ecosystem now does per tile. This node is a more direct way to investigate the source: it analyzes all 125 individual tensor units in the FLUX.2 VAE (a batch norm, 70 decoder units, 54 encoder units - every conv, norm, shortcut, and q/k/v/proj individually) and can measure exactly which ones are pushing color, rather than papering over the symptom downstream.

This is a specialist diagnostic tool, not something you run on every generation. Reach for it when a color cast, a saturation bias, or some other pixel-level artifact keeps showing up across many outputs and you want to find the actual source in the VAE rather than keep correcting for it after decode.

How it works

Two modes, with a real cost difference between them:

Weight analysis (default, fast) looks at weight norms, distributions, and max values across all 125 tensor units - it tells you which convs and norms dominate the VAE's own weights, without needing to actually decode anything.

Decode ablation (slow, needs a latent) is the one that actually answers "which unit is responsible for this color problem": it weakens each decoder unit individually, decodes, and measures the pixel-level and per-channel R/G/B impact of removing it. That gives you a precise importance ranking for color tuning specifically - exactly the kind of evidence you'd want before reaching for the matching VAE Deep Debiaser node to correct a specific unit.

Inputs and outputs that matter

• vae - the FLUX.2 VAE to analyze.
• analysis_mode - weight analysis (fast) or decode ablation (slow, requires a latent input).
• ablation_strength (0–1, default 0.5) - how hard each unit gets weakened during the ablation test.
• Optional latent - required only for decode ablation mode; without it, you're limited to the weight-analysis pass.

Outputs: report (a human-readable breakdown, including per-channel R/G/B impact in ablation mode) and json_data (the same analysis in a structured form).

Installing it

No external backend needed - this node works the moment the pack is installed. ComfyUI Manager (search "Realtime LoRA Trainer") or:

cd ComfyUI/custom_nodes
git clone https://github.com/ShootTheSound/comfyUI-Realtime-Lora

Restart ComfyUI.

Common issues

Decode ablation is very slow. Expected - it's decoding 125 times, once per unit, to isolate each one's individual contribution. Run weight analysis first to narrow down which region of the VAE (encoder vs decoder, and roughly where) looks unusual, and only run ablation if you need the precise per-unit ranking.

Ran decode ablation with no latent connected. This mode needs an actual latent to decode repeatedly while ablating - wire one in from your sampler, or stay in weight-analysis mode if you don't have one handy.

Results don't point at an obvious fix. This node identifies which tensor unit correlates with a color or pixel-level effect - it doesn't change anything itself. Pair the findings with VAE Deep Debiaser (Flux 2 Klein - 125 Tensors) to actually scale the implicated unit down (or up) and test whether that resolves the artifact.

Expected this to work on a non-Flux VAE. It's built and mapped specifically against the FLUX.2 Klein VAE's 125-tensor structure - a different architecture's VAE won't have the same unit layout, so don't expect this node's block map to transfer.