feat(stablediffusion-ggml): make VAE tiling configurable (#11216)

GenerateImage hardcoded TilingParamsSetEnabled(vaep, false), so tiled VAE
decoding was unreachable from a model config even though all four upstream
setters were already bound in main.go.

Sampling runs in latent space, but the final VAE decode expands to full
resolution and needs one large compute buffer. At 1024x1024 that buffer
exceeds 8GB, which fails on two kinds of device: cards without the VRAM
for a full-frame decode, and drivers that cap a single allocation
regardless of how much memory is free. Mesa RADV reports a 4GiB
maxMemoryAllocationSize, so a Radeon 8060S with 74GiB of device-local
heap still cannot serve that decode:

    [INFO ] sampling completed, taking 251.82s
    [INFO ] decoding 1 latents
    ggml_vulkan: Requested buffer size exceeds device buffer size limit:
                 ErrorOutOfDeviceMemory
    [ERROR] vae: failed to allocate the compute buffer
    [ERROR] decode_first_stage failed for latent 1

Every sampling step completes and then the run is discarded at the last
stage, so the whole generation is wasted.

Add three options, parsed in Load and applied per generation:

    vae_tiling:true            enable tiled decoding (bare flag also works)
    vae_tile_size:512          tile size, or 512x384 for a rectangle
    vae_tile_overlap:0.25      overlap between tiles

Tiling stays off unless requested, so existing models are unaffected. Tile
size and overlap only reach the library when the operator set them, which
keeps upstream's defaults rather than pushing a zero, and an unparseable
value is treated as absent for the same reason.

Truthy spellings match what load_model already accepts for its own bool
options, and the bare-flag form matches diffusion_model, so no new
convention is introduced.

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>

backend/go/stablediffusion-ggml/gosd.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
 	"unsafe"
 
@@ -18,6 +19,7 @@ type SDGGML struct {
 	threads      int
 	sampleMethod string
 	cfgScale     float32
+	vaeTiling    vaeTiling
 }
 
 var (
@@ -43,6 +45,62 @@ var (
 	VidGenParamsSetVideoFrames func(params uintptr, n int)
 )
 
+type vaeTiling struct {
+	enabled       bool
+	tileSizeX     int
+	tileSizeY     int
+	hasTileSize   bool
+	targetOverlap float32
+	hasOverlap    bool
+}
+
+func parseVAETiling(options []string) vaeTiling {
+	var t vaeTiling
+	for _, op := range options {
+		name, value, hasValue := strings.Cut(op, ":")
+		switch name {
+		case "vae_tiling":
+			// A bare flag reads as "on", matching "diffusion_model". The truthy
+			// spellings are the ones load_model already accepts for its own
+			// bool options, so an author does not have to remember two
+			// conventions.
+			t.enabled = !hasValue || value == "true" || value == "1"
+		case "vae_tile_size":
+			if x, y, ok := parseTileSize(value); ok {
+				t.tileSizeX, t.tileSizeY, t.hasTileSize = x, y, true
+			}
+		case "vae_tile_overlap":
+			if f, err := strconv.ParseFloat(value, 32); err == nil && f >= 0 {
+				t.targetOverlap, t.hasOverlap = float32(f), true
+			}
+		}
+	}
+	return t
+}
+
+// parseTileSize accepts "512" for a square tile and "512x384" for a
+// rectangular one.
+//
+// A value it cannot make sense of is reported as absent rather than as a zero.
+// The caller only calls the upstream setter when a size was given, so a typo
+// leaves the library's own default in place instead of installing a degenerate
+// tiling that would fail at generation time.
+func parseTileSize(value string) (int, int, bool) {
+	xs, ys, split := strings.Cut(value, "x")
+	if !split {
+		ys = xs
+	}
+	x, err := strconv.Atoi(xs)
+	if err != nil || x <= 0 {
+		return 0, 0, false
+	}
+	y, err := strconv.Atoi(ys)
+	if err != nil || y <= 0 {
+		return 0, 0, false
+	}
+	return x, y, true
+}
+
 // Copied from Purego internal/strings
 // TODO: We should upstream sending []string
 func hasSuffix(s, suffix string) bool {
@@ -100,6 +158,9 @@ func (sd *SDGGML) Load(opts *pb.ModelOptions) error {
 	}
 
 	sd.cfgScale = opts.CFGScale
+	// Read from the unfiltered list: none of the tiling options name a path, so
+	// the resolution pass above neither rewrites nor drops them.
+	sd.vaeTiling = parseVAETiling(opts.Options)
 
 	ret := LoadModel(modelFile, modelPathC, options, opts.Threads, diffusionModel)
 	runtime.KeepAlive(keepAlive)
@@ -148,8 +209,22 @@ func (sd *SDGGML) GenerateImage(opts *pb.GenerateImageRequest) error {
 	ImgGenParamsSetPrompts(p, t, negative)
 	ImgGenParamsSetDimensions(p, int(opts.Width), int(opts.Height))
 	ImgGenParamsSetSeed(p, int64(opts.Seed))
+	// Tiling decodes the latent in overlapping tiles, so the VAE compute buffer
+	// scales with the tile rather than with the image. That is the difference
+	// between working and failing on any device that caps a single allocation
+	// (RADV reports a 4GiB maxMemoryAllocationSize, for one) or that simply
+	// does not have the VRAM for a full-frame decode at high resolution.
+	//
+	// Only the setters the operator configured are called, so an unset tile
+	// size or overlap keeps the library's own default.
 	vaep := ImgGenParamsGetVaeTilingParams(p)
-	TilingParamsSetEnabled(vaep, false)
+	TilingParamsSetEnabled(vaep, sd.vaeTiling.enabled)
+	if sd.vaeTiling.hasTileSize {
+		TilingParamsSetTileSizes(vaep, sd.vaeTiling.tileSizeX, sd.vaeTiling.tileSizeY)
+	}
+	if sd.vaeTiling.hasOverlap {
+		TilingParamsSetTargetOverlap(vaep, sd.vaeTiling.targetOverlap)
+	}
 
 	ret := GenImage(p, int(opts.Step), dst, sd.cfgScale, srcImage, strength, maskImage, refImages, refImagesCount)
 	runtime.KeepAlive(keepAlive)

backend/go/stablediffusion-ggml/gosd_test.go

@@ -0,0 +1,180 @@
+package main
+
+import (
+	"testing"
+
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestStableDiffusionGGML(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "stablediffusion-ggml backend test suite")
+}
+
+var _ = DescribeTable("parseVAETiling enablement",
+	func(options []string, want bool) {
+		Expect(parseVAETiling(options).enabled).To(Equal(want))
+	},
+	Entry("explicit true", []string{"vae_tiling:true"}, true),
+	Entry("one is truthy", []string{"vae_tiling:1"}, true),
+	// "diffusion_model" is already a bare flag in this option list, so accept
+	// the same shape rather than making vae_tiling the one option that demands
+	// a value.
+	Entry("bare flag", []string{"vae_tiling"}, true),
+	Entry("explicit false", []string{"vae_tiling:false"}, false),
+	Entry("anything else is false", []string{"vae_tiling:maybe"}, false),
+	// Tiling trades a little quality at the tile seams for a much smaller
+	// compute buffer, so it must not switch itself on for the many models that
+	// never needed it.
+	Entry("absent leaves it off", []string{"diffusion_model", "sampler:euler"}, false),
+	Entry("nil options", []string(nil), false),
+)
+
+var _ = DescribeTable("parseVAETiling tile size",
+	func(options []string, wantSet bool, wantX, wantY int) {
+		got := parseVAETiling(options)
+		Expect(got.hasTileSize).To(Equal(wantSet))
+		if wantSet {
+			Expect(got.tileSizeX).To(Equal(wantX))
+			Expect(got.tileSizeY).To(Equal(wantY))
+		}
+	},
+	Entry("one number is a square tile", []string{"vae_tile_size:512"}, true, 512, 512),
+	Entry("rectangular", []string{"vae_tile_size:512x384"}, true, 512, 384),
+	// Absent must stay absent: the caller only invokes the upstream setter when
+	// a size was given, so this is what preserves the library's own default
+	// instead of pushing a zero.
+	Entry("absent", []string{"vae_tiling:true"}, false, 0, 0),
+	// A typo must not silently become a zero tile size and break a model that
+	// would otherwise have worked on the default.
+	Entry("unparseable", []string{"vae_tile_size:banana"}, false, 0, 0),
+	Entry("zero rejected", []string{"vae_tile_size:0"}, false, 0, 0),
+	Entry("negative rejected", []string{"vae_tile_size:-8"}, false, 0, 0),
+	Entry("half unparseable", []string{"vae_tile_size:512xbanana"}, false, 0, 0),
+)
+
+var _ = DescribeTable("parseVAETiling target overlap",
+	func(options []string, wantSet bool, want float32) {
+		got := parseVAETiling(options)
+		Expect(got.hasOverlap).To(Equal(wantSet))
+		if wantSet {
+			Expect(got.targetOverlap).To(Equal(want))
+		}
+	},
+	Entry("fraction", []string{"vae_tile_overlap:0.25"}, true, float32(0.25)),
+	Entry("zero is a legitimate overlap", []string{"vae_tile_overlap:0"}, true, float32(0)),
+	Entry("absent", []string{"vae_tiling:true"}, false, float32(0)),
+	Entry("unparseable", []string{"vae_tile_overlap:banana"}, false, float32(0)),
+	Entry("negative rejected", []string{"vae_tile_overlap:-0.5"}, false, float32(0)),
+)
+
+// fakeSDLib swaps the purego bindings for recorders so the wiring between the
+// parsed options and the upstream tiling setters can be exercised without the
+// shared library. The bindings are package-level vars, which is the only seam
+// available here; every one the code under test touches must be set or the
+// call panics on a nil func.
+type fakeSDLib struct {
+	tilingEnabled bool
+	tileSizeCalls int
+	tileSizeX     int
+	tileSizeY     int
+	overlapCalls  int
+	targetOverlap float32
+}
+
+// install points the bindings at the recorder and restores them afterwards, so
+// one spec cannot leak fakes into the next.
+func (f *fakeSDLib) install() {
+	savedImgGenParamsNew := ImgGenParamsNew
+	savedImgGenParamsSetPrompts := ImgGenParamsSetPrompts
+	savedImgGenParamsSetDimensions := ImgGenParamsSetDimensions
+	savedImgGenParamsSetSeed := ImgGenParamsSetSeed
+	savedImgGenParamsGetVaeTilingParams := ImgGenParamsGetVaeTilingParams
+	savedTilingParamsSetEnabled := TilingParamsSetEnabled
+	savedTilingParamsSetTileSizes := TilingParamsSetTileSizes
+	savedTilingParamsSetTargetOverlap := TilingParamsSetTargetOverlap
+	savedGenImage := GenImage
+	savedLoadModel := LoadModel
+
+	DeferCleanup(func() {
+		ImgGenParamsNew = savedImgGenParamsNew
+		ImgGenParamsSetPrompts = savedImgGenParamsSetPrompts
+		ImgGenParamsSetDimensions = savedImgGenParamsSetDimensions
+		ImgGenParamsSetSeed = savedImgGenParamsSetSeed
+		ImgGenParamsGetVaeTilingParams = savedImgGenParamsGetVaeTilingParams
+		TilingParamsSetEnabled = savedTilingParamsSetEnabled
+		TilingParamsSetTileSizes = savedTilingParamsSetTileSizes
+		TilingParamsSetTargetOverlap = savedTilingParamsSetTargetOverlap
+		GenImage = savedGenImage
+		LoadModel = savedLoadModel
+	})
+
+	ImgGenParamsNew = func() uintptr { return 1 }
+	ImgGenParamsSetPrompts = func(uintptr, string, string) {}
+	ImgGenParamsSetDimensions = func(uintptr, int, int) {}
+	ImgGenParamsSetSeed = func(uintptr, int64) {}
+	ImgGenParamsGetVaeTilingParams = func(uintptr) uintptr { return 2 }
+	TilingParamsSetEnabled = func(_ uintptr, enabled bool) { f.tilingEnabled = enabled }
+	TilingParamsSetTileSizes = func(_ uintptr, x, y int) {
+		f.tileSizeCalls++
+		f.tileSizeX, f.tileSizeY = x, y
+	}
+	TilingParamsSetTargetOverlap = func(_ uintptr, o float32) {
+		f.overlapCalls++
+		f.targetOverlap = o
+	}
+	GenImage = func(uintptr, int, string, float32, string, float32, string, []uintptr, int) int { return 0 }
+	LoadModel = func(string, string, []uintptr, int32, int) int { return 0 }
+}
+
+var _ = Describe("GenerateImage VAE tiling", func() {
+	var fake *fakeSDLib
+
+	BeforeEach(func() {
+		fake = &fakeSDLib{}
+		fake.install()
+	})
+
+	// generate drives the real Load and GenerateImage so the options travel the
+	// path they travel in production, with only the C boundary faked.
+	generate := func(options []string) {
+		sd := &SDGGML{}
+		Expect(sd.Load(&pb.ModelOptions{Options: options})).To(Succeed())
+		Expect(sd.GenerateImage(&pb.GenerateImageRequest{Width: 1024, Height: 1024})).To(Succeed())
+	}
+
+	It("enables tiling when the model asks for it", func() {
+		generate([]string{"vae_tiling:true"})
+
+		Expect(fake.tilingEnabled).To(BeTrue())
+	})
+
+	// The pre-existing behaviour: every model that never asked for tiling must
+	// still get it switched off.
+	It("leaves tiling off by default", func() {
+		generate([]string{"sampler:euler"})
+
+		Expect(fake.tilingEnabled).To(BeFalse())
+	})
+
+	It("applies a configured tile size and overlap", func() {
+		generate([]string{"vae_tiling:true", "vae_tile_size:512x384", "vae_tile_overlap:0.25"})
+
+		Expect(fake.tileSizeCalls).To(Equal(1))
+		Expect(fake.tileSizeX).To(Equal(512))
+		Expect(fake.tileSizeY).To(Equal(384))
+		Expect(fake.overlapCalls).To(Equal(1))
+		Expect(fake.targetOverlap).To(Equal(float32(0.25)))
+	})
+
+	// Not calling the setters is what preserves the library's own defaults, so
+	// unconfigured values must leave them untouched rather than send a zero.
+	It("leaves unset tiling parameters alone", func() {
+		generate([]string{"vae_tiling:true"})
+
+		Expect(fake.tileSizeCalls).To(BeZero())
+		Expect(fake.overlapCalls).To(BeZero())
+	})
+})

docs/content/features/image-generation.md

@@ -107,6 +107,34 @@ options:
 `vae_decode_only` is still accepted for backwards compatibility but is now a no-op: upstream removed the flag and the model decides automatically.
 {{% /notice %}}
 
+#### VAE tiling
+
+Sampling runs in a small latent space, but the final VAE decode expands it to full resolution and needs one large compute buffer. At 1024x1024 that buffer can exceed 8 GB, which is enough to fail on a GPU that has plenty of memory overall:
+
+```
+ggml_vulkan: Requested buffer size exceeds device buffer size limit: ErrorOutOfDeviceMemory
+[ERROR] vae: failed to allocate the compute buffer
+[ERROR] stable-diffusion.cpp - decode_first_stage failed for latent 1
+```
+
+Tiling decodes the latent in overlapping tiles, so peak memory scales with the tile instead of the image. It costs a little time and can leave faint seams, so it is off unless you ask for it.
+
+| Option | Example | Description |
+|--------|---------|-------------|
+| `vae_tiling` | `vae_tiling:true` | Enable tiled VAE decoding. The bare form `vae_tiling` also works. |
+| `vae_tile_size` | `vae_tile_size:512` or `vae_tile_size:512x384` | Tile dimensions in pixels. One number means a square tile. Omit to use the upstream default. |
+| `vae_tile_overlap` | `vae_tile_overlap:0.25` | Overlap between neighbouring tiles. More overlap hides seams better and costs more time. Omit to use the upstream default. |
+
+```yaml
+options:
+- "diffusion_model"
+- "sampler:euler"
+- "vae_tiling:true"
+- "vae_tile_size:512"
+```
+
+Reach for this when the decode fails but sampling completed, which the backend log shows as `sampling completed` followed by a VAE allocation error. Two cases hit it: cards without the VRAM for a full-frame decode, and drivers that cap a single allocation regardless of free memory (Mesa RADV reports a 4 GiB `maxMemoryAllocationSize`, so a large decode fails there even with tens of GB free). Lowering the output resolution avoids it too, at the cost of the resolution.
+
 #### Distributed inference (RPC workers)
 
 The `stablediffusion-ggml` backend can offload computation to remote `ggml` RPC workers, sharding a model that does not fit on a single machine. It reuses the **same backend-agnostic `rpc-server` workers as the llama.cpp backend**, so one worker pool can serve both.