vae tiling improvements: encoding support and adaptative overlap

[stduhpf]

• Added a boolean argument sd_tiling() function, specifying if the output tile is scaled up (for decoding) or if the input is (encoding).
• Added a mechanism to automatically fix the tile overlaps in cases of mismatch between image size and tile size.

With these changes, vae encoding can now be tiled, and a small (unreported?) issue with the upscaler (or tae) over-brightening some parts of the image is now fixed. This could also be a step toward tiled diffusion support.

Fixes https://github.com/leejet/stable-diffusion.cpp/issues/588, https://github.com/leejet/stable-diffusion.cpp/issues/564 and https://github.com/leejet/stable-diffusion.cpp/issues/666

[Green-Sky]

• Added a mechanism to automatically fix the tile overlaps in cases of mismatch between image size and tile size.

With these changes, vae encoding can now be tiled, and a small (unreported?) issue with the upscaler over-brightening some parts of the image is now fixed. This could also be a step toward tiled diffusion support.

Nice! It would also be interesting to see larger tile sizes. Depending on VAE, larger tiles would still fit into the memory budget and would speed things up.

[masamaru-san]

I am not familiar with AI programming, and this is a machine translation, so please ignore me if I say something strange.

The last_x used in the loop at line 568 (or line 629 in this PR) in ggml_extend.hpp is only initialized outside the loop. Locally, it seems that this PR with initializing last_x = false; before the for (int x = 0; ~ loop would improve the generation results.

However, the brightness issue (issue #588) seems to remain.

[stduhpf]

The last_x used in the loop at line 568 (or line 629 in this PR) in ggml_extend.hpp is only initialized outside the loop. Locally, it seems that this PR with initializing last_x = false; before the for (int x = 0; ~ loop would improve the generation results.

I'm sorry, I don't quite understand what you mean. last_x is already set to false when the loop starts.

However, the brightness issue (issue https://github.com/leejet/stable-diffusion.cpp/issues/588) seems to remain.

You mean this PR doesn't fix it for you? If so this is strange, I can't reproduce the issue with this PR...

master	PR

As you can see the image made with the master branch is much brighter than the one with this PR. (ignore the slight variation of the composition, it's just a little unrelated bug of the vulkan backend that happens at non standard resolutions, that's fixed by updating GGML)

[masamaru-san]

The last_x used in the loop at line 568 (or line 629 in this PR) in ggml_extend.hpp is only initialized outside the loop. Locally, it seems that this PR with initializing last_x = false; before the for (int x = 0; ~ loop would improve the generation results.

I'm sorry, I don't quite understand what you mean. last_x is already set to false when the loop starts.

Umm.. Sometimes the bottom of the generated image is grayed out; this did not happen when I initialized last_x in the y-loop. It could be caused by something else.

You mean this PR doesn't fix it for you? If so this is strange, I can't reproduce the issue with this PR...

I think I have misled you, but there is no doubt that this PR is very good, as you have demonstrated in your examples. What I meant to say is that in your right cat PR image, for example, the position around the neck is slightly whiter than without vae-tiling; if it's a technical problem with the GGML backend or tiling (I usually use the vulkan backend), I'll find a better way.

[stduhpf]

Umm.. Sometimes the bottom of the generated image is grayed out; this did not happen when I initialized last_x in the y-loop. It could be caused by something else.

Hmm that's really weird. I don't see how this could happen. last_x is initialized to false before the loop (https://github.com/stduhpf/stable-diffusion.cpp/blob/d103632919ace8a6cc56326a684b56324fa6aaed/ggml_extend.hpp#L627), and it's always reset to false after the x-loop ends (https://github.com/stduhpf/stable-diffusion.cpp/blob/d103632919ace8a6cc56326a684b56324fa6aaed/ggml_extend.hpp#L658). This should be strictly equivalent to initializing it to false at the beginning of the y-loop...

I think I have misled you, but there is no doubt that this PR is very good, as you have demonstrated in your examples. What I meant to say is that in your right cat PR image, for example, the position around the neck is slightly whiter than without vae-tiling; if it's a technical problem with the GGML backend or tiling (I usually use the vulkan backend), I'll find a better way.

The result with tiled VAE is always going to be slightly lower quality than the full VAE, because it's lacking information about the whole image. (That's especially true with SD1.x/SD2.x models, because the VAE is quirky ). When using TAE (or with SDXL/SD3 or FLUX), the results are almost completely identical with the PR, while on master, there is the "overexposure" issue on the tiled TAE.

PR: vae	PR: tiled vae	absolute difference	absolute difference *2

(there is a lot of difference here, mostly because of the kl-f8 vae doesn't handle tiling very well)

PR: tae	PR: tiled tae	absolute difference	absolute difference *127

(these are almost identical however)

absolute difference between tae and vae	x2

(no tiling)

[divVerent]

I can't vouch for the new code being 100% correct, not knowing the code base, but the approach looks reasonable (but does it have to match how the upscaler model was trained, and if so, does it match? I do not know enough for that). But I can definitely vouch for the old code being incorrect. It assumes that tile overlap is equal in x and y direction, and that's unlikely assuming "somewhat arbitrary" total sizes and a fixed tile size.

And I can definitely confirm that https://github.com/leejet/stable-diffusion.cpp/issues/666 is fixed by this change - all five test cases pass (up from only the base case), and the more complex case I discovered this with appears to be fixed as well. Now the difference pictures from pre- and post-upscale pictures (after scaling both back to the same size) look like an edge detector without any noticeable artifacts/biases - precisely what's expected.

[wbruna]

@leejet , just to let you know: this PR has been included in Koboldcpp since a few releases ago (1.94 - 1.96.2) to fix #588, with no issues reported so far, so it should be safe to merge.

[wbruna]

@stduhpf , would you please rebase this onto master? There are a few conflicts due to the VAE changes for Wan support.

[stduhpf]

I haven't tested it, I know it builds fine, but I'm not sure if it actually works (I can't test it right now)

[wbruna]

I haven't tested it, I know it builds fine, but I'm not sure if it actually works (I can't test it right now)

No problem, I'll test it here.

[stduhpf]

Encode is also broken, fixing it

[leejet]

Thank you for your contribution. Please change the environment variables to command-line arguments. Once this change is completed, I will merge the PR.

[wbruna]

@leejet , a quick question: the VAE tiling doesn't really change the loaded models, and its usefulness depends on the chosen image resolution (we even have a workaround on Koboldcpp to enable or disable it at inference time). Could we move it, together with the new tiling size parameters, to the sd_img_gen_params_t struct (possibly a new sd_tiling_params_t struct)? (I can do this change, either separately or on top of this PR)

[leejet]

@leejet , a quick question: the VAE tiling doesn't really change the loaded models, and its usefulness depends on the chosen image resolution (we even have a workaround on Koboldcpp to enable or disable it at inference time). Could we move it, together with the new tiling size parameters, to the sd_img_gen_params_t struct (possibly a new sd_tiling_params_t struct)? (I can do this change, either separately or on top of this PR)

This is what I expect.

[stduhpf]

I can do this change, either separately or on top of this PR

@wbruna do you want to do these changes yourself or should I go ahead and start working on it?

[wbruna]

@wbruna do you want to do these changes yourself or should I go ahead and start working on it?

I won't be able to until Friday, so feel free to do it. From what I looked, the hardest part will be choosing how to implement the tile size parameters: that env var syntax kind of does too much 😃

[stduhpf]

Thanks for the review! I probably won't be able to work on these changes before tomorrow though.

[leejet]

I think this PR can be merged now. Thank you all for your contributions.