[Bug] Flux.2-Klein VAE Decode Artefacts

[cholleme]

Git commit

git: 636d3cb

Operating System & Version

Ubuntu 24.04.3 LTS

GGML backends

CUDA

Command-line arguments used

[ "--diffusion-model", "/home/charles/coding/models/flux2/flux-2-klein-4b.safetensors", "--vae", "/home/charles/coding/models/flux2/flux2-vae.safetensors", "--llm", "/home/charles/coding/models/flux2/Qwen3-4B-Q6_K.gguf", "--cfg-scale", "1.0", "--steps", "4", "-v", "--diffusion-fa", "--width", "1024", "--height", "1024", "-p", ""A cat holding a beachball on the river bank."" ]

Steps to reproduce

Get the latents exported by diffusers from latents_3.zip.

Add the following code to load the repro data:

    LOG_INFO("decoding %zu latents", final_latents.size());

    FILE* f=fopen("/home/charles/coding/wishingwell/scripts/latents_3.raw", "rb");
    fseek(f, 0, SEEK_END);
    long fsize = ftell(f);
    fseek(f, 0, SEEK_SET);

    float *data = (float*)malloc(fsize + 1);
    fread(data, fsize, 1, f);
    fclose(f);

    GGML_ASSERT(fsize == ggml_nbytes(final_latents[0]));
    LOG_INFO("latent: %zu %zu %zu %zu", final_latents[0]->ne[0], final_latents[0]->ne[1], final_latents[0]->ne[2], final_latents[0]->ne[3]); 
    auto hack_latent = ggml_dup_tensor(work_ctx, final_latents[0]); 
    memcpy(hack_latent->data, data, fsize);
    final_latents[0] = hack_latent;

After the line

    LOG_INFO("generating %" PRId64 " latent images completed, taking %.2fs", final_latents.size(), (t3 - t1) * 1.0f / 1000);

In stable-diffusion.cpp

(also attached the python code to generate this file)
simple_inference.py
latents_3.zip

What you expected to happen

When working with flux2.klein I noticed artifacts I didn't see when using the python diffuses version before. I suspected the VAE so made a quick repro where diffusers outputs the raw VAE and then load it in stable-diffusion.cpp for decoding. The results are worse in particular notice the gradients on the ball.

This seems beyond expected difference of the two implementations and leads to unusable quality.

What actually happened

Noticeable quality degradation see Images. Same style artifacts happen when directly generating images from prompts. They seem mostly visible in smooth gradient areas.

This one is generated using stock stable-diffusion.cpp

Logs / error messages / stack trace

None

Additional context / environment details

NVIDIA GB10 Driver Version: 580.95.05 CUDA Version: 13.0

[stduhpf]

I've also noticed the desaturated colors when asking the model to "Recreate this exact same image" with a reference image. I assumed it was an issue with the model, but now it's clear something is not right.

This effect seems to be more prononced at higher resolutions, can you confirm?

[stduhpf]

The issue doesn't happen with TAEF2 (https://huggingface.co/madebyollin/taef2/tree/main):

[jcluts]

I noticed this exact same issue and spent days trying to fix it myself. I was at least able to isolate it to being specific to CUDA and the VAE decode, like cholleme was able to do. As best I could tell it was in ggml internals in code I lack the expertise to troubleshoot.

My current workarounds are to use Vulkan or --vae-on-cpu, which bypasses CUDA for the VAE decode. In both cases I get image quality that is excellent in comparison to the CUDA decode.

[stduhpf]

It also happens on ROCm backend (which is based on CUDA, so it makes sense)

[stduhpf]

This is with Vulkan backend (with --vae-conv-direct to avoid OOM):

With --vae-on-cpu:

I don't think it's specific to CUDA backend

[jcluts]

You're right, I can reproduce on Vulkan. I've been defaulting to always use --fa on every generation since I struggled to not get OOM without it. I went to Q4 for everything so I could get generations through without --fa, and sure enough it has the same desaturation/artifacting as CUDA without --fa. Same thing with --vae-on-cpu, I still get the same artifacting with it both on and off, with Vulkan.

So to summarize what I see for Vulkan when I isolate flags:
--vae-on-cpu: bad
--diffusion-fa: bad
--fa: good

For some reason CUDA is now silently failing for me and I haven't had a chance to troubleshoot why, so I can't run similar tests on the affect of of --fa and --diffusion-fa there. But it did seem that --vae-on-cpu was fixing CUDA for me when I was troubleshooting this last week.

[JohnLoveJoy]

Yeah. I noticed in Z Image Turbo that the image appears slightly sharper when using the VAE on the CPU.

However, when Conv2D Direct is enabled, the image seems a bit more blurred. Should this affect image quality?

[stduhpf]

@JohnLoveJoy Z Image Turbo uses Flux1 VAE, it shouldn't have this same issue...

[masamaru-san]

For reference: I’m seeing severe VAE quality degradation on my setup, so I locally reverted commit 885e62e.

• backend: Vulkan
  • AMD Ryzen 7 5700U with Radeon Graphics + 32GB UMA memory
  • Windows 11 25H2 with AMD Software: Adrenalin Edition 25.8.1 (vulkan driver v2.0.279)
• Models:
  • sd3.5m_turbo.safetensors
  • stable-diffusion-v1-5-pruned-emaonly-Q8_0.gguf
  • Any SD1.x based models like anything-v4.5-pruned.safetensors or so.

[leejet]

#1297 should fix this issue.

[JohnLoveJoy]

For reference: I’m seeing severe VAE quality degradation on my setup, so I locally reverted commit 885e62e.

• backend: Vulkan

  • AMD Ryzen 7 5700U with Radeon Graphics + 32GB UMA memory
  • Windows 11 25H2 with AMD Software: Adrenalin Edition 25.8.1 (vulkan driver v2.0.279)

• Models:

  • sd3.5m_turbo.safetensors
  • stable-diffusion-v1-5-pruned-emaonly-Q8_0.gguf
  • Any SD1.x based models like anything-v4.5-pruned.safetensors or so.

Simply disabling FA should fix that, don't you think? The current SD.cpp' FA implementation does not work well for GPUs prior to the RDNA gen.

[masamaru-san]

Thanks to #1297, at least sd1.x and sd3.5m seem to have been fixed. 😄

@JohnLoveJoy I forgot to mention that I'm not using Flash Attention.