diff --git a/backend/go/stablediffusion-ggml/cpp/gosd.cpp b/backend/go/stablediffusion-ggml/cpp/gosd.cpp index 676163a8c..f324df28e 100644 --- a/backend/go/stablediffusion-ggml/cpp/gosd.cpp +++ b/backend/go/stablediffusion-ggml/cpp/gosd.cpp @@ -1,5 +1,7 @@ #include "stable-diffusion.h" #include +#include +#include #include #define GGML_MAX_NAME 128 @@ -1452,3 +1454,66 @@ int unload() { free_sd_ctx(sd_c); return 0; } + + +// Own copies even if upstream retains constructor arguments: never retain a Go +// string pointer across the ABI call. +struct gosd_upscaler { + std::string path, backend, params_backend; + upscaler_ctx_t* ctx = nullptr; + ~gosd_upscaler() { if (ctx) free_upscaler_ctx(ctx); } +}; + +void* gosd_upscaler_create(const char* path, bool direct, int threads, int tile_size, + const char* backend, const char* params_backend) { + if (!path || !*path || tile_size < 0) return nullptr; + try { + auto h = std::make_unique(); + h->path = path; + h->backend = backend ? backend : ""; + h->params_backend = params_backend ? params_backend : ""; + h->ctx = new_upscaler_ctx(h->path.c_str(), direct, threads, tile_size, + h->backend.c_str(), h->params_backend.c_str()); + if (!h->ctx) return nullptr; + return h.release(); + } catch (...) { return nullptr; } +} + +int gosd_upscaler_scale(void* handle) { + if (!handle) return 0; + try { return get_upscale_factor(static_cast(handle)->ctx); } + catch (...) { return 0; } +} + +void gosd_upscaler_destroy(void* handle) { + delete static_cast(handle); +} + +int gosd_upscaler_run(void* handle, const char* src, const char* dst, int scale) { + if (!handle || !src || !*src || !dst || !*dst || scale <= 0) return 1; + // RAII also covers exceptions from native inference. + struct images { + sd_image_t* data = nullptr; + int count = 0; + ~images() { if (data) free_sd_images(data, count); } + } output; + int width = 0, height = 0, channels = 0; + std::unique_ptr input( + stbi_load(src, &width, &height, &channels, 3), stbi_image_free); + if (!input || width <= 0 || height <= 0) return 2; + const uint64_t out_width = uint64_t(width) * scale; + const uint64_t out_height = uint64_t(height) * scale; + if (out_width > INT_MAX / 3 || out_height > INT_MAX || + out_width * out_height > SIZE_MAX / 3) return 4; + sd_image_t image = {uint32_t(width), uint32_t(height), 3, input.get()}; + try { + if (!upscale(static_cast(handle)->ctx, image, scale, + &output.data, &output.count)) return 3; + if (!output.data || output.count != 1 || !output.data[0].data || + output.data[0].width != out_width || output.data[0].height != out_height || + output.data[0].channel != 3) return 4; + if (!stbi_write_png(dst, int(out_width), int(out_height), 3, + output.data[0].data, int(out_width) * 3)) return 5; + return 0; + } catch (...) { return 3; } +} diff --git a/backend/go/stablediffusion-ggml/cpp/gosd.h b/backend/go/stablediffusion-ggml/cpp/gosd.h index 4b94b6009..d0b790998 100644 --- a/backend/go/stablediffusion-ggml/cpp/gosd.h +++ b/backend/go/stablediffusion-ggml/cpp/gosd.h @@ -5,6 +5,13 @@ extern "C" { #endif +// Opaque ESRGAN context. Run returns 0, or argument/decode/inference/output/write +// errors (1..5); all image allocations remain native. +void* gosd_upscaler_create(const char* path, bool direct, int threads, int tile_size, const char* backend, const char* params_backend); +int gosd_upscaler_scale(void* handle); +int gosd_upscaler_run(void* handle, const char* src, const char* dst, int scale); +void gosd_upscaler_destroy(void* handle); + void sd_tiling_params_set_enabled(sd_tiling_params_t *params, bool enabled); void sd_tiling_params_set_tile_sizes(sd_tiling_params_t *params, int tile_size_w, int tile_size_h); void sd_tiling_params_set_rel_sizes(sd_tiling_params_t *params, float rel_size_w, float rel_size_h); diff --git a/backend/go/stablediffusion-ggml/gosd.go b/backend/go/stablediffusion-ggml/gosd.go index 9970cefd5..24d5492bd 100644 --- a/backend/go/stablediffusion-ggml/gosd.go +++ b/backend/go/stablediffusion-ggml/gosd.go @@ -8,6 +8,7 @@ import ( "runtime" "strconv" "strings" + "sync" "unsafe" "github.com/mudler/LocalAI/pkg/grpc/base" @@ -17,6 +18,9 @@ import ( type SDGGML struct { base.SingleThread + upscaleMu sync.Mutex + upscaler uintptr + upscale upscaleSettings threads int sampleMethod string cfgScale float32 @@ -25,9 +29,13 @@ type SDGGML struct { } var ( - LoadModel func(model, model_apth string, options []uintptr, threads int32, diff int) int - GenImage func(params uintptr, steps int, dst string, cfgScale float32, srcImage string, strength float32, maskImage string, refImages []uintptr, refImagesCount int) int - GenVideo func(params uintptr, steps int, dst string, cfgScale float32, fps int, initImage string, endImage string) int + UpscalerCreate func(string, bool, int32, int32, string, string) uintptr + UpscalerScale func(uintptr) int32 + UpscalerRun func(uintptr, string, string, int32) int32 + UpscalerDestroy func(uintptr) + LoadModel func(model, model_apth string, options []uintptr, threads int32, diff int) int + GenImage func(params uintptr, steps int, dst string, cfgScale float32, srcImage string, strength float32, maskImage string, refImages []uintptr, refImagesCount int) int + GenVideo func(params uintptr, steps int, dst string, cfgScale float32, fps int, initImage string, endImage string) int TilingParamsSetEnabled func(params uintptr, enabled bool) TilingParamsSetTileSizes func(params uintptr, tileSizeX int, tileSizeY int) @@ -154,6 +162,22 @@ func CString(name string) *byte { } func (sd *SDGGML) Load(opts *pb.ModelOptions) error { + sd.upscaleMu.Lock() + defer sd.upscaleMu.Unlock() + if opts == nil { + return fmt.Errorf("missing model options") + } + settings, err := parseUpscaleSettings(opts) + if err != nil { + return err + } + if sd.upscaler != 0 { + return fmt.Errorf("free the existing upscaler before loading another model") + } + sd.upscale = settings + if settings.enabled { + return nil + } sd.threads = int(opts.Threads) @@ -336,3 +360,175 @@ func (sd *SDGGML) GenerateVideo(opts *pb.GenerateVideoRequest) error { } return nil } + +// Presence of the backend option upscale_scale selects upscale mode. +type upscaleSettings struct { + enabled bool + path string + scale, tile, threads int32 + direct bool + backend, paramsBackend string +} + +func parseUpscaleSettings(opts *pb.ModelOptions) (upscaleSettings, error) { + s := upscaleSettings{threads: opts.Threads} + image := false + for _, op := range opts.Options { + k, _, _ := strings.Cut(op, ":") + if k == "upscale_scale" { + s.enabled = true + } + if k == "diffusion_model" { + image = true + } + } + if !s.enabled { + return s, nil + } + if image { + return s, fmt.Errorf("upscale_scale cannot be combined with diffusion_model") + } + seen := make(map[string]bool) + for _, op := range opts.Options { + k, v, _ := strings.Cut(op, ":") + switch k { + case "upscale_scale", "upscale_tile_size", "upscale_direct", "backend", "params_backend": + if seen[k] { + return s, fmt.Errorf("duplicate upscale setting %q", k) + } + seen[k] = true + } + switch k { + case "upscale_scale", "upscale_tile_size": + n, err := strconv.ParseInt(v, 10, 32) + if err != nil || (k == "upscale_scale" && n <= 0) || (k == "upscale_tile_size" && n < 0) { + if k == "upscale_tile_size" { + return s, fmt.Errorf("%s must be a non-negative integer", k) + } + return s, fmt.Errorf("%s must be a positive integer", k) + } + if k == "upscale_scale" { + s.scale = int32(n) + } else { + s.tile = int32(n) + } + case "upscale_direct": + b, err := strconv.ParseBool(v) + if err != nil { + return s, fmt.Errorf("invalid upscale_direct: %w", err) + } + s.direct = b + case "backend": + s.backend = v + case "params_backend": + s.paramsBackend = v + } + } + if s.scale <= 0 { + return s, fmt.Errorf("upscale_scale must be configured and positive") + } + modelFile := opts.ModelFile + if modelFile == "" { + return s, fmt.Errorf("upscale model path is empty") + } + if filepath.IsAbs(modelFile) { + s.path = filepath.Clean(modelFile) + } else { + // Verify the untrusted relative name before joining, then verify the + // resolved result as well. VerifyPath intentionally treats absolute + // inputs as relative, so absolute model files use VerifyResolvedPath. + if err := utils.VerifyPath(modelFile, opts.ModelPath); err != nil { + return s, fmt.Errorf("upscale model path is outside model path: %w", err) + } + s.path = filepath.Join(opts.ModelPath, modelFile) + } + if err := utils.VerifyResolvedPath(s.path, opts.ModelPath); err != nil { + return s, fmt.Errorf("upscale model path is outside model path: %w", err) + } + // Lexical containment alone does not prevent a model symlink (or a + // symlinked parent directory) from escaping the trusted root. Resolve + // both, requiring an existing target, and cache only the real model path. + realRoot, err := filepath.EvalSymlinks(opts.ModelPath) + if err != nil { + return s, fmt.Errorf("resolve upscale model root: %w", err) + } + realPath, err := filepath.EvalSymlinks(s.path) + if err != nil { + return s, fmt.Errorf("resolve upscale model file: %w", err) + } + if err := utils.VerifyResolvedPath(realPath, realRoot); err != nil { + return s, fmt.Errorf("upscale model path is outside real model path: %w", err) + } + info, err := os.Stat(realPath) + if err != nil { + return s, fmt.Errorf("stat upscale model file: %w", err) + } + if !info.Mode().IsRegular() { + return s, fmt.Errorf("upscale model path must be a regular file") + } + s.path = realPath + for _, v := range []string{s.path, s.backend, s.paramsBackend} { + if strings.ContainsRune(v, 0) { + return s, fmt.Errorf("upscale settings contain NUL") + } + } + return s, nil +} + +func (sd *SDGGML) UpscaleImage(opts *pb.UpscaleImageRequest) error { + sd.upscaleMu.Lock() + defer sd.upscaleMu.Unlock() + s := sd.upscale + if !s.enabled { + return fmt.Errorf("no upscale-only model loaded (configure upscale_scale)") + } + if opts == nil || opts.Scale <= 0 || opts.Src == "" || opts.Dst == "" || strings.ContainsRune(opts.Src, 0) || strings.ContainsRune(opts.Dst, 0) { + return fmt.Errorf("upscale requires source, destination and positive requested scale") + } + if sd.upscaler == 0 { + h := UpscalerCreate(s.path, s.direct, s.threads, s.tile, s.backend, s.paramsBackend) + runtime.KeepAlive(s) + if h == 0 { + return fmt.Errorf("could not initialize ESRGAN upscaler") + } + detected := UpscalerScale(h) + if detected <= 0 || detected != s.scale { + UpscalerDestroy(h) + return fmt.Errorf("configured upscale scale %d differs from detected native scale %d", s.scale, detected) + } + sd.upscaler = h + } + if opts.Scale != s.scale { + return fmt.Errorf("requested upscale scale %d differs from native scale %d", opts.Scale, s.scale) + } + code := UpscalerRun(sd.upscaler, opts.Src, opts.Dst, opts.Scale) + runtime.KeepAlive(opts) + switch code { + case 0: + return nil + case 1: + return fmt.Errorf("invalid native upscale arguments") + case 2: + return fmt.Errorf("upscale source decode failed") + case 3: + return fmt.Errorf("upscale inference failed") + case 4: + return fmt.Errorf("upscale output validation failed") + case 5: + return fmt.Errorf("upscale PNG write failed") + default: + return fmt.Errorf("upscale failed (native code %d)", code) + } +} + +func (sd *SDGGML) Free() error { + sd.upscaleMu.Lock() + defer sd.upscaleMu.Unlock() + if sd.upscaler != 0 { + UpscalerDestroy(sd.upscaler) + sd.upscaler = 0 + } + sd.upscale = upscaleSettings{} + // Deliberately leave the existing diffusion sd_c lifecycle unchanged. + return nil +} diff --git a/backend/go/stablediffusion-ggml/gosd_test.go b/backend/go/stablediffusion-ggml/gosd_test.go index beca7b3da..4b0cc0b59 100644 --- a/backend/go/stablediffusion-ggml/gosd_test.go +++ b/backend/go/stablediffusion-ggml/gosd_test.go @@ -1,6 +1,8 @@ package main import ( + "os" + "path/filepath" "testing" "unsafe" @@ -222,3 +224,313 @@ var _ = Describe("GenerateImage VAE tiling", func() { Expect(fake.overlapCalls).To(BeZero()) }) }) + +var _ = Describe("ESRGAN lifecycle", func() { + var sd *SDGGML + var root, modelPath string + var creates, destroys, runs, loads int + var scale, code int32 + var fail bool + var entered, release chan struct{} + BeforeEach(func() { + sd = &SDGGML{} + root = GinkgoT().TempDir() + var err error + root, err = filepath.EvalSymlinks(root) + Expect(err).NotTo(HaveOccurred()) + modelPath = filepath.Join(root, "esrgan.gguf") + Expect(os.WriteFile(modelPath, []byte("model"), 0600)).To(Succeed()) + creates, destroys, runs, loads = 0, 0, 0, 0 + scale, code, fail = 4, 0, false + entered, release = nil, nil + oldCreate, oldScale, oldRun, oldDestroy, oldLoad := UpscalerCreate, UpscalerScale, UpscalerRun, UpscalerDestroy, LoadModel + DeferCleanup(func() { + UpscalerCreate, UpscalerScale, UpscalerRun, UpscalerDestroy, LoadModel = oldCreate, oldScale, oldRun, oldDestroy, oldLoad + }) + UpscalerCreate = func(path string, direct bool, threads, tile int32, backend, params string) uintptr { + creates++ + Expect(path).To(Equal(modelPath)) + Expect(tile).To(Equal(int32(128))) + if fail { + return 0 + } + return 42 + } + UpscalerScale = func(uintptr) int32 { return scale } + UpscalerRun = func(uintptr, string, string, int32) int32 { + runs++ + if entered != nil { + close(entered) + <-release + } + return code + } + UpscalerDestroy = func(h uintptr) { Expect(h).To(Equal(uintptr(42))); destroys++ } + LoadModel = func(string, string, []uintptr, int32, int) int { loads++; return 0 } + }) + load := func(extra ...string) error { + return sd.Load(&pb.ModelOptions{ModelPath: root, ModelFile: "esrgan.gguf", Options: append([]string{"upscale_scale:4", "upscale_tile_size:128"}, extra...)}) + } + run := func(s int32) error { + return sd.UpscaleImage(&pb.UpscaleImageRequest{Src: "in.png", Dst: "out.png", Scale: s}) + } + It("skips diffusion, creates lazily, reuses and frees once", func() { + Expect(load()).To(Succeed()) + Expect(loads).To(BeZero()) + Expect(creates).To(BeZero()) + Expect(run(4)).To(Succeed()) + Expect(run(4)).To(Succeed()) + Expect(creates).To(Equal(1)) + Expect(runs).To(Equal(2)) + Expect(sd.Free()).To(Succeed()) + Expect(sd.Free()).To(Succeed()) + Expect(destroys).To(Equal(1)) + Expect(run(4)).NotTo(Succeed()) + }) + It("rejects invalid configuration", func() { + Expect(load("upscale_scale:0")).NotTo(Succeed()) + Expect(load("upscale_tile_size:-1")).NotTo(Succeed()) + }) + It("rejects detected scale mismatch without caching", func() { + Expect(load()).To(Succeed()) + scale = 2 + Expect(run(4)).To(MatchError(ContainSubstring("configured"))) + Expect(destroys).To(Equal(1)) + scale = 4 + Expect(run(4)).To(Succeed()) + Expect(creates).To(Equal(2)) + }) + It("rejects requested scale mismatch", func() { + Expect(load()).To(Succeed()) + Expect(run(2)).To(MatchError(ContainSubstring("requested"))) + Expect(runs).To(BeZero()) + }) + It("does not cache failed initialization", func() { + Expect(load()).To(Succeed()) + fail = true + Expect(run(4)).NotTo(Succeed()) + fail = false + Expect(run(4)).To(Succeed()) + Expect(creates).To(Equal(2)) + }) + It("rejects undetected scales", func() { + Expect(load()).To(Succeed()) + scale = 0 + Expect(run(4)).NotTo(Succeed()) + Expect(destroys).To(Equal(1)) + }) + It("reports native failures", func() { + Expect(load()).To(Succeed()) + for _, c := range []int32{1, 2, 3, 4, 5} { + code = c + Expect(run(4)).NotTo(Succeed()) + } + }) + It("waits for inference before Free", func() { + Expect(load()).To(Succeed()) + Expect(run(4)).To(Succeed()) + entered, release = make(chan struct{}), make(chan struct{}) + done := make(chan error, 1) + freed := make(chan error, 1) + go func() { done <- run(4) }() + <-entered + go func() { freed <- sd.Free() }() + Consistently(freed, "50ms").ShouldNot(Receive()) + close(release) + Eventually(done).Should(Receive(BeNil())) + Eventually(freed).Should(Receive(BeNil())) + Expect(destroys).To(Equal(1)) + }) +}) + +func TestUpscaleSettings(t *testing.T) { + root := upscaleTestRoot(t) + opts := &pb.ModelOptions{ModelPath: root, ModelFile: filepath.Join(root, "model.gguf"), Threads: 7, Options: []string{ + "upscale_scale:4", "upscale_tile_size:256", + "upscale_direct:true", "backend:Vulkan0", "params_backend:CPU", + }} + got, err := parseUpscaleSettings(opts) + if err != nil { + t.Fatal(err) + } + if !got.enabled || got.path != opts.ModelFile || got.scale != 4 || got.tile != 256 || got.threads != 7 || !got.direct || got.backend != "Vulkan0" || got.paramsBackend != "CPU" { + t.Fatalf("wrong settings: %+v", got) + } + for _, options := range [][]string{ + {"upscale_scale:banana"}, {"upscale_scale:2147483648"}, + {"upscale_scale:4", "upscale_tile_size:-1"}, + {"upscale_scale:4", "diffusion_model"}, {"upscale_scale:4", "upscale_direct:maybe"}, + } { + opts.Options = options + if _, err := parseUpscaleSettings(opts); err == nil { + t.Fatalf("accepted invalid options %v", options) + } + } +} + +func TestUpscaleModeSelection(t *testing.T) { + root := upscaleTestRoot(t) + + upscale, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: "model.gguf", Options: []string{"upscale_scale:4"}}) + if err != nil || !upscale.enabled { + t.Fatalf("upscale_scale did not enable upscale mode: settings=%+v err=%v", upscale, err) + } + + knownUsecases, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: "model.gguf", Options: []string{"known_usecases:upscale"}}) + if err != nil || knownUsecases.enabled { + t.Fatalf("known_usecases must not enable upscale mode: settings=%+v err=%v", knownUsecases, err) + } + + if _, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: "model.gguf", Options: []string{"upscale_scale:4", "diffusion_model"}}); err == nil { + t.Fatal("accepted upscale_scale combined with diffusion_model") + } + + savedLoadModel := LoadModel + defer func() { LoadModel = savedLoadModel }() + loads := 0 + diffusion := -1 + LoadModel = func(_ string, _ string, _ []uintptr, _ int32, diffusionModel int) int { + loads++ + diffusion = diffusionModel + return 0 + } + if err := (&SDGGML{}).Load(&pb.ModelOptions{ModelPath: root, ModelFile: "model.gguf", Options: []string{"diffusion_model"}}); err != nil { + t.Fatalf("Load without upscale_scale: %v", err) + } + if loads != 1 || diffusion != 1 { + t.Fatalf("Load without upscale_scale did not use diffusion path: loads=%d diffusion=%d", loads, diffusion) + } +} + +func TestUpscaleSettingsFullFrameTile(t *testing.T) { + root := upscaleTestRoot(t) + for _, options := range [][]string{ + {"upscale_scale:4"}, + {"upscale_scale:4", "upscale_tile_size:0"}, + } { + got, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: "model.gguf", Options: options}) + if err != nil { + t.Fatalf("parseUpscaleSettings(%v): %v", options, err) + } + if got.tile != 0 { + t.Fatalf("tile for %v = %d, want full-frame 0", options, got.tile) + } + } +} + +func TestUpscaleSettingsPathContainment(t *testing.T) { + root := upscaleTestRoot(t) + if err := os.Mkdir(filepath.Join(root, "nested"), 0700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "nested/model.gguf"), []byte("model"), 0600); err != nil { + t.Fatal(err) + } + for _, tc := range []struct { + name, modelFile, want string + wantErr bool + }{ + {"safe nested relative path", "nested/model.gguf", filepath.Join(root, "nested/model.gguf"), false}, + {"relative traversal", "../outside.gguf", "", true}, + {"allowed in-root absolute path", filepath.Join(root, "model.gguf"), filepath.Join(root, "model.gguf"), false}, + {"outside absolute path", "/other/model.gguf", "", true}, + } { + t.Run(tc.name, func(t *testing.T) { + got, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: tc.modelFile, Options: []string{"upscale_scale:4"}}) + if tc.wantErr { + if err == nil { + t.Fatal("accepted path outside model path") + } + return + } + if err != nil { + t.Fatal(err) + } + if got.path != tc.want { + t.Fatalf("path = %q, want %q", got.path, tc.want) + } + }) + } +} + +func TestUpscaleSettingsRejectsDuplicateSettings(t *testing.T) { + root := upscaleTestRoot(t) + for _, setting := range []string{"upscale_scale:4", "upscale_tile_size:128", "upscale_direct:true", "backend:Vulkan", "params_backend:CPU"} { + t.Run(setting, func(t *testing.T) { + _, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: "model.gguf", Options: []string{"upscale_scale:4", setting, setting}}) + if err == nil { + t.Fatalf("accepted duplicate %s", setting) + } + }) + } +} + +func upscaleTestRoot(t *testing.T) string { + t.Helper() + root, err := filepath.EvalSymlinks(t.TempDir()) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "model.gguf"), []byte("model"), 0600); err != nil { + t.Fatal(err) + } + return root +} + +func TestUpscaleSettingsRealPathContainment(t *testing.T) { + root := upscaleTestRoot(t) + outside := upscaleTestRoot(t) + for _, tc := range []struct { + name, target string + wantErr bool + }{ + {"in-root symlink", filepath.Join(root, "model.gguf"), false}, + {"outside symlink", filepath.Join(outside, "model.gguf"), true}, + {"dangling symlink", filepath.Join(root, "missing.gguf"), true}, + {"outside directory symlink", outside, true}, + } { + t.Run(tc.name, func(t *testing.T) { + link := filepath.Join(root, tc.name) + if err := os.Symlink(tc.target, link); err != nil { + t.Skipf("symlink creation unsupported: %v", err) + } + file := link + if tc.name == "outside directory symlink" { + file = filepath.Join(link, "model.gguf") + } + got, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: file, Options: []string{"upscale_scale:4"}}) + if tc.wantErr { + if err == nil { + t.Fatal("accepted unsafe symlink") + } + return + } + if err != nil { + t.Fatal(err) + } + if got.path != tc.target { + t.Fatalf("cached path %q, want real path %q", got.path, tc.target) + } + }) + } + t.Run("symlinked trusted root", func(t *testing.T) { + link := filepath.Join(t.TempDir(), "root") + if err := os.Symlink(root, link); err != nil { + t.Skipf("symlink creation unsupported: %v", err) + } + got, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: link, ModelFile: "model.gguf", Options: []string{"upscale_scale:4"}}) + if err != nil { + t.Fatal(err) + } + if got.path != filepath.Join(root, "model.gguf") { + t.Fatalf("not resolved: %q", got.path) + } + }) + for _, file := range []string{"missing.gguf", "."} { + t.Run(file, func(t *testing.T) { + if _, err := parseUpscaleSettings(&pb.ModelOptions{ModelPath: root, ModelFile: file, Options: []string{"upscale_scale:4"}}); err == nil { + t.Fatal("accepted missing or non-file model") + } + }) + } +} diff --git a/backend/go/stablediffusion-ggml/main.go b/backend/go/stablediffusion-ggml/main.go index 2f2badca8..be288fe69 100644 --- a/backend/go/stablediffusion-ggml/main.go +++ b/backend/go/stablediffusion-ggml/main.go @@ -35,6 +35,10 @@ func main() { } libFuncs := []LibFuncs{ + {&UpscalerCreate, "gosd_upscaler_create"}, + {&UpscalerScale, "gosd_upscaler_scale"}, + {&UpscalerRun, "gosd_upscaler_run"}, + {&UpscalerDestroy, "gosd_upscaler_destroy"}, {&LoadModel, "load_model"}, {&GenImage, "gen_image"}, {&GenVideo, "gen_video"}, diff --git a/core/backend/options_internal_test.go b/core/backend/options_internal_test.go index 13ad13b57..a83ccef69 100644 --- a/core/backend/options_internal_test.go +++ b/core/backend/options_internal_test.go @@ -451,3 +451,21 @@ var _ = Describe("effectiveThreads", func() { Expect(effectiveThreads(config.ModelConfig{}, 0)).To(Equal(1)) }) }) + +var _ = Describe("grpcModelOpts option forwarding", func() { + It("preserves explicit upscale backend options exactly", func() { + threads := 1 + options := []string{"upscale_scale:2", "upscale_scale:8", "upscale_tile_size:-1", "backend:CPU"} + cfg := config.ModelConfig{Threads: &threads, Backend: "stablediffusion-ggml", KnownUsecases: config.GetUsecasesFromYAML([]string{"upscale", "image"}), Options: options} + Expect(grpcModelOpts(cfg, "/models").Options).To(Equal(options)) + Expect(cfg.Options).To(Equal(options)) + }) + + It("leaves generic SD options unchanged", func() { + threads := 1 + cfg := config.ModelConfig{Threads: &threads, Backend: "stablediffusion-ggml", Options: []string{"backend:CPU", "diffusion_model:model.gguf"}} + Expect(grpcModelOpts(cfg, "/models").Options).To(Equal(cfg.Options)) + cfg.KnownUsecases = config.GetUsecasesFromYAML([]string{"image"}) + Expect(grpcModelOpts(cfg, "/models").Options).To(Equal(cfg.Options)) + }) +}) diff --git a/core/backend/upscale.go b/core/backend/upscale.go index ad3885174..48c12491a 100644 --- a/core/backend/upscale.go +++ b/core/backend/upscale.go @@ -2,6 +2,7 @@ package backend import ( "context" + "fmt" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/pkg/grpc/proto" @@ -24,7 +25,7 @@ func ImageUpscale(ctx context.Context, src, dst string, scale int, loader *model return err } defer release() - _, err = inferenceModel.UpscaleImage( + result, err := inferenceModel.UpscaleImage( ctx, &proto.UpscaleImageRequest{ Src: src, @@ -32,7 +33,19 @@ func ImageUpscale(ctx context.Context, src, dst string, scale int, loader *model Scale: int32(scale), }, ) - return err + if err != nil { + return err + } + if result == nil { + return fmt.Errorf("image upscale returned nil result") + } + if !result.Success { + if result.Message != "" { + return fmt.Errorf("image upscale failed: %s", result.Message) + } + return fmt.Errorf("image upscale failed") + } + return nil } return fn, nil diff --git a/core/backend/upscale_test.go b/core/backend/upscale_test.go new file mode 100644 index 000000000..7e654ea7c --- /dev/null +++ b/core/backend/upscale_test.go @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: MIT + +package backend_test + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + grpcPkg "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + ggrpc "google.golang.org/grpc" +) + +type upscaleBackend struct { + grpcPkg.Backend + result *pb.Result + err error + request *pb.UpscaleImageRequest +} + +func (*upscaleBackend) HealthCheck(context.Context) (bool, error) { return true, nil } +func (*upscaleBackend) IsBusy() bool { return false } +func (b *upscaleBackend) UpscaleImage(_ context.Context, in *pb.UpscaleImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + b.request = in + return b.result, b.err +} + +func TestImageUpscaleResult(t *testing.T) { + transportErr := errors.New("transport failed") + for _, tc := range []struct { + name string + result *pb.Result + err error + want string + }{ + {name: "nil", want: "nil"}, + {name: "unsuccessful", result: &pb.Result{}, want: "failed"}, + {name: "message", result: &pb.Result{Message: "native mismatch"}, want: "native mismatch"}, + {name: "transport", err: transportErr, want: "transport failed"}, + {name: "success", result: &pb.Result{Success: true}}, + } { + t.Run(tc.name, func(t *testing.T) { + fake := &upscaleBackend{result: tc.result, err: tc.err} + loader := model.NewModelLoader(&system.SystemState{}) + loader.SetModelRouter(func(_ context.Context, id string, _, _, _, _ string, _ *pb.ModelOptions, _ bool) (*model.Model, error) { + return model.NewModelWithClient(id, "test://upscale", fake), nil + }) + cfg := config.ModelConfig{Name: "upscale", Backend: "stub"} + cfg.SetDefaults() + fn, err := backend.ImageUpscale(context.Background(), "source", "destination", 3, loader, cfg, config.NewApplicationConfig(config.WithSystemState(&system.SystemState{}))) + if err != nil { + t.Fatal(err) + } + err = fn() + if tc.want == "" { + if err != nil { + t.Fatal(err) + } + } else if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("error = %v, want %q", err, tc.want) + } + if tc.err != nil && !errors.Is(err, tc.err) { + t.Fatalf("lost transport error: %v", err) + } + if fake.request == nil || fake.request.Scale != 3 || fake.request.Src != "source" || fake.request.Dst != "destination" { + t.Fatalf("unexpected request: %v", fake.request) + } + }) + } +} diff --git a/core/config/backend_capabilities.go b/core/config/backend_capabilities.go index 96d53c4e2..9de40568c 100644 --- a/core/config/backend_capabilities.go +++ b/core/config/backend_capabilities.go @@ -17,6 +17,7 @@ const ( UsecaseEmbeddings = "embeddings" UsecaseTokenize = "tokenize" UsecaseImage = "image" + UsecaseUpscale = "upscale" UsecaseVideo = "video" Usecase3D = "3d" Usecase3DAnimation = "3d_animation" @@ -122,6 +123,11 @@ var UsecaseInfoMap = map[string]UsecaseInfo{ GRPCMethod: MethodTokenizeString, Description: "Tokenization via the TokenizeString RPC without running inference.", }, + UsecaseUpscale: { + Flag: FLAG_UPSCALE, + GRPCMethod: MethodUpscaleImage, + Description: "Image upscaling via the UpscaleImage RPC.", + }, UsecaseImage: { Flag: FLAG_IMAGE, GRPCMethod: MethodGenerateImage, @@ -406,7 +412,7 @@ var BackendCapabilities = map[string]BackendCapability{ // --- Image/video generation backends --- "diffusers": { GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodUpscaleImage, MethodGenerateVideo, MethodSoundGeneration}, - PossibleUsecases: []string{UsecaseImage, UsecaseVideo, UsecaseSoundGeneration}, + PossibleUsecases: []string{UsecaseImage, UsecaseUpscale, UsecaseVideo, UsecaseSoundGeneration}, DefaultUsecases: []string{UsecaseImage}, Description: "HuggingFace diffusers — image, video, and sound generation", }, @@ -432,8 +438,8 @@ var BackendCapabilities = map[string]BackendCapability{ Description: "Stable Diffusion native backend", }, "stablediffusion-ggml": { - GRPCMethods: []GRPCMethod{MethodGenerateImage}, - PossibleUsecases: []string{UsecaseImage}, + GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodUpscaleImage}, + PossibleUsecases: []string{UsecaseImage, UsecaseUpscale}, DefaultUsecases: []string{UsecaseImage}, Description: "Stable Diffusion via GGML quantized models", }, diff --git a/core/config/backend_capabilities_test.go b/core/config/backend_capabilities_test.go index 6f683c0bf..8ac46c6ae 100644 --- a/core/config/backend_capabilities_test.go +++ b/core/config/backend_capabilities_test.go @@ -445,3 +445,15 @@ var _ = Describe("AllBackendNames", func() { Expect(slices.IsSorted(names)).To(BeTrue()) }) }) + +var _ = Describe("Upscale backend declaration", func() { + It("maps upscale to its dedicated RPC and advertises SD GGML support", func() { + info := UsecaseInfoMap[UsecaseUpscale] + Expect(info.Flag).To(Equal(FLAG_UPSCALE)) + Expect(info.GRPCMethod).To(Equal(MethodUpscaleImage)) + cap := GetBackendCapability("stablediffusion-ggml") + Expect(cap.GRPCMethods).To(ContainElement(MethodUpscaleImage)) + Expect(cap.PossibleUsecases).To(ContainElement(UsecaseUpscale)) + Expect(cap.DefaultUsecases).NotTo(ContainElement(UsecaseUpscale)) + }) +}) diff --git a/core/config/meta/constants.go b/core/config/meta/constants.go index cc1027f2b..008afe64f 100644 --- a/core/config/meta/constants.go +++ b/core/config/meta/constants.go @@ -62,6 +62,7 @@ var UsecaseOptions = []FieldOption{ {Value: "rerank", Label: "Rerank"}, {Value: "score", Label: "Score (Router Classifier)"}, {Value: "image", Label: "Image"}, + {Value: "upscale", Label: "Image upscaling"}, {Value: "vision", Label: "Vision"}, {Value: "detection", Label: "Detection"}, {Value: "depth", Label: "Depth"}, diff --git a/core/config/meta/registry_test.go b/core/config/meta/registry_test.go index 221533050..bd8c43bb8 100644 --- a/core/config/meta/registry_test.go +++ b/core/config/meta/registry_test.go @@ -85,3 +85,15 @@ var _ = Describe("Decisions router metadata", func() { } }) }) + +var _ = Describe("backend-only upscale options", func() { + It("does not expose scale or tile size as generic root configuration", func() { + reg := meta.DefaultRegistry() + Expect(reg).NotTo(HaveKey("upscale_scale")) + Expect(reg).NotTo(HaveKey("upscale_tile_size")) + md := meta.BuildForTest(reflect.TypeOf(config.ModelConfig{}), reg) + for _, field := range md.Fields { + Expect(field.Path).NotTo(BeElementOf("upscale_scale", "upscale_tile_size")) + } + }) +}) diff --git a/core/config/model_capabilities.go b/core/config/model_capabilities.go index 36a2f18da..5d99510f3 100644 --- a/core/config/model_capabilities.go +++ b/core/config/model_capabilities.go @@ -162,6 +162,7 @@ func (c *ModelConfig) Capabilities() []string { add(c.HasUsecases(FLAG_TTS), UsecaseTTS) add(c.HasUsecases(FLAG_SOUND_GENERATION), UsecaseSoundGeneration) add(c.HasUsecases(FLAG_IMAGE), UsecaseImage) + add(c.HasUsecases(FLAG_UPSCALE), UsecaseUpscale) add(c.HasUsecases(FLAG_VIDEO), UsecaseVideo) add(c.HasUsecases(FLAG_3D), Usecase3D) add(c.HasUsecases(FLAG_3D_ANIMATION), Usecase3DAnimation) @@ -207,7 +208,7 @@ func (c *ModelConfig) InputModalities() []string { // Ollama surface); detection/depth/face/3D models consume images directly. imageIn := (chatish && c.VisionSupported()) || c.LimitMMPerPrompt.LimitImagePerPrompt > 0 || c.HasUsecases(FLAG_DETECTION) || c.HasUsecases(FLAG_DEPTH) || c.HasUsecases(FLAG_FACE_RECOGNITION) || - c.HasUsecases(FLAG_3D) + c.HasUsecases(FLAG_3D) || c.HasUsecases(FLAG_UPSCALE) audioIn := c.AudioInputSupported() || c.HasUsecases(FLAG_TRANSCRIPT) || c.HasUsecases(FLAG_AUDIO_TRANSFORM) || c.HasUsecases(FLAG_REALTIME_AUDIO) || c.HasUsecases(FLAG_VAD) || c.HasUsecases(FLAG_DIARIZATION) || @@ -228,7 +229,7 @@ func (c *ModelConfig) OutputModalities() []string { modalities := declaredModalities(c.KnownOutputModalities) textOut := c.HasUsecases(FLAG_CHAT) || c.HasUsecases(FLAG_COMPLETION) || c.HasUsecases(FLAG_EDIT) || c.HasUsecases(FLAG_TRANSCRIPT) - imageOut := c.HasUsecases(FLAG_IMAGE) + imageOut := c.HasUsecases(FLAG_IMAGE) || c.HasUsecases(FLAG_UPSCALE) audioOut := c.HasUsecases(FLAG_TTS) || c.HasUsecases(FLAG_SOUND_GENERATION) || c.HasUsecases(FLAG_AUDIO_TRANSFORM) || c.HasUsecases(FLAG_REALTIME_AUDIO) videoOut := c.HasUsecases(FLAG_VIDEO) diff --git a/core/config/model_capabilities_test.go b/core/config/model_capabilities_test.go index e228d1cb8..e86460b52 100644 --- a/core/config/model_capabilities_test.go +++ b/core/config/model_capabilities_test.go @@ -207,3 +207,17 @@ var _ = Describe("Model capabilities derivation", func() { }) }) }) + +var _ = Describe("Upscale capabilities", func() { + It("keeps dedicated upscalers image-only without guessing generation", func() { + for _, backend := range []string{"stablediffusion-ggml", "llama-cpp", ""} { + cfg := &ModelConfig{Backend: backend, KnownUsecases: GetUsecasesFromYAML([]string{"upscale"})} + Expect(cfg.Capabilities()).To(Equal([]string{UsecaseUpscale})) + Expect(cfg.InputModalities()).To(Equal([]string{ModalityImage})) + Expect(cfg.OutputModalities()).To(Equal([]string{ModalityImage})) + Expect(cfg.HasUsecases(FLAG_IMAGE)).To(BeFalse()) + Expect((&ModelConfig{Backend: backend}).HasUsecases(FLAG_UPSCALE)).To(BeFalse()) + } + Expect(*GetUsecasesFromYAML([]string{"FLAG_UPSCALE"})).To(Equal(FLAG_UPSCALE)) + }) +}) diff --git a/core/config/model_config.go b/core/config/model_config.go index b2b727628..9b280a7d1 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -2088,6 +2088,9 @@ const ( // pickers it cannot serve. FLAG_DECISIONS ModelConfigUsecase = 1 << 25 + // Marks a model as an image upscaler served by the UpscaleImage RPC. + FLAG_UPSCALE ModelConfigUsecase = 1 << 26 + // Common Subsets FLAG_LLM ModelConfigUsecase = FLAG_CHAT | FLAG_COMPLETION | FLAG_EDIT ) @@ -2151,6 +2154,7 @@ func GetAllModelConfigUsecases() map[string]ModelConfigUsecase { "FLAG_3D": FLAG_3D, "FLAG_3D_ANIMATION": FLAG_3D_ANIMATION, "FLAG_DECISIONS": FLAG_DECISIONS, + "FLAG_UPSCALE": FLAG_UPSCALE, } } @@ -2179,7 +2183,7 @@ func GetUsecasesFromYAML(input []string) *ModelConfigUsecase { // // Declared known_usecases are normally additive — the guessing heuristic // still adds whatever it can infer from backend/templates. The exceptions -// are FLAG_SCORE, FLAG_TOKEN_CLASSIFY and FLAG_DECISIONS: when the operator +// are FLAG_SCORE, FLAG_TOKEN_CLASSIFY, FLAG_DECISIONS and FLAG_UPSCALE: when the operator // declared any of them, they reserved the model for a direct-decode primitive // (the router classifier, the PII NER tier, or a decision head). Letting GuessUsecases // paint chat/completion/embeddings on top would surface it in pickers it @@ -2191,7 +2195,7 @@ func (c *ModelConfig) HasUsecases(u ModelConfigUsecase) bool { if (u & *c.KnownUsecases) == u { return true } - if (*c.KnownUsecases & (FLAG_SCORE | FLAG_TOKEN_CLASSIFY | FLAG_DECISIONS)) != 0 { + if (*c.KnownUsecases & (FLAG_SCORE | FLAG_TOKEN_CLASSIFY | FLAG_DECISIONS | FLAG_UPSCALE)) != 0 { return false } } @@ -2202,6 +2206,10 @@ func (c *ModelConfig) HasUsecases(u ModelConfigUsecase) bool { // In its current state, this function should ideally check for properties of the config like templates, rather than the direct backend name checks for the lower half. // This avoids the maintenance burden of updating this list for each new backend - but unfortunately, that's the best option for some services currently. func (c *ModelConfig) GuessUsecases(u ModelConfigUsecase) bool { + // Upscaling requires a dedicated model declaration, not just an image backend. + if u&FLAG_UPSCALE != 0 { + return false + } // Backends that are clearly not text-generation nonTextGenBackends := []string{ "whisper", "piper", "kokoro", diff --git a/core/config/model_config_test.go b/core/config/model_config_test.go index b9d46f1c2..96dec68c5 100644 --- a/core/config/model_config_test.go +++ b/core/config/model_config_test.go @@ -52,6 +52,25 @@ parameters: Expect(valid).To(BeTrue()) }) + It("parses upscale backend options and explicit usecase from YAML", func() { + raw := []byte(` +name: realesrgan +backend: stablediffusion-ggml +known_usecases: [upscale] +options: + - upscale_scale:4 + - upscale_tile_size:128 +parameters: + model: realesrgan-x4plus.gguf +`) + var cfg ModelConfig + Expect(yaml.Unmarshal(raw, &cfg)).To(Succeed()) + Expect(cfg.Options).To(Equal([]string{"upscale_scale:4", "upscale_tile_size:128"})) + Expect(cfg.KnownUsecases).NotTo(BeNil()) + Expect(*cfg.KnownUsecases & FLAG_UPSCALE).To(Equal(FLAG_UPSCALE)) + Expect(cfg.KnownUsecaseStrings).To(ContainElement("FLAG_UPSCALE")) + }) + It("round-trips context compression settings", func() { raw := []byte(` name: compressed-chat diff --git a/core/config/upscale_options.go b/core/config/upscale_options.go new file mode 100644 index 000000000..71ce44c55 --- /dev/null +++ b/core/config/upscale_options.go @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: MIT +package config + +import ( + "strconv" + "strings" +) + +// UpscaleScaleFromOptions reports the declared native scale for client metadata. +// Match the backend's int32 range and reject duplicates rather than advertising +// a value that the backend will reject. This does not replace backend validation. +func UpscaleScaleFromOptions(c ModelConfig) int { + var scale int + seen := false + for _, option := range c.Options { + key, value, _ := strings.Cut(option, ":") + if key != "upscale_scale" { + continue + } + if seen { + return 0 + } + seen = true + n, err := strconv.ParseInt(value, 10, 32) + if err != nil || n <= 0 { + return 0 + } + scale = int(n) + } + return scale +} diff --git a/core/config/upscale_options_test.go b/core/config/upscale_options_test.go new file mode 100644 index 000000000..5745b7203 --- /dev/null +++ b/core/config/upscale_options_test.go @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: MIT +package config + +import "testing" + +func TestUpscaleScaleFromOptions(t *testing.T) { + for _, tc := range []struct { + name string + options []string + want int + }{ + {"absent", nil, 0}, + {"unrelated", []string{"backend:CPU", "upscale_tile_size:128"}, 0}, + {"valid", []string{"upscale_scale:4"}, 4}, + {"malformed", []string{"upscale_scale:banana"}, 0}, + {"missing value", []string{"upscale_scale"}, 0}, + {"empty", []string{"upscale_scale:"}, 0}, + {"zero", []string{"upscale_scale:0"}, 0}, + {"negative", []string{"upscale_scale:-1"}, 0}, + {"overflow", []string{"upscale_scale:2147483648"}, 0}, + {"huge", []string{"upscale_scale:999999999999999999999999999"}, 0}, + {"duplicate", []string{"upscale_scale:4", "upscale_scale:4"}, 0}, + {"invalid duplicate", []string{"upscale_scale:4", "upscale_scale:bad"}, 0}, + {"extra separator", []string{"upscale_scale:4:2"}, 0}, + {"whitespace", []string{"upscale_scale: 4"}, 0}, + } { + t.Run(tc.name, func(t *testing.T) { + cfg := ModelConfig{Options: tc.options} + if got := UpscaleScaleFromOptions(cfg); got != tc.want { + t.Fatalf("got %d, want %d", got, tc.want) + } + }) + } +} diff --git a/core/gallery/esrgan_entries_test.go b/core/gallery/esrgan_entries_test.go new file mode 100644 index 000000000..b89740724 --- /dev/null +++ b/core/gallery/esrgan_entries_test.go @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: MIT +package gallery_test + +import ( + "github.com/mudler/LocalAI/core/config" + "os" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" +) + +var _ = Describe("gallery/index.yaml ESRGAN entry", func() { + It("installs the official native x4 checkpoint as an upscale-only model", func() { + entries, err := loadGalleryIndex() + Expect(err).NotTo(HaveOccurred()) + entry, ok := indexEntriesByName(entries)["realesrgan-x4plus-anime-6b"] + Expect(ok).To(BeTrue()) + Expect(entry.URL).To(Equal("github:mudler/LocalAI/gallery/virtual.yaml@master")) + Expect(entry.License).To(Equal("BSD-3-Clause")) + Expect(entry.Overrides["backend"]).To(Equal("stablediffusion-ggml")) + Expect(entry.Overrides["known_usecases"]).To(ConsistOf("upscale")) + Expect(entry.Overrides).NotTo(HaveKey("upscale_scale")) + Expect(entry.Overrides).NotTo(HaveKey("upscale_tile_size")) + Expect(entry.Overrides["options"]).To(ConsistOf("upscale_scale:4", "upscale_tile_size:128")) + Expect(entry.AdditionalFiles).To(HaveLen(1)) + file := entry.AdditionalFiles[0] + Expect(file.Filename).To(Equal("RealESRGAN_x4plus_anime_6B.pth")) + Expect(entry.Overrides["parameters"]).To(HaveKeyWithValue("model", file.Filename)) + Expect(file.URI).To(Equal("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/" + file.Filename)) + Expect(file.SHA256).To(Equal("f872d837d3c90ed2e05227bed711af5671a6fd1c9f7d7e91c911a61f155e99da")) + + // Keep the copy-paste configuration in the documentation aligned with the + // installable entry, including YAML scalar types. + doc, err := os.ReadFile("../../docs/content/features/image-generation.md") + Expect(err).NotTo(HaveOccurred()) + const start = "```yaml\nname: realesrgan-x4plus-anime-6b\n" + _, example, found := strings.Cut(string(doc), start) + Expect(found).To(BeTrue()) + example, _, found = strings.Cut(example, "```") + Expect(found).To(BeTrue()) + var settings map[string]any + Expect(yaml.Unmarshal([]byte(example), &settings)).To(Succeed()) + Expect(settings).To(Equal(entry.Overrides)) + }) +}) + +var _ = Describe("Diffusers upscaler gallery capability", func() { + It("explicitly declares upscaling rather than image generation", func() { + entries, err := loadGalleryIndex() + Expect(err).NotTo(HaveOccurred()) + entry, ok := indexEntriesByName(entries)["stable-diffusion-x4-upscaler"] + Expect(ok).To(BeTrue()) + data, err := yaml.Marshal(entry.Overrides) + Expect(err).NotTo(HaveOccurred()) + var cfg config.ModelConfig + Expect(yaml.Unmarshal(data, &cfg)).To(Succeed()) + Expect(cfg.HasUsecases(config.FLAG_UPSCALE)).To(BeTrue()) + Expect(cfg.HasUsecases(config.FLAG_IMAGE)).To(BeFalse()) + Expect(config.GetBackendCapability(cfg.Backend).GRPCMethods).To(ContainElement(config.MethodUpscaleImage)) + }) +}) diff --git a/core/http/endpoints/openai/list_capabilities.go b/core/http/endpoints/openai/list_capabilities.go index b7694a393..745663804 100644 --- a/core/http/endpoints/openai/list_capabilities.go +++ b/core/http/endpoints/openai/list_capabilities.go @@ -44,6 +44,7 @@ func ListModelCapabilitiesEndpoint(bcl *config.ModelConfigLoader, ml *model.Mode cfg.ContextSize = &appConfig.ContextSize } entry.Capabilities = cfg.Capabilities() + entry.UpscaleScale = config.UpscaleScaleFromOptions(cfg) // Generation aliases inherit target capabilities, but the router's // native classifier loads the named config directly (no alias resolution). original, exists := bcl.GetModelConfig(m) diff --git a/core/http/endpoints/openai/list_capabilities_test.go b/core/http/endpoints/openai/list_capabilities_test.go index 47fe7c94b..eae2d5200 100644 --- a/core/http/endpoints/openai/list_capabilities_test.go +++ b/core/http/endpoints/openai/list_capabilities_test.go @@ -98,6 +98,31 @@ var _ = Describe("ListModelCapabilitiesEndpoint", func() { } }) + It("exposes lowercase upscale and backend option scale without text modalities", func() { + writeConfig("upscaler", "name: upscaler\nbackend: stablediffusion-ggml\nknown_usecases: [upscale]\noptions:\n - upscale_scale:4\n") + entry := entryFor(call(), "upscaler") + Expect(entry).NotTo(BeNil()) + Expect(entry.Capabilities).To(Equal([]string{"upscale"})) + Expect(entry.InputModalities).To(Equal([]string{"image"})) + Expect(entry.OutputModalities).To(Equal([]string{"image"})) + Expect(entry.UpscaleScale).To(Equal(4)) + raw, err := json.Marshal(entry) + Expect(err).NotTo(HaveOccurred()) + Expect(string(raw)).To(ContainSubstring(`"upscale_scale":4`)) + }) + + It("omits absent or invalid backend scale metadata", func() { + for _, options := range []string{"", "options: [upscale_scale:bad]\n", "options: [upscale_scale:0]\n", "options: [upscale_scale:-1]\n", "options: [upscale_scale:2147483648]\n", "options: [upscale_scale:4, upscale_scale:4]\n"} { + writeConfig("upscaler", "name: upscaler\nbackend: stablediffusion-ggml\nknown_usecases: [upscale]\n"+options) + entry := entryFor(call(), "upscaler") + Expect(entry).NotTo(BeNil()) + Expect(entry.UpscaleScale).To(BeZero()) + raw, err := json.Marshal(entry) + Expect(err).NotTo(HaveOccurred()) + Expect(string(raw)).NotTo(ContainSubstring(`"upscale_scale"`)) + } + }) + It("returns the list envelope even with no models", func() { resp := call() Expect(resp.Object).To(Equal("list")) diff --git a/core/http/endpoints/openai/upscale.go b/core/http/endpoints/openai/upscale.go index eab000f8e..f63aa33d2 100644 --- a/core/http/endpoints/openai/upscale.go +++ b/core/http/endpoints/openai/upscale.go @@ -21,6 +21,10 @@ import ( model "github.com/mudler/LocalAI/pkg/model" ) +// maxUpscaleScale limits output area amplification to 256x and fits the int32 +// protobuf field. The backend still enforces the model's exact native scale. +const maxUpscaleScale = 16 + // UpscaleEndpoint handles POST /v1/images/upscale // // @Summary Image upscaling @@ -30,7 +34,7 @@ import ( // @Produce application/json // @Param model formData string true "Upscaler model identifier (e.g. stable-diffusion-x4-upscaler)" // @Param image formData file true "Input image file" -// @Param scale formData int false "Upscale factor: 2 or 4 (default 2)" +// @Param scale formData int true "Upscale factor: integer from 1 to 16; must match the model native scale" // @Success 200 {object} schema.OpenAIResponse // @Failure 400 {object} map[string]string // @Failure 500 {object} map[string]string @@ -45,13 +49,12 @@ func UpscaleEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appCon return echo.NewHTTPError(http.StatusBadRequest, "missing model") } - scale := 2 - if scaleStr != "" { - v, err := strconv.Atoi(scaleStr) - if err != nil || (v != 2 && v != 4) { - return echo.NewHTTPError(http.StatusBadRequest, "scale must be 2 or 4") - } - scale = v + if scaleStr == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing scale") + } + scale, err := strconv.Atoi(scaleStr) + if err != nil || scale < 1 || scale > maxUpscaleScale { + return echo.NewHTTPError(http.StatusBadRequest, "scale must be an integer from 1 to 16") } // Read uploaded image @@ -103,19 +106,26 @@ func UpscaleEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appCon id := uuid.New().String() dstPath := filepath.Join(tmpDir, fmt.Sprintf("upscale_%s.png", id)) + // Only this endpoint owns the generated destination; the generic backend + // must not delete arbitrary caller-supplied paths on failure. + completed := false + defer func() { + if !completed { + _ = os.Remove(dstPath) + } + }() + fn, err := backend.ImageUpscaleFunc(c.Request().Context(), srcPath, dstPath, scale, ml, *cfg, appConfig) if err != nil { return err } if err := fn(); err != nil { - _ = os.Remove(dstPath) return err } baseURL := middleware.BaseURL(c) imgURL, err := url.JoinPath(baseURL, "generated-images", filepath.Base(dstPath)) if err != nil { - _ = os.Remove(dstPath) return err } @@ -129,6 +139,10 @@ func UpscaleEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appCon }, } - return c.JSON(http.StatusOK, resp) + if err := c.JSON(http.StatusOK, resp); err != nil { + return err + } + completed = true + return nil } } diff --git a/core/http/endpoints/openai/upscale_test.go b/core/http/endpoints/openai/upscale_test.go index 2e13c8fc0..37fd9a6e4 100644 --- a/core/http/endpoints/openai/upscale_test.go +++ b/core/http/endpoints/openai/upscale_test.go @@ -1,13 +1,18 @@ package openai import ( - "bytes" "context" "encoding/json" + "errors" + grpcPkg "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/system" + ggrpc "google.golang.org/grpc" "net/http" "net/http/httptest" "os" "path/filepath" + "strconv" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/backend" @@ -29,7 +34,7 @@ var _ = Describe("Image upscaling", func() { var err error tmpDir, err = os.MkdirTemp("", "upscale") Expect(err).ToNot(HaveOccurred()) - appConfig = config.NewApplicationConfig(config.WithGeneratedContentDir(tmpDir)) + appConfig = config.NewApplicationConfig(config.WithGeneratedContentDir(tmpDir), config.WithSystemState(&system.SystemState{})) }) AfterEach(func() { @@ -39,7 +44,7 @@ var _ = Describe("Image upscaling", func() { It("stores the result in the directory served by /generated-images", func() { original := backend.ImageUpscaleFunc backend.ImageUpscaleFunc = func(_ context.Context, _, dst string, scale int, _ *model.ModelLoader, _ config.ModelConfig, _ *config.ApplicationConfig) (func() error, error) { - Expect(scale).To(Equal(4)) + Expect(scale).To(Equal(3)) return func() error { return os.WriteFile(dst, []byte("PNGDATA"), 0o644) }, nil @@ -47,7 +52,7 @@ var _ = Describe("Image upscaling", func() { DeferCleanup(func() { backend.ImageUpscaleFunc = original }) req, _ := makeMultipartRequest( - map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "4"}, + map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "3"}, map[string][]byte{"image": []byte("IMAGEDATA")}, ) rec := httptest.NewRecorder() @@ -66,24 +71,92 @@ var _ = Describe("Image upscaling", func() { contents, err := os.ReadFile(filepath.Join(tmpDir, "images", filename)) Expect(err).ToNot(HaveOccurred()) Expect(contents).To(Equal([]byte("PNGDATA"))) + entries, err := os.ReadDir(filepath.Join(tmpDir, "images")) + Expect(err).NotTo(HaveOccurred()) + Expect(entries).To(HaveLen(1)) }) - It("rejects unsupported scale factors", func() { - req, _ := makeMultipartRequest( - map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "3"}, - map[string][]byte{"image": []byte("IMAGEDATA")}, - ) + DescribeTable("rejects missing or invalid scales", func(scale string) { + fields := map[string]string{"model": "upscaler"} + if scale != "missing" { + fields["scale"] = scale + } + req, _ := makeMultipartRequest(fields, map[string][]byte{"image": []byte("IMAGE")}) rec := httptest.NewRecorder() - ctx := echo.New().NewContext(req, rec) - ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Backend: "diffusers"}) - + e := echo.New() + ctx := e.NewContext(req, rec) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{}) err := UpscaleEndpoint(nil, nil, appConfig)(ctx) - var httpErr *echo.HTTPError - Expect(err).To(MatchError(ContainSubstring("scale must be 2 or 4"))) - Expect(err).To(BeAssignableToTypeOf(httpErr)) - httpErr = err.(*echo.HTTPError) - Expect(httpErr.Code).To(Equal(http.StatusBadRequest)) - Expect(httpErr.Message).To(Equal("scale must be 2 or 4")) - Expect(bytes.TrimSpace(rec.Body.Bytes())).To(BeEmpty()) - }) + Expect(err).To(HaveOccurred()) + e.HTTPErrorHandler(err, ctx) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + }, + Entry("missing", "missing"), Entry("empty", ""), Entry("text", "abc"), + Entry("fraction", "2.5"), Entry("zero", "0"), Entry("negative", "-1"), + // 16 caps area amplification at 256x; int32 overflow must also fail. + Entry("above resource bound", "17"), Entry("int32 overflow", "2147483648"), + Entry("integer overflow", "999999999999999999999999")) + + DescribeTable("cleans only request files on failure", func(outcome string) { + sentinel := filepath.Join(tmpDir, "unrelated.png") + Expect(os.WriteFile(sentinel, []byte("keep"), 0600)).To(Succeed()) + var srcPath, dstPath string + original := backend.ImageUpscaleFunc + DeferCleanup(func() { backend.ImageUpscaleFunc = original }) + backend.ImageUpscaleFunc = func(ctx context.Context, src, dst string, scale int, _ *model.ModelLoader, cfg config.ModelConfig, app *config.ApplicationConfig) (func() error, error) { + srcPath, dstPath = src, dst + Expect(os.WriteFile(dst, []byte("partial"), 0600)).To(Succeed()) + if outcome == "load" { + return nil, errors.New("load failed") + } + fake := &upscaleTestBackend{outcome: outcome} + loader := model.NewModelLoader(&system.SystemState{}) + loader.SetModelRouter(func(_ context.Context, id string, _, _, _, _ string, _ *pb.ModelOptions, _ bool) (*model.Model, error) { + return model.NewModelWithClient(id, "test://upscale", fake), nil + }) + return backend.ImageUpscale(ctx, src, dst, scale, loader, cfg, app) + } + req, _ := makeMultipartRequest(map[string]string{"model": "upscaler", "scale": strconv.Itoa(16)}, map[string][]byte{"image": []byte("IMAGE")}) + rec := httptest.NewRecorder() + e := echo.New() + ctx := e.NewContext(req, rec) + cfg := &config.ModelConfig{Name: "upscaler", Backend: "stub"} + cfg.SetDefaults() + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, cfg) + err := UpscaleEndpoint(nil, nil, appConfig)(ctx) + Expect(err).To(HaveOccurred()) + if outcome == "unsuccessful" { + Expect(err.Error()).To(ContainSubstring("native scale mismatch")) + } + e.HTTPErrorHandler(err, ctx) + Expect(rec.Code).To(Equal(http.StatusInternalServerError)) + Expect(srcPath).NotTo(BeEmpty()) + Expect(dstPath).NotTo(BeEmpty()) + _, err = os.Stat(srcPath) + Expect(os.IsNotExist(err)).To(BeTrue()) + _, err = os.Stat(dstPath) + Expect(os.IsNotExist(err)).To(BeTrue()) + contents, err := os.ReadFile(sentinel) + Expect(err).NotTo(HaveOccurred()) + Expect(string(contents)).To(Equal("keep")) + }, Entry("model loading", "load"), Entry("transport", "transport"), Entry("nil result", "nil"), Entry("unsuccessful result", "unsuccessful")) }) + +type upscaleTestBackend struct { + grpcPkg.Backend + outcome string +} + +func (*upscaleTestBackend) HealthCheck(context.Context) (bool, error) { return true, nil } +func (*upscaleTestBackend) IsBusy() bool { return false } +func (b *upscaleTestBackend) UpscaleImage(_ context.Context, in *pb.UpscaleImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + Expect(in.Scale).To(Equal(int32(16))) + switch b.outcome { + case "transport": + return nil, errors.New("transport failed") + case "nil": + return nil, nil + default: + return &pb.Result{Success: false, Message: "native scale mismatch"}, nil + } +} diff --git a/core/http/react-ui/e2e/studio-fixtures.js b/core/http/react-ui/e2e/studio-fixtures.js index 15f33b7b4..cf7554ae7 100644 --- a/core/http/react-ui/e2e/studio-fixtures.js +++ b/core/http/react-ui/e2e/studio-fixtures.js @@ -54,11 +54,13 @@ export const GPU_RESOURCES = { } // Stub everything the front page reads. `types` are the types that have a -// model. `installs` collects the names POSTed to the install endpoint; after an -// install the type gains its model on the next capabilities read when -// `installOnPost` is set. +// model; `capabilityModels` supplies complete capability records when a test +// needs fields beyond a model type. `installs` collects the names POSTed to the +// install endpoint; after an install the type gains its model on the next +// capabilities read when `installOnPost` is set. export async function mockStudio(page, { types = ALL_TYPES, + capabilityModels, resources = GPU_RESOURCES, operations = [], installOnPost = false, @@ -69,7 +71,7 @@ export async function mockStudio(page, { await page.route('**/api/models/capabilities', route => { state.capabilityCalls += 1 if (capabilitiesStatus !== 200) return route.fulfill({ status: capabilitiesStatus, json: { error: 'down' } }) - return route.fulfill({ json: capabilities(state.types) }) + return route.fulfill({ json: { data: capabilityModels || capabilities(state.types).data } }) }) // The v1 list is the fallback useModels takes when the capabilities call // fails; make it fail too so an error state is reachable. diff --git a/core/http/react-ui/e2e/studio-lineage.spec.js b/core/http/react-ui/e2e/studio-lineage.spec.js index b64f8dcc5..dc9b18d6d 100644 --- a/core/http/react-ui/e2e/studio-lineage.spec.js +++ b/core/http/react-ui/e2e/studio-lineage.spec.js @@ -243,3 +243,81 @@ test.describe('Lineage view: phone', () => { } }) }) + +test.describe('Inline upscale', () => { + async function setup(page, { models = true, failure } = {}) { + await mockStudio(page, { capabilityModels: models ? [ + { id: 'up-4', capabilities: ['FLAG_UPSCALE'], upscaleScale: 4 }, + { id: 'up-2', capabilities: ['FLAG_UPSCALE'], upscaleScale: 2 }, + { id: 'invalid', capabilities: ['FLAG_UPSCALE'], upscaleScale: 0 }, + { id: 'text-image', capabilities: ['FLAG_IMAGE'] }, + ] : [] }) + await stubMedia(page) + const history = sampleHistory() + if (failure === 'cross-origin') history.image.find(e => e.id === 'img-bowls').results[0].url = 'https://other.test/source.png' + await seedHistory(page, history) + if (['fetch', 'non-image', 'source-decode'].includes(failure)) { + await page.route('**/generated-images/bowls.png', route => route.fulfill({ + status: failure === 'fetch' ? 500 : 200, + contentType: failure === 'non-image' ? 'text/html' : 'image/png', body: 'bad image', + })) + } + let requests = 0 + await page.route('**/v1/images/upscale', async route => { + requests++ + const request = route.request() + expect(request.headers()['content-type']).toContain('multipart/form-data; boundary=') + const body = request.postDataBuffer().toString() + expect(body).toContain('name="model"\r\n\r\nup-2') + expect(body).toContain('name="scale"\r\n\r\n2') + expect(body).toContain('name="image"; filename="bowls.png"') + await route.fulfill({ status: failure === 'api' ? 500 : 200, json: failure === 'api' + ? { error: 'upscale unavailable' } + : { data: [{ url: '/generated-images/upscaled.png', width: 99999, height: 99999 }] } }) + }) + if (failure === 'result-decode') await page.route('**/generated-images/upscaled.png', route => route.fulfill({ contentType: 'image/png', body: 'broken' })) + await page.goto('/app/studio?work=img-bowls') + await page.getByTestId('studio-dock-toggle').click() + return () => requests + } + + test('absent without an eligible model', async ({ page }) => { + await setup(page, { models: false }) + await expect(page.getByTestId('step-upscale')).toHaveCount(0) + }) + + test('filtered models, fixed scale, multipart and a persisted image child with repeat suggestions', async ({ page }) => { + await setup(page) + await page.getByTestId('step-upscale').click() + const composer = page.getByTestId('upscale-composer') + await expect(composer.locator('img')).toBeVisible() + await expect(composer.locator('option')).toHaveText(['up-4', 'up-2']) + await composer.getByLabel('Upscaling model').selectOption('up-2') + await expect(composer.getByLabel('Scale', { exact: true })).toHaveValue('2×') + await expect(composer.getByLabel('Scale', { exact: true })).toHaveAttribute('readonly', '') + await composer.getByRole('button', { name: 'Upscale', exact: true }).click() + await expect(page.getByTestId('lineage-node')).toHaveCount(2) + for (const edge of ['animate', 'to-3d', 'variation', 'upscale']) await expect(page.getByTestId(`step-${edge}`)).toBeVisible() + await expect.poll(() => page.evaluate(() => JSON.parse(localStorage.getItem('localai_image_history')).find(e => e.edge === 'upscale'))).toMatchObject({ + parentId: 'img-bowls', model: 'up-2', params: { scale: 2, sourceDimensions: { width: 640, height: 420 }, outputDimensions: { width: 640, height: 420 } }, + }) + await expect(page.locator('.studio-edge-label')).toContainText(['upscale']) + await page.getByTestId('step-upscale').click() + await expect(page.getByTestId('upscale-composer').locator('img')).toHaveAttribute('src', '/generated-images/upscaled.png') + }) + + for (const failure of ['cross-origin', 'non-image', 'fetch', 'api', 'source-decode', 'result-decode']) { + test(`${failure} failure preserves selection and creates no child`, async ({ page }) => { + const requests = await setup(page, { failure }) + await page.getByTestId('step-upscale').click() + const composer = page.getByTestId('upscale-composer') + await composer.getByLabel('Upscaling model').selectOption('up-2') + await composer.getByRole('button', { name: 'Upscale', exact: true }).click() + await expect(composer.getByRole('alert')).toContainText('Upscale failed') + await expect(page.getByTestId('lineage-node')).toHaveCount(1) + await expect(node(page, 'img-bowls')).toHaveAttribute('data-selected', 'true') + expect(requests()).toBe(['api', 'result-decode'].includes(failure) ? 1 : 0) + expect(await page.evaluate(() => JSON.parse(localStorage.getItem('localai_image_history')).some(e => e.edge === 'upscale'))).toBe(false) + }) + } +}) diff --git a/core/http/react-ui/public/locales/en/media.json b/core/http/react-ui/public/locales/en/media.json index 614d8268d..8e2eae66c 100644 --- a/core/http/react-ui/public/locales/en/media.json +++ b/core/http/react-ui/public/locales/en/media.json @@ -206,7 +206,8 @@ "transform": "transform", "diarize": "who spoke", "soundtrack": "soundtrack", - "turntable": "turntable" + "turntable": "turntable", + "upscale": "upscale" }, "steps": { "animate": { @@ -240,6 +241,10 @@ "turntable": { "title": "Turntable clip", "why": "Show the object from every side" + }, + "upscale": { + "title": "Upscale", + "why": "Increase the resolution of this image." } }, "kinds": { @@ -364,6 +369,14 @@ "notKept": "The recording and the transcript are not kept, so only the summary of this run is shown. Run it again with the file to see the timeline.", "empty": "Choose a recording to see who spoke and when" } + }, + "upscale": { + "model": "Upscaling model", + "scale": "Scale", + "source": "Source image", + "loading": "Upscaling…", + "failed": "Upscale failed: {{message}}", + "noResult": "No output image was returned." } }, "image": { diff --git a/core/http/react-ui/src/components/studio/LineageView.jsx b/core/http/react-ui/src/components/studio/LineageView.jsx index 95cb2a829..77bbbb90d 100644 --- a/core/http/react-ui/src/components/studio/LineageView.jsx +++ b/core/http/react-ui/src/components/studio/LineageView.jsx @@ -1,5 +1,7 @@ import { useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react' import { useTranslation } from 'react-i18next' +// eslint-disable-next-line no-unused-vars +import UpscaleComposer from './UpscaleComposer' import Icon from '../Icon' // eslint-disable-next-line no-unused-vars import InstallNote from './InstallNote' @@ -31,10 +33,10 @@ function isTyping(el) { // are installed. // // "Run as a new take" and "Branch from here" open the workspace with the source -// and the words filled in. They do not run anything here. A step the destination +// and the words filled in. Upscale instead runs in the inline composer. A step the destination // workspace cannot start from yet is listed but disabled, with the reason. export default function LineageView({ - items, workId, modalities, defaultModel, onClose, onToggleFavourite, onHandoff, onModelsChanged, + items, workId, modalities, upscalers = [], onAddImage, defaultModel, onClose, onToggleFavourite, onHandoff, onModelsChanged, }) { const { t } = useTranslation('media') const members = useMemo(() => projectOf(items, workId), [items, workId]) @@ -137,7 +139,7 @@ export default function LineageView({ } const root = members[0] - const steps = NEXT_STEPS[selected.type] || [] + const steps = [...(NEXT_STEPS[selected.type] || []), ...(selected.type === 'images' && upscalers.length ? [{ to: 'images', edge: 'upscale', supported: true }] : [])] const takeOk = canTake(selected) const branchStep = suggestion || steps.find(s => s.supported) || null const draftInstalled = draft ? installed.has(draft.to) : true @@ -189,7 +191,7 @@ export default function LineageView({ ) } if (n.kind === 'ghost') { - const missing = !installed.has(n.ghost.to) + const missing = n.ghost.edge !== 'upscale' && !installed.has(n.ghost.to) return ( - + } ) : ( <> @@ -284,7 +286,13 @@ export default function LineageView({

)} - {expanded && (draft ? ( + {expanded && (draft?.edge === 'upscale' ? ( + { + const id = onAddImage(entry) + setSelectedId(id) + setDraft(null) + }} /> + ) : draft ? (
{t('studio.lineage.makeFrom')}
@@ -296,7 +304,7 @@ export default function LineageView({ type="button" className="dk-chip studio-type" aria-pressed={draft.edge === step.edge} - data-missing={(!installed.has(step.to)) || undefined} + data-missing={(step.edge !== 'upscale' && !installed.has(step.to)) || undefined} disabled={bad} data-testid={`draft-target-${step.edge}`} onClick={() => startDraft(step)} @@ -330,7 +338,7 @@ export default function LineageView({ {steps.length === 0 &&
  • {t('studio.lineage.noBranch', { type: label(selected.type) })}
    • } {steps.map(step => { const bad = !step.supported - const missing = !bad && !installed.has(step.to) + const missing = !bad && step.edge !== 'upscale' && !installed.has(step.to) return (
  • +
    • + + ) +} diff --git a/core/http/react-ui/src/hooks/useStudioWork.js b/core/http/react-ui/src/hooks/useStudioWork.js index 90b86b894..4ecb5abc8 100644 --- a/core/http/react-ui/src/hooks/useStudioWork.js +++ b/core/http/react-ui/src/hooks/useStudioWork.js @@ -1,5 +1,5 @@ import { useCallback, useEffect, useMemo, useState } from 'react' -import { clearAllMediaHistory, FAVOURITES_KEY, readAllMediaHistory } from './useMediaHistory' +import { clearAllMediaHistory, FAVOURITES_KEY, readAllMediaHistory, useMediaHistory } from './useMediaHistory' import { use3DHistory } from './use3DHistory' import { collectWork, pruneIds, toggleId } from '../utils/studioWork' @@ -21,12 +21,14 @@ function writeFavourites(ids) { // Everything the Studio front page shows as "your work": the entries each // workspace stored, in one list, with favourites. The browser-storage lists are -// read once on mount (they change only while a workspace is open, and this page -// mounts fresh when you come back); 3D arrives asynchronously from IndexedDB. +// read once on mount, except images: inline upscaling writes through the same +// image history hook and merges its live entries. 3D arrives from IndexedDB. // // `ready` is false until 3D has answered, so the page can show a skeleton // instead of flashing an empty state for work that is a moment away. export function useStudioWork() { + const { addEntry: addImageEntry, historyProps: imageHistory } = useMediaHistory('image') + const clearImages = imageHistory.onClearAll const [media, setMedia] = useState(() => readAllMediaHistory()) const [favourites, setFavourites] = useState(readFavourites) const { entries: threeD, clearAll: clearThreeD } = use3DHistory() @@ -41,8 +43,8 @@ export function useStudioWork() { useEffect(() => { if (threeD.length > 0) setReady(true) }, [threeD]) const items = useMemo( - () => collectWork(media, threeD, { favourites }), - [media, threeD, favourites], + () => collectWork({ ...media, image: imageHistory.entries }, threeD, { favourites }), + [media, imageHistory.entries, threeD, favourites], ) const toggleFavourite = useCallback((id) => { @@ -54,11 +56,12 @@ export function useStudioWork() { }, [items]) const clearHistory = useCallback(async () => { + clearImages() clearAllMediaHistory() setMedia(readAllMediaHistory()) setFavourites([]) await clearThreeD() - }, [clearThreeD]) + }, [clearThreeD, clearImages]) - return { items, ready, toggleFavourite, clearHistory } + return { items, ready, addImageEntry, toggleFavourite, clearHistory } } diff --git a/core/http/react-ui/src/pages/Studio.jsx b/core/http/react-ui/src/pages/Studio.jsx index 7f49d6920..4ed1165df 100644 --- a/core/http/react-ui/src/pages/Studio.jsx +++ b/core/http/react-ui/src/pages/Studio.jsx @@ -14,6 +14,7 @@ import { useModels } from '../hooks/useModels' import { CAP_DIARIZATION, CAP_IMAGE, CAP_VIDEO, CAP_3D, CAP_3D_ANIMATION, CAP_TTS, CAP_SOUND_GENERATION, CAP_AUDIO_TRANSFORM, } from '../utils/capabilities' +import { upscaleModels } from '../utils/upscale' import Icon from '../components/Icon' import '../components/studio/studio.css' @@ -132,6 +133,7 @@ export default function Studio() { ) : ( { + const form = new FormData() + form.append('model', model) + form.append('image', image) + form.append('scale', String(scale)) + return handleResponse(await fetch(apiUrl('/v1/images/upscale'), { method: 'POST', body: form })) + }, generate: (body) => postJSON(API_CONFIG.endpoints.imageGenerations, body), } diff --git a/core/http/react-ui/src/utils/capabilities.js b/core/http/react-ui/src/utils/capabilities.js index 722c85842..cd0681e2b 100644 --- a/core/http/react-ui/src/utils/capabilities.js +++ b/core/http/react-ui/src/utils/capabilities.js @@ -9,6 +9,7 @@ export const CAP_EDIT = 'FLAG_EDIT' export const CAP_EMBEDDINGS = 'FLAG_EMBEDDINGS' export const CAP_RERANK = 'FLAG_RERANK' export const CAP_IMAGE = 'FLAG_IMAGE' +export const CAP_UPSCALE = 'FLAG_UPSCALE' export const CAP_TRANSCRIPT = 'FLAG_TRANSCRIPT' export const CAP_TTS = 'FLAG_TTS' export const CAP_SOUND_GENERATION = 'FLAG_SOUND_GENERATION' diff --git a/core/http/react-ui/src/utils/studioWork.js b/core/http/react-ui/src/utils/studioWork.js index 2d05ce2a4..d6b9fa178 100644 --- a/core/http/react-ui/src/utils/studioWork.js +++ b/core/http/react-ui/src/utils/studioWork.js @@ -39,7 +39,7 @@ const TYPE_BY_MEDIA = Object.fromEntries( Object.entries(TYPE_INFO).filter(([, v]) => v.media).map(([key, v]) => [v.media, key]), ) -export const EDGE_KINDS = ['take', 'animate', 'to-3d', 'variation', 'transform', 'diarize'] +export const EDGE_KINDS = ['upscale', 'take', 'animate', 'to-3d', 'variation', 'transform', 'diarize'] // What each kind of result can be sent on to. `supported` is whether the // destination workspace accepts that result as a starting point today. The @@ -121,7 +121,8 @@ export function parseSize(size) { // Width over height for a tile. Only what the entry recorded: an entry with no // size is drawn square rather than guessed. export function aspectOf(item) { - const size = parseSize(item?.params?.size) + const output = item?.params?.outputDimensions + const size = output?.width > 0 && output?.height > 0 ? output : parseSize(item?.params?.size) return size ? size.width / size.height : 1 } diff --git a/core/http/react-ui/src/utils/upscale.js b/core/http/react-ui/src/utils/upscale.js new file mode 100644 index 000000000..0e64b0df5 --- /dev/null +++ b/core/http/react-ui/src/utils/upscale.js @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: MIT +import { CAP_UPSCALE } from './capabilities.js' + +export function upscaleModels(models) { + return models.filter(m => !m.disabled && m.capabilities?.includes(CAP_UPSCALE) && + Number.isFinite(m.upscaleScale) && m.upscaleScale > 0) + .map(({ id, upscaleScale }) => ({ id, upscaleScale })) +} + +export async function sourceImageFile(source, base = window.location.href, fetcher = fetch) { + if (source instanceof File) return source + const url = new URL(source, base) + if (url.origin !== new URL(base).origin || !['http:', 'https:'].includes(url.protocol)) { + throw new Error('Source image must be same-origin') + } + const response = await fetcher(url.href, { credentials: 'same-origin', mode: 'same-origin', redirect: 'error' }) + if (!response.ok) throw new Error(`Image fetch failed (HTTP ${response.status})`) + if (!response.headers.get('content-type')?.toLowerCase().startsWith('image/')) throw new Error('Source is not an image') + const blob = await response.blob() + return new File([blob], url.pathname.split('/').pop() || 'source.png', { type: blob.type }) +} + +export function imageDimensions(source) { + return new Promise((resolve, reject) => { + const objectUrl = source instanceof Blob ? URL.createObjectURL(source) : null + const image = new Image() + const finish = (error) => { + clearTimeout(timer) + image.onload = image.onerror = null + if (objectUrl) URL.revokeObjectURL(objectUrl) + if (error || !image.naturalWidth || !image.naturalHeight) reject(new Error('Could not decode image dimensions')) + else resolve({ width: image.naturalWidth, height: image.naturalHeight }) + } + const timer = setTimeout(() => finish(true), 30000) + image.onload = () => finish(false) + image.onerror = () => finish(true) + image.src = objectUrl || source + }) +} diff --git a/core/http/react-ui/src/utils/upscale.test.js b/core/http/react-ui/src/utils/upscale.test.js new file mode 100644 index 000000000..7f3cbe627 --- /dev/null +++ b/core/http/react-ui/src/utils/upscale.test.js @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: MIT +import test from 'node:test' +import assert from 'node:assert/strict' +import { upscaleModels, sourceImageFile } from './upscale.js' +import { toWorkItem, EDGE_KINDS } from './studioWork.js' + +test('only enabled upscalers with positive finite scale are offered', () => { + assert.deepEqual(upscaleModels([ + { id: 'yes', capabilities: ['FLAG_UPSCALE'], upscaleScale: 4 }, + { id: 'image', capabilities: ['FLAG_IMAGE'], upscaleScale: 4 }, + ...[0, -1, undefined, Infinity].map(upscaleScale => ({ capabilities: ['FLAG_UPSCALE'], upscaleScale })), + { capabilities: ['FLAG_UPSCALE'], upscaleScale: 2, disabled: true }, + ]), [{ id: 'yes', upscaleScale: 4 }]) +}) + +test('URL source is same-origin, checked and named', async () => { + let called = false + const fetcher = async (url, options) => { + called = true + assert.equal(url, 'https://local.test/generated/source.png') + assert.equal(options.credentials, 'same-origin') + return new Response('png', { headers: { 'Content-Type': 'image/png' } }) + } + const file = await sourceImageFile('/generated/source.png', 'https://local.test/app/studio', fetcher) + assert.equal(file.name, 'source.png') + assert.equal(file.type, 'image/png') + called = false + await assert.rejects(sourceImageFile('https://other.test/a.png', 'https://local.test/', fetcher), /origin/) + assert.equal(called, false) + for (const response of [new Response('bad', { status: 500 }), new Response('html', { headers: { 'Content-Type': 'text/html' } })]) { + await assert.rejects(sourceImageFile('/a', 'https://local.test/', async () => response)) + } + await assert.rejects(sourceImageFile('/a', 'https://local.test/', async () => { throw new Error('offline') })) +}) + +test('upscale history normalizes as a child image', () => { + assert.ok(EDGE_KINDS.includes('upscale')) + const item = toWorkItem('images', { id: 'child', parentId: 'source', edge: 'upscale', results: [{ url: '/out.png' }] }) + assert.equal(item.type, 'images') + assert.equal(item.parentId, 'source') + assert.equal(item.edge, 'upscale') +}) diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index fc4831b04..a4066ae7f 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -1,6 +1,8 @@ package routes import ( + "net/http" + "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/config" @@ -267,10 +269,31 @@ func RegisterOpenAIRoutes(app *echo.Echo, app.POST("/v1/images/inpainting", inpaintingHandler, imageMiddleware...) app.POST("/images/inpainting", inpaintingHandler, imageMiddleware...) - // upscale endpoint - reuse same middleware config as images + // Upscaling has its own model usecase; generation models are not defaults. + upscaleMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, + traceMiddleware, + // Default: use the first available upscaling model + re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_UPSCALE)), + re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }), + func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + // The default-model filter does not validate explicitly named models. + cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || cfg == nil || !cfg.HasUsecases(config.FLAG_UPSCALE) { + return echo.NewHTTPError(http.StatusBadRequest, "model does not support image upscaling") + } + if err := re.SetOpenAIRequest(c); err != nil { + return err + } + return next(c) + } + }, + } + upscaleHandler := openai.UpscaleEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) - app.POST("/v1/images/upscale", upscaleHandler, imageMiddleware...) - app.POST("/images/upscale", upscaleHandler, imageMiddleware...) + app.POST("/v1/images/upscale", upscaleHandler, upscaleMiddleware...) + app.POST("/images/upscale", upscaleHandler, upscaleMiddleware...) // List models app.GET("/v1/models", openai.ListModelsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig(), application.AuthDB())) diff --git a/core/http/routes/ui_api.go b/core/http/routes/ui_api.go index 08ced3bc0..78ff873aa 100644 --- a/core/http/routes/ui_api.go +++ b/core/http/routes/ui_api.go @@ -822,6 +822,7 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model NodeStatus string `json:"node_status"` } type modelCapability struct { + UpscaleScale int `json:"upscaleScale,omitempty"` ThreeDOperations []schema.ThreeDOperation `json:"three_d_operations,omitempty"` ID string `json:"id"` Capabilities []string `json:"capabilities"` @@ -873,6 +874,7 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model result = append(result, modelCapability{ ID: cfg.Name, Capabilities: cfg.KnownUsecaseStrings, + UpscaleScale: config.UpscaleScaleFromOptions(cfg), ThreeDOperations: cfg.ThreeDOperations(), Backend: cfg.Backend, Disabled: cfg.IsDisabled(), diff --git a/core/http/routes/ui_api_capabilities_test.go b/core/http/routes/ui_api_capabilities_test.go new file mode 100644 index 000000000..5714000de --- /dev/null +++ b/core/http/routes/ui_api_capabilities_test.go @@ -0,0 +1,80 @@ +package routes_test + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/routes" + "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" +) + +func TestUIUpscaleCapabilities(t *testing.T) { + dir := t.TempDir() + st, err := system.GetSystemState(system.WithModelPath(dir)) + if err != nil { + t.Fatal(err) + } + cl := config.NewModelConfigLoader(dir) + for name, options := range map[string]string{ + "upscaler": "options: [upscale_scale:4]\n", + "absent": "", + "malformed": "options: [upscale_scale:bad]\n", + "nonpositive": "options: [upscale_scale:0]\n", + "negative": "options: [upscale_scale:-1]\n", + "overflow": "options: [upscale_scale:2147483648]\n", + "duplicate": "options: [upscale_scale:4, upscale_scale:4]\n", + } { + path := filepath.Join(dir, name+".yaml") + if err := os.WriteFile(path, []byte("name: "+name+"\nbackend: stablediffusion-ggml\nknown_usecases: [upscale]\n"+options), 0600); err != nil { + t.Fatal(err) + } + if err := cl.ReadModelConfig(path); err != nil { + t.Fatal(err) + } + } + ac := config.NewApplicationConfig() + ac.SystemState = st + ml := model.NewModelLoader(st) + gs := galleryop.NewGalleryService(ac, ml) + app := echo.New() + routes.RegisterUIAPIRoutes(app, cl, ml, ac, gs, galleryop.NewOpCache(gs), &application.Application{}, func(next echo.HandlerFunc) echo.HandlerFunc { return next }) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/models/capabilities", nil)) + if rec.Code != http.StatusOK { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + var response struct { + Data []struct { + ID string `json:"id"` + Capabilities []string `json:"capabilities"` + Scale *int `json:"upscaleScale"` + } `json:"data"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if len(response.Data) != 7 { + t.Fatalf("unexpected model count: %d", len(response.Data)) + } + for _, entry := range response.Data { + if len(entry.Capabilities) != 1 || entry.Capabilities[0] != "FLAG_UPSCALE" { + t.Fatalf("unexpected capabilities: %+v", entry) + } + if entry.ID == "upscaler" { + if entry.Scale == nil || *entry.Scale != 4 { + t.Fatalf("unexpected scale: %+v", entry) + } + } else if entry.Scale != nil { + t.Fatalf("invalid scale exposed: %+v", entry) + } + } +} diff --git a/core/http/routes/upscale_test.go b/core/http/routes/upscale_test.go new file mode 100644 index 000000000..fa9f2f741 --- /dev/null +++ b/core/http/routes/upscale_test.go @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: MIT +package routes_test + +import ( + "bytes" + "context" + "mime/multipart" + "net/http" + "net/http/httptest" + "os" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/http/routes" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Upscale route model selection", func() { + It("uses upscale-only models on both aliases without admitting generation models", func() { + root := GinkgoT().TempDir() + app, err := application.New(config.WithDataPath(root), config.WithGeneratedContentDir(root), config.WithDisableLocalAIAssistant(true), config.WithSystemState(&system.SystemState{Model: system.Model{ModelsPath: root}, Backend: system.Backend{BackendsPath: root}})) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(app.Shutdown()).To(Succeed()) }) + upscale, image := config.FLAG_UPSCALE, config.FLAG_IMAGE + configs := []config.ModelConfig{ + {Name: "esrgan", Backend: "stablediffusion-ggml", KnownUsecases: &upscale}, + {Name: "stable-diffusion-x4-upscaler", Backend: "diffusers", KnownUsecases: &upscale}, + {Name: "ordinary-image", Backend: "diffusers", KnownUsecases: &image}, + {Name: "incompatible", Backend: "whisper"}, + } + for i := range configs { + configs[i].SetDefaults() + configs[i].Model = configs[i].Name + } + original := backend.ImageUpscaleFunc + DeferCleanup(func() { backend.ImageUpscaleFunc = original }) + served := "" + backend.ImageUpscaleFunc = func(_ context.Context, _, dst string, _ int, _ *model.ModelLoader, cfg config.ModelConfig, _ *config.ApplicationConfig) (func() error, error) { + served = cfg.Name + return func() error { return os.WriteFile(dst, []byte("PNG"), 0600) }, nil + } + e := echo.New() + selected := "" + e.Use(func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + err := next(c) + if cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig); ok && cfg != nil { + selected = cfg.Name + } + return err + } + }) + routes.RegisterOpenAIRoutes(e, middleware.NewRequestExtractor(app.ModelConfigLoader(), app.ModelLoader(), app.ApplicationConfig()), app) + for _, path := range []string{"/v1/images/upscale", "/images/upscale"} { + for _, cfg := range configs { + app.ModelConfigLoader().ReplaceModelConfigs([]config.ModelConfig{cfg}) + for _, explicit := range []bool{true, false} { + served, selected = "", "" + var body bytes.Buffer + form := multipart.NewWriter(&body) + if explicit { + Expect(form.WriteField("model", cfg.Name)).To(Succeed()) + } + Expect(form.WriteField("scale", "4")).To(Succeed()) + file, err := form.CreateFormFile("image", "image.png") + Expect(err).NotTo(HaveOccurred()) + _, err = file.Write([]byte("IMAGE")) + Expect(err).NotTo(HaveOccurred()) + Expect(form.Close()).To(Succeed()) + req := httptest.NewRequest(http.MethodPost, path, &body) + req.Header.Set("Content-Type", form.FormDataContentType()) + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + eligible := cfg.HasUsecases(config.FLAG_UPSCALE) + if explicit && eligible { + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + Expect(served).To(Equal(cfg.Name)) + } else { + // The endpoint still requires a form model even when middleware selects a default. + Expect(rec.Code).To(Equal(http.StatusBadRequest), rec.Body.String()) + Expect(served).To(BeEmpty()) + } + if !explicit { + if eligible { + Expect(selected).To(Equal(cfg.Name)) + } else { + Expect(selected).NotTo(Equal(cfg.Name)) + } + } + } + } + } + }) +}) diff --git a/core/schema/openai.go b/core/schema/openai.go index ff8b3b5c6..62b245262 100644 --- a/core/schema/openai.go +++ b/core/schema/openai.go @@ -275,6 +275,7 @@ type ModelsDataResponse struct { // served by the LocalAI-specific /v1/models/capabilities endpoint so clients can // route attachments (image/audio/video) to a model only when it can handle them. type ModelCapabilities struct { + UpscaleScale int `json:"upscale_scale,omitempty"` ThreeDOperations []ThreeDOperation `json:"three_d_operations,omitempty"` ID string `json:"id"` Object string `json:"object"` diff --git a/docs/content/features/image-generation.md b/docs/content/features/image-generation.md index 065e3ce4c..03018d088 100644 --- a/docs/content/features/image-generation.md +++ b/docs/content/features/image-generation.md @@ -135,6 +135,83 @@ options: 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. +#### ESRGAN image upscaling + +The `stablediffusion-ggml` backend supports standalone ESRGAN upscaling through +`POST /v1/images/upscale`. This is a separate, upscale-only model, not a diffusion +pipeline. Image generation does **not** automatically upscale its results. + +Install the gallery model: + +```bash +local-ai models install realesrgan-x4plus-anime-6b +``` + +This uses the official [RealESRGAN_x4plus_anime_6B.pth release](https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.2.2.4), +a six-block model for anime images with a native **4x** scale. The PyTorch `.pth` +checkpoint is supported directly by [stable-diffusion.cpp](https://github.com/leejet/stable-diffusion.cpp/blob/f89d9b13d730eabeede7314ce49dacd18d3c90c2/docs/esrgan.md); +no GGUF conversion is needed. The upstream project is licensed under +[BSD-3-Clause](https://github.com/xinntao/Real-ESRGAN/blob/v0.2.2.4/LICENSE). + +For manual installation, download that checkpoint into your models directory and +save the following as `realesrgan-x4plus-anime-6b.yaml` alongside it (the gallery +installer configures these settings for you): + +```yaml +name: realesrgan-x4plus-anime-6b +backend: stablediffusion-ggml +parameters: + model: RealESRGAN_x4plus_anime_6B.pth +known_usecases: [upscale] +options: + - upscale_scale:4 + - upscale_tile_size:128 +``` + +`upscale_scale` and `upscale_tile_size` are backend options, expressed as +`key:value` strings in `options`, not root-level configuration fields. +`upscale_scale` must be a positive integer. Keep `known_usecases` explicitly +set to `[upscale]`: do not add `image` or `diffusion_model` options to this +configuration. Options are forwarded unchanged; the backend rejects invalid +or duplicate recognized settings. + +The multipart fields `model`, `image`, and **`scale` are required**. The HTTP +endpoint accepts only integer scales from **1 through 16**, rejecting missing, +malformed, zero, negative, or larger values with HTTP 400. This range is not a +promise that each model supports every scale: for ESRGAN, the configured +`upscale_scale`, the loaded model's detected native scale, and the request's +`scale` must all match exactly. Mismatches fail rather than silently resizing or +running multiple passes. For this model, always request `scale=4`. + +```bash +curl http://localhost:8080/v1/images/upscale \ + -F model=realesrgan-x4plus-anime-6b \ + -F scale=4 \ + -F image=@input.png +``` + +The response follows the image-generation response format, with the output URL +in `data[0].url` under `/generated-images`. A successful 4x upscale multiplies +both width and height by four. + +`upscale_tile_size` controls ESRGAN tiling in pixels: omit it or set it to +`0` for stable-diffusion.cpp full-frame processing; set a positive value (such +as the gallery entry's `128`) to enable tiling. Smaller tiles reduce working +memory at the cost of more tile work. Full-frame processing can use much more +memory. `upscale_direct` is a separate convolution-algorithm flag, not a tile +control. Tile settings apply to ESRGAN, not the diffusion model's separate VAE +tiling settings. + +In **Studio**, open an existing image, choose **Upscale**, select an installed +upscaler, and submit. The scale is read-only and comes from the selected model's +configured scale. A successful result becomes an **Upscale child node** of the +source image, leaving the original intact. Open that child to continue the +workflow with the available image actions, including another upscale or image +editing; each upscale is an explicit operation, not a post-generation default. + +Unlike the Diffusers upscaling described below, this path uses ESRGAN inference +and does not fall back to Lanczos resizing. + #### 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. @@ -166,8 +243,10 @@ By default the RPC devices join the pool and participate in placement; combine w #### Image upscaling -LocalAI can upscale an uploaded image by a factor of 2 or 4 through -`POST /v1/images/upscale`. Install the included Stable Diffusion x4 upscaler +The Diffusers backend can upscale an uploaded image through +`POST /v1/images/upscale`. The multipart `scale` field is required and must be an +integer from 1 through 16 at the HTTP boundary; actual scaling behavior depends +on the selected pipeline. Install the included Stable Diffusion x4 upscaler gallery model first: ```bash @@ -186,7 +265,12 @@ curl http://localhost:8080/v1/images/upscale \ The response uses the same format as image generation and returns the generated image under `/generated-images`. The `diffusers` backend uses a loaded `StableDiffusionUpscalePipeline` or `StableDiffusionLatentUpscalePipeline` when -configured. Other diffusers pipelines fall back to Lanczos resizing. +configured. Other diffusers pipelines fall back to Lanczos resizing. Upscale +routes select models with `known_usecases: [upscale]`, not ordinary image +generation models. If you installed the x4 upscaler before it declared this +usecase, update its model YAML to `known_usecases: [upscale]` or reinstall the +gallery entry. Custom Diffusers configurations intended for upscaling should +also explicitly declare the `upscale` usecase. #### Model setup diff --git a/gallery/index.yaml b/gallery/index.yaml index 400149a0d..944c96774 100644 --- a/gallery/index.yaml +++ b/gallery/index.yaml @@ -21331,6 +21331,36 @@ - filename: llama-cpp/models/deepseek-ai.DeepSeek-V3.2.Q4_K_M-00029-of-00029.gguf sha256: 013af4e9d2f84e484f77c7bae2a02652607f0f0179bd2815ffdf401c3ada5184 uri: https://huggingface.co/DevQuasar/deepseek-ai.DeepSeek-V3.2-GGUF/resolve/main/deepseek-ai.DeepSeek-V3.2.Q4_K_M-00029-of-00029.gguf +# Official PyTorch checkpoint recommended by stable-diffusion.cpp docs/esrgan.md. +# SHA256 measured from the release download (17,938,799 bytes). +- name: realesrgan-x4plus-anime-6b + url: github:mudler/LocalAI/gallery/virtual.yaml@master + urls: + - https://github.com/xinntao/Real-ESRGAN + - https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.2.2.4 + - https://github.com/xinntao/Real-ESRGAN/blob/v0.2.2.4/LICENSE + - https://github.com/leejet/stable-diffusion.cpp/blob/f89d9b13d730eabeede7314ce49dacd18d3c90c2/docs/esrgan.md + description: | + RealESRGAN x4plus anime 6B is a six-block ESRGAN model for anime images, + with a native 4x scale. Uses the official PyTorch .pth checkpoint through + stable-diffusion.cpp, without conversion. Upscale-only: not a text-to-image model. + license: BSD-3-Clause + tags: + - image-upscaling + - super-resolution + - stablediffusion-ggml + overrides: + backend: stablediffusion-ggml + known_usecases: [upscale] + options: + - upscale_scale:4 + - upscale_tile_size:128 + parameters: + model: RealESRGAN_x4plus_anime_6B.pth + files: + - filename: RealESRGAN_x4plus_anime_6B.pth + uri: https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth + sha256: f872d837d3c90ed2e05227bed711af5671a6fd1c9f7d7e91c911a61f155e99da - name: stable-diffusion-x4-upscaler url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: @@ -21350,7 +21380,7 @@ diffusers: pipeline_type: StableDiffusionUpscalePipeline known_usecases: - - image + - upscale parameters: model: stabilityai/stable-diffusion-x4-upscaler artifacts: