fix(flux2): Fix FLUX.2 Klein image generation quality

[Pfannkuchensack]

Summary

This commit addresses severe quality degradation and artifacts when generating images with FLUX.2 Klein models.

Root causes fixed:

1. BN Normalization on noise input (flux2_denoise.py)

   • Random noise was incorrectly normalized before denoising
   • Diffusers only normalizes image latents from VAE.encode(), not random noise
   • Random noise is already N(0,1) distributed - normalization corrupted it

2. Missing BN Denormalization on output (flux2_denoise.py)

   • Output latents must be denormalized before VAE decode
   • Converts from transformer space to VAE-expected space
   • Without this, images appear washed-out with lost colors

3. Dynamic mu calculation causing over-shifted schedules (sampling_utils.py)

   • mu was resolution-dependent, causing extreme values at high resolutions
   • At 2048x2048: mu=3.23 compressed almost all sigmas above 0.9
   • Fixed to constant 2.02 (matching ComfyUI's proven configuration)

Impact:

• All resolutions: Proper color reproduction, no washed-out images
• High resolutions (>1024x1024): No grid/diamond artifacts
• FLUX.2 Klein 4B/9B: Both models now produce correct output
• All generation modes: T2I and reference images

Related Issues / Discussions

#8832

Community-reported quality degradation issues on Discord:

• Garbage bin examples and discussion
• Additional garbage bin examples
• Devchat discussion with examples
• ComfyUI comparison showing issue is InvokeAI-specific

QA Instructions

Test Setup

Use FLUX.2 Klein models (both 4B and 9B variants) with the following configurations:

Critical Test Cases (Priority)

1. Baseline - Standard Resolution:

• Generate at 1024×1024 with both Klein 4B and 9B
• Use 4 steps for distilled models, 28 steps for base model
• Expected: Full color, proper detail, no artifacts

2. High Resolution Test:

• Test resolutions: 1536×1536, 2048×2048
• Klein 9B was heavily affected here before fix
• Expected: Consistent quality, no grid/diamond patterns

3. Color Reproduction:

• Generate colorful subjects (e.g., "orange tiger in green forest")
• Verify colors are vibrant, not washed-out or desaturated
• Expected: Full color saturation throughout denoising

4. Both Generation Modes:

• Text-to-Image (T2I) generation
• Reference image editing (Kontext)

What to Look For

Improvements:

• ✅ No grid/diamond artifacts at any resolution
• ✅ Full color reproduction (no washed-out images)
• ✅ Proper detail preservation and coherence
• ✅ Consistent quality scaling with resolution
• ✅ Stable denoising process (colors maintained throughout)

Technical Verification:

• BN stats should be found from VAE (check logs if issues)
• Latent values should be properly scaled after denormalization

Merge Plan

No special merge considerations needed.

This is a targeted bug fix with:

• No breaking changes
• No API changes
• No database schema changes
• No migration required

Safe to merge after QA verification.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable) - Not applicable (algorithmic fix)
• ❗Changes to a redux slice have a corresponding migration - Not applicable (backend-only changes)
• Documentation added / updated (if applicable) - Not applicable (internal algorithm fix, code comments added)
• Updated What's New copy (if doing a release after this PR) - Recommended: "Fixed FLUX.2 Klein image quality issues (color loss, artifacts)"

[lstein]

I can confirm that the PR does not degrade image quality or performance, but I'm having a difficult time reproducing the image quality degradation on the main branch for comparison. Is it possible that the issue only appears with certain video cards?

I am going to ask other developers on Discord to do some testing.

[lstein]

Just updating the comment thread to match conclusions on Discord.

This PR does not fix the underlying issue. At higher step counts, larger image sizes, and/or shorter prompts, both flux.2 klein models produce marked image artifacts. The effect is easily seen at 1024x1024 by generating without a prompt, which produces images with prominent diagonal elements.

[JPPhoto]

Results are worlds better than they were, great job!