diff --git a/backend/go/stablediffusion-ggml/gosd.go b/backend/go/stablediffusion-ggml/gosd.go index 219b78470..e1567bd18 100644 --- a/backend/go/stablediffusion-ggml/gosd.go +++ b/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) diff --git a/backend/go/stablediffusion-ggml/gosd_test.go b/backend/go/stablediffusion-ggml/gosd_test.go new file mode 100644 index 000000000..3476700c6 --- /dev/null +++ b/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()) + }) +}) diff --git a/docs/content/features/image-generation.md b/docs/content/features/image-generation.md index 722bdb510..aca97b236 100644 --- a/docs/content/features/image-generation.md +++ b/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.