vae tiling improvements: encoding support and adaptative overlap #484

stduhpf · 2024-11-26T15:12:54Z

Added a mechanism to automatically fix the tile overlaps in cases of mismatch between image size and tile size.
Added a way to change tile size for VAE decoding with the SD_TILE_SIZE environment variable. (tile size for vae encoding will be about 1.3x bigger, for the same memory requirement)

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 #588, #564 and #666

Green-Sky · 2024-11-28T10:38:15Z

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 · 2025-02-13T04:34:52Z

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 · 2025-02-13T17:02:54Z

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 #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 · 2025-02-15T02:35:02Z

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 · 2025-02-15T13:32:17Z

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 · 2025-04-25T20:38:30Z

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 #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 · 2025-08-02T19:56:22Z

@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 · 2025-09-06T12:05:12Z

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

stduhpf · 2025-09-06T17:51:52Z

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 · 2025-09-06T18:13:47Z

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.

wbruna

I've only noticed a single, very familiar issue 😀

stable-diffusion.cpp

stduhpf · 2025-09-08T11:10:40Z

Encode is also broken, fixing it

wbruna

There is some unrelated whitespace changes, but apart from that it looks good to me. Working fine with image to image in SDXL, SD3 and Chroma, and image edit with Kontext.

leejet · 2025-09-10T14:00:53Z

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 · 2025-09-10T15:19:43Z

@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 · 2025-09-10T18:09:02Z

@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 · 2025-09-10T20:46:48Z

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 · 2025-09-10T21:39:33Z

@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 😃

* refactor tile number calculation * support non-square tiles * add env var to change tile overlap * add safeguards and better error messages for SD_TILE_OVERLAP * add safeguards and include overlapping factor for SD_TILE_SIZE * avoid rounding issues when specifying SD_TILE_SIZE as a factor * lower SD_TILE_OVERLAP limit * zero-init empty output buffer

wbruna

Looking good, and working fine on a few tests. The only issue I notice was a crash on an invalid tile size (more detailed comments below).

wbruna · 2025-09-12T12:41:51Z

stable-diffusion.h

+    int tile_size_x;
+    int tile_size_y;
+    float target_overlap;
+    bool relative;


Perhaps this boolean could be dropped, and we apply the same logic as the command-line: if rel_size_ is bigger than 0.0, it overrides the tile_size_ ?

wbruna · 2025-09-12T12:56:03Z

stable-diffusion.cpp

        int64_t t0 = ggml_time_ms();
        if (!use_tiny_autoencoder) {
+            float tile_overlap = vae_tiling_params.target_overlap;
+            int tile_size_x    = vae_tiling_params.tile_size_x;
+            int tile_size_y    = vae_tiling_params.tile_size_y;


A relative size of 0.0 applies a 4x4 tile (doesn't look good, but works); but a zeroed-out tile size crashes. Maybe an invalid value could just mean the default?

int tile_size_x = vae_tiling_params.tile_size_x < 4 ? 32 : vae_tiling_params.tile_size_x; int tile_size_y = vae_tiling_params.tile_size_y < 4 ? 32 : vae_tiling_params.tile_size_y;

Then we could also simplify the default values for the param struct, by using zeroed-out fields.

Same issue with the tile_overlap above.

stduhpf · 2025-09-12T13:47:30Z

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

* avoid crash with invalid tile sizes, use 0 for default * refactor default tile size, limit overlap factor * remove explicit parameter for relative tile size * limit encoding tile to latent size

stduhpf force-pushed the tiled-vae-encode branch from 9edc59f to af0491b Compare November 26, 2024 15:36

stduhpf changed the title ~~Implement vae tiling for encoding~~ vae tiling improvements: encoding support and adaptative overlap Nov 27, 2024

stduhpf force-pushed the tiled-vae-encode branch from 6f49af0 to 6f96903 Compare November 27, 2024 02:27

stduhpf force-pushed the tiled-vae-encode branch from 5495c93 to e9a8d17 Compare December 8, 2024 01:06

stduhpf force-pushed the tiled-vae-encode branch from e9a8d17 to d103632 Compare December 28, 2024 21:51

stduhpf mentioned this pull request Jan 10, 2025

VAE tiling seems to skip the last several tiles during processing #564

Open

stduhpf closed this Feb 15, 2025

stduhpf reopened this Feb 15, 2025

stduhpf mentioned this pull request Apr 25, 2025

Upscaling broken for non-512x512 images #666

Open

stduhpf force-pushed the tiled-vae-encode branch from d103632 to a3fabad Compare June 7, 2025 00:06

This was referenced Jun 13, 2025

Configurable VAE threshold limit LostRuins/koboldcpp#1601

Merged

Chroma support (pruned Flux model) #696

Merged

stduhpf force-pushed the tiled-vae-encode branch 2 times, most recently from d1fcf84 to fb5f5b7 Compare July 10, 2025 17:28

wbruna mentioned this pull request Aug 7, 2025

Can't use any sdxl lora with the OpenCL backend: GGML_ASSERT(src0->type == src1->type) failed #757

Open

stduhpf force-pushed the tiled-vae-encode branch from fb5f5b7 to c0c0097 Compare September 6, 2025 17:50

stduhpf force-pushed the tiled-vae-encode branch from c0c0097 to e78d630 Compare September 6, 2025 17:55

wbruna suggested changes Sep 7, 2025

View reviewed changes

stable-diffusion.cpp Outdated Show resolved Hide resolved

stduhpf force-pushed the tiled-vae-encode branch from e78d630 to d7b671c Compare September 8, 2025 13:30

stduhpf requested a review from wbruna September 8, 2025 13:43

wbruna approved these changes Sep 8, 2025

View reviewed changes

wbruna mentioned this pull request Sep 9, 2025

update stable-diffusion.cpp to master-c648001 (+fixes) LostRuins/koboldcpp#1732

Draft

stduhpf force-pushed the tiled-vae-encode branch 2 times, most recently from 101dcbb to 0a2b4e4 Compare September 11, 2025 10:18

stduhpf and others added 15 commits September 11, 2025 22:29

implement tiling vae encode support

751f501

Tiling (vae/upscale): adaptative overlap

b409c9c

Tiling: fix edge case

b340e42

Tiling: fix crash when less than 2 tiles per dim

22dac4f

remove extra dot

0dfd15a

Tiling: fix edge cases for adaptative overlap

687ee72

tiling: fix edge case

957f009

set vae tile size via env var

5054cf4

vae tiling: refactor again, base on smaller buffer for alignment

2af9294

Use bigger tiles for encode (to match compute buffer size)

f3349d0

Fix edge case when tile is bigger than latent

532aacb

Fix decode latent size

f08d696

fix encode

606976f

tile size params instead of env

987ced8

stduhpf force-pushed the tiled-vae-encode branch from 0a2b4e4 to 987ced8 Compare September 11, 2025 20:29

wbruna suggested changes Sep 12, 2025

View reviewed changes

Tiled vae parameter validation (#6)

2995c92

* avoid crash with invalid tile sizes, use 0 for default * refactor default tile size, limit overlap factor * remove explicit parameter for relative tile size * limit encoding tile to latent size

vae tiling improvements: encoding support and adaptative overlap #484

Are you sure you want to change the base?

vae tiling improvements: encoding support and adaptative overlap #484

Conversation

stduhpf commented Nov 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Green-Sky commented Nov 28, 2024

Uh oh!

masamaru-san commented Feb 13, 2025

Uh oh!

stduhpf commented Feb 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

masamaru-san commented Feb 15, 2025

Uh oh!

stduhpf commented Feb 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

divVerent commented Apr 25, 2025

Uh oh!

wbruna commented Aug 2, 2025

Uh oh!

wbruna commented Sep 6, 2025

Uh oh!

stduhpf commented Sep 6, 2025

Uh oh!

wbruna commented Sep 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wbruna left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

stduhpf commented Sep 8, 2025

Uh oh!

wbruna left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leejet commented Sep 10, 2025

Uh oh!

wbruna commented Sep 10, 2025

Uh oh!

leejet commented Sep 10, 2025

Uh oh!

stduhpf commented Sep 10, 2025

Uh oh!

wbruna commented Sep 10, 2025

Uh oh!

wbruna left a comment

Choose a reason for hiding this comment

Uh oh!

wbruna Sep 12, 2025

Choose a reason for hiding this comment

Uh oh!

wbruna Sep 12, 2025

Choose a reason for hiding this comment

Uh oh!

wbruna Sep 12, 2025

Choose a reason for hiding this comment

Uh oh!

stduhpf commented Sep 12, 2025

Uh oh!

Uh oh!

stduhpf commented Nov 26, 2024 •

edited

Loading

stduhpf commented Feb 13, 2025 •

edited

Loading

stduhpf commented Feb 15, 2025 •

edited

Loading

wbruna commented Sep 6, 2025 •

edited

Loading

wbruna left a comment •

edited

Loading