From 0177e4988ca23e43ede8a440c844060d60620439 Mon Sep 17 00:00:00 2001 From: mudler-agent Date: Sun, 11 Oct 2026 00:30:09 +0200 Subject: [PATCH] feat(images): add ESRGAN upscaling and Studio workflow (#12620) * feat(images): add upscale model capabilities * fix(config): keep SD GGML upscale opt-in * feat(stablediffusion): add lazy ESRGAN upscale bridge Keep native image ownership in C++ and serialize context creation, inference and teardown. Use upscale_scale in backend options as the explicit upscale-only marker because ModelOptions does not carry usecases. Assisted-by: Codex:gpt-5 * fix(stablediffusion): validate ESRGAN settings * fix(upscale): forward typed settings and contain resolved model paths Typed settings replace legacy options without duplicates. Resolve model and trusted root symlinks before caching the model path, rejecting missing files and escapes. Assisted-by: Codex:gpt-5 * fix(upscale): validate scale and clean failed HTTP outputs Require an explicit scale bounded to 1-16, check backend results, and keep destination cleanup owned by the allocating endpoint. Cover validation, result failures, cleanup, and successful URL responses. Assisted-by: Codex:gpt-5 * feat(studio): add inline image upscale workflow Assisted-by: OpenAI:gpt-5 * test(studio): serve upscale capability fixtures * fix(stablediffusion): default ESRGAN tiling to full frame * feat(gallery): add RealESRGAN anime upscaler * fix(stablediffusion): allow full-frame ESRGAN tiles * fix(routes): select explicit upscale models Give image upscaling dedicated usecase-filtered middleware and keep ordinary image models out of the endpoint. Mark and verify the legacy Diffusers x4 gallery model as an explicit upscaler. Assisted-by: OpenAI:gpt-5 [shell] * refactor(upscale): keep settings in backend options * test(config): keep upscale options backend-specific * fix: keep upscale routing metadata out of backend options --------- Co-authored-by: Ettore Di Giacinto --- backend/go/stablediffusion-ggml/cpp/gosd.cpp | 65 ++++ backend/go/stablediffusion-ggml/cpp/gosd.h | 7 + backend/go/stablediffusion-ggml/gosd.go | 202 +++++++++++- backend/go/stablediffusion-ggml/gosd_test.go | 312 ++++++++++++++++++ backend/go/stablediffusion-ggml/main.go | 4 + core/backend/options_internal_test.go | 18 + core/backend/upscale.go | 17 +- core/backend/upscale_test.go | 76 +++++ core/config/backend_capabilities.go | 12 +- core/config/backend_capabilities_test.go | 12 + core/config/meta/constants.go | 1 + core/config/meta/registry_test.go | 12 + core/config/model_capabilities.go | 5 +- core/config/model_capabilities_test.go | 14 + core/config/model_config.go | 12 +- core/config/model_config_test.go | 19 ++ core/config/upscale_options.go | 31 ++ core/config/upscale_options_test.go | 34 ++ core/gallery/esrgan_entries_test.go | 63 ++++ .../endpoints/openai/list_capabilities.go | 1 + .../openai/list_capabilities_test.go | 25 ++ core/http/endpoints/openai/upscale.go | 36 +- core/http/endpoints/openai/upscale_test.go | 113 +++++-- core/http/react-ui/e2e/studio-fixtures.js | 10 +- core/http/react-ui/e2e/studio-lineage.spec.js | 78 +++++ .../react-ui/public/locales/en/media.json | 15 +- .../src/components/studio/LineageView.jsx | 26 +- .../src/components/studio/UpscaleComposer.jsx | 62 ++++ core/http/react-ui/src/hooks/useStudioWork.js | 17 +- core/http/react-ui/src/pages/Studio.jsx | 2 + .../react-ui/src/pages/StudioOverview.jsx | 6 +- core/http/react-ui/src/utils/api.js | 7 + core/http/react-ui/src/utils/capabilities.js | 1 + core/http/react-ui/src/utils/studioWork.js | 5 +- core/http/react-ui/src/utils/upscale.js | 39 +++ core/http/react-ui/src/utils/upscale.test.js | 42 +++ core/http/routes/openai.go | 29 +- core/http/routes/ui_api.go | 2 + core/http/routes/ui_api_capabilities_test.go | 80 +++++ core/http/routes/upscale_test.go | 100 ++++++ core/schema/openai.go | 1 + docs/content/features/image-generation.md | 90 ++++- gallery/index.yaml | 32 +- 43 files changed, 1660 insertions(+), 75 deletions(-) create mode 100644 core/backend/upscale_test.go create mode 100644 core/config/upscale_options.go create mode 100644 core/config/upscale_options_test.go create mode 100644 core/gallery/esrgan_entries_test.go create mode 100644 core/http/react-ui/src/components/studio/UpscaleComposer.jsx create mode 100644 core/http/react-ui/src/utils/upscale.js create mode 100644 core/http/react-ui/src/utils/upscale.test.js create mode 100644 core/http/routes/ui_api_capabilities_test.go create mode 100644 core/http/routes/upscale_test.go 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: