From 8a68f3571cbea78308cf3f7ee3de49a6843a642a Mon Sep 17 00:00:00 2001 From: Pete Date: Mon, 3 Aug 2026 06:27:22 -0700 Subject: [PATCH] feat(api): add POST /v1/images/upscale endpoint (#10227) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(api): add POST /v1/images/upscale endpoint Add a new image upscaling endpoint that accepts a source image and returns an upscaled version. Supports selectable upscaler models (e.g. realesrgan) and a configurable scale factor (2x or 4x). - backend.proto: add UpscaleImage RPC and UpscaleImageRequest message - pkg/grpc: implement UpscaleImage in Backend interface, client, server and embed shim - core/backend/upscale.go: new backend helper (mirrors ImageGeneration) - core/http/endpoints/openai/upscale.go: new multipart/form-data handler - core/http/routes/openai.go: register POST /v1/images/upscale - core/http/auth/features.go: gate upscale routes under FeatureImages - backend/python/diffusers/backend.py: implement UpscaleImage — uses diffusers upscale pipeline when loaded, falls back to Lanczos resize * fix(grpc): add UpscaleImage stub to Base backend All Go backends embedding Base now satisfy the AIModel interface without needing to implement UpscaleImage explicitly. * fix(images): complete upscale endpoint integration Store generated upscales under the served images directory, validate scale factors, document and advertise the endpoint, and add a functional Stable Diffusion x4 gallery model. Assisted-by: Codex:gpt-5 --------- Co-authored-by: localai-org-maint-bot <306269227+localai-org-maint-bot@users.noreply.github.com> --- backend/backend.proto | 7 ++ backend/python/diffusers/backend.py | 28 +++++ core/backend/upscale.go | 37 ++++++ core/config/backend_capabilities.go | 3 +- core/http/auth/features.go | 2 + core/http/endpoints/openai/upscale.go | 134 +++++++++++++++++++++ core/http/endpoints/openai/upscale_test.go | 89 ++++++++++++++ core/http/routes/openai.go | 5 + core/services/nodes/health_mock_test.go | 3 + core/services/nodes/inflight.go | 6 + core/services/nodes/inflight_test.go | 4 + docs/content/features/image-generation.md | 24 ++++ gallery/index.yaml | 28 +++++ pkg/grpc/backend.go | 1 + pkg/grpc/base/base.go | 4 + pkg/grpc/client.go | 17 +++ pkg/grpc/embed.go | 4 + pkg/grpc/interface.go | 1 + pkg/grpc/server.go | 12 ++ swagger/docs.go | 63 ++++++++++ swagger/swagger.json | 63 ++++++++++ swagger/swagger.yaml | 43 +++++++ 22 files changed, 577 insertions(+), 1 deletion(-) create mode 100644 core/backend/upscale.go create mode 100644 core/http/endpoints/openai/upscale.go create mode 100644 core/http/endpoints/openai/upscale_test.go diff --git a/backend/backend.proto b/backend/backend.proto index 5d926b644..2de6c8711 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -15,6 +15,7 @@ service Backend { rpc PredictStream(PredictOptions) returns (stream Reply) {} rpc Embedding(PredictOptions) returns (EmbeddingResult) {} rpc GenerateImage(GenerateImageRequest) returns (Result) {} + rpc UpscaleImage(UpscaleImageRequest) returns (Result) {} rpc GenerateVideo(GenerateVideoRequest) returns (Result) {} rpc Generate3D(Generate3DRequest) returns (Result) {} rpc AudioTranscription(TranscriptRequest) returns (TranscriptResult) {} @@ -637,6 +638,12 @@ message GenerateImageRequest { string ModelIdentity = 13; } +message UpscaleImageRequest { + string src = 1; // input image path + string dst = 2; // output image path + int32 scale = 3; // upscale factor (e.g. 2 or 4) +} + message GenerateVideoRequest { string prompt = 1; string negative_prompt = 2; // Negative prompt for video generation diff --git a/backend/python/diffusers/backend.py b/backend/python/diffusers/backend.py index cb43c88c0..539ce5444 100755 --- a/backend/python/diffusers/backend.py +++ b/backend/python/diffusers/backend.py @@ -883,6 +883,34 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return backend_pb2.Result(message="Media generated", success=True) + def UpscaleImage(self, request, context): + try: + if not request.src: + return backend_pb2.Result(success=False, message="No source image provided") + if not request.dst: + return backend_pb2.Result(success=False, message="No destination path provided") + + scale = request.scale if request.scale > 0 else 2 + image = Image.open(request.src).convert("RGB") + + # If the loaded pipeline supports upscaling (e.g. StableDiffusionUpscalePipeline), + # use it; otherwise fall back to high-quality Lanczos resize. + if self.pipe is not None and self.PipelineType in ("StableDiffusionUpscalePipeline", "StableDiffusionLatentUpscalePipeline"): + print(f"UpscaleImage: using diffusers upscale pipeline ({self.PipelineType})", file=sys.stderr) + upscaled = self.pipe(prompt="", image=image).images[0] + else: + # Fallback: high-quality Lanczos resize + print(f"UpscaleImage: no upscale pipeline loaded, using Lanczos resize (scale={scale})", file=sys.stderr) + new_w = image.width * scale + new_h = image.height * scale + upscaled = image.resize((new_w, new_h), Image.LANCZOS) + + upscaled.save(request.dst) + return backend_pb2.Result(message="Image upscaled", success=True) + except Exception as e: + print(f"UpscaleImage error: {e}", file=sys.stderr) + return backend_pb2.Result(success=False, message=str(e)) + def GenerateVideo(self, request, context): try: prompt = request.prompt diff --git a/core/backend/upscale.go b/core/backend/upscale.go new file mode 100644 index 000000000..c821f56cd --- /dev/null +++ b/core/backend/upscale.go @@ -0,0 +1,37 @@ +package backend + +import ( + "context" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/grpc/proto" + model "github.com/mudler/LocalAI/pkg/model" +) + +// ImageUpscale loads the model specified in modelConfig and calls UpscaleImage +// on the backend, writing the result to dst. +func ImageUpscale(ctx context.Context, src, dst string, scale int, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (func() error, error) { + opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx)) + inferenceModel, err := loader.Load(opts...) + if err != nil { + recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil) + return nil, err + } + + fn := func() error { + _, err := inferenceModel.UpscaleImage( + ctx, + &proto.UpscaleImageRequest{ + Src: src, + Dst: dst, + Scale: int32(scale), + }, + ) + return err + } + + return fn, nil +} + +// ImageUpscaleFunc is a test-friendly indirection. +var ImageUpscaleFunc = ImageUpscale diff --git a/core/config/backend_capabilities.go b/core/config/backend_capabilities.go index 844d418ca..13901d35b 100644 --- a/core/config/backend_capabilities.go +++ b/core/config/backend_capabilities.go @@ -44,6 +44,7 @@ const ( MethodPredictStream GRPCMethod = "PredictStream" MethodEmbedding GRPCMethod = "Embedding" MethodGenerateImage GRPCMethod = "GenerateImage" + MethodUpscaleImage GRPCMethod = "UpscaleImage" MethodGenerateVideo GRPCMethod = "GenerateVideo" MethodGenerate3D GRPCMethod = "Generate3D" MethodAudioTranscription GRPCMethod = "AudioTranscription" @@ -348,7 +349,7 @@ var BackendCapabilities = map[string]BackendCapability{ // --- Image/video generation backends --- "diffusers": { - GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodGenerateVideo}, + GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodUpscaleImage, MethodGenerateVideo}, PossibleUsecases: []string{UsecaseImage, UsecaseVideo}, DefaultUsecases: []string{UsecaseImage}, Description: "HuggingFace diffusers — Stable Diffusion, Flux, video generation", diff --git a/core/http/auth/features.go b/core/http/auth/features.go index 36fdafe12..dac7259fe 100644 --- a/core/http/auth/features.go +++ b/core/http/auth/features.go @@ -39,6 +39,8 @@ var RouteFeatureRegistry = []RouteFeature{ {"POST", "/images/generations", FeatureImages}, {"POST", "/v1/images/inpainting", FeatureImages}, {"POST", "/images/inpainting", FeatureImages}, + {"POST", "/v1/images/upscale", FeatureImages}, + {"POST", "/images/upscale", FeatureImages}, // Audio transcription {"POST", "/v1/audio/transcriptions", FeatureAudioTranscription}, diff --git a/core/http/endpoints/openai/upscale.go b/core/http/endpoints/openai/upscale.go new file mode 100644 index 000000000..eab000f8e --- /dev/null +++ b/core/http/endpoints/openai/upscale.go @@ -0,0 +1,134 @@ +package openai + +import ( + "fmt" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + "time" + + "github.com/google/uuid" + "github.com/labstack/echo/v4" + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + model "github.com/mudler/LocalAI/pkg/model" +) + +// UpscaleEndpoint handles POST /v1/images/upscale +// +// @Summary Image upscaling +// @Description Upscale an image using a specified model (e.g. stable-diffusion-x4-upscaler). Accepts multipart/form-data. +// @Tags images +// @Accept multipart/form-data +// @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)" +// @Success 200 {object} schema.OpenAIResponse +// @Failure 400 {object} map[string]string +// @Failure 500 {object} map[string]string +// @Router /v1/images/upscale [post] +func UpscaleEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return func(c echo.Context) error { + modelName := c.FormValue("model") + scaleStr := c.FormValue("scale") + + if modelName == "" { + xlog.Error("Upscale Endpoint - missing model") + 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 + } + + // Read uploaded image + imageFile, err := c.FormFile("image") + if err != nil { + xlog.Error("Upscale Endpoint - missing image file", "error", err) + return echo.NewHTTPError(http.StatusBadRequest, "missing image file") + } + + imgSrc, err := imageFile.Open() + if err != nil { + return err + } + defer imgSrc.Close() + imgBytes, err := io.ReadAll(imgSrc) + if err != nil { + return err + } + + // Get model config from middleware context + cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || cfg == nil { + xlog.Error("Upscale Endpoint - model config not found in context") + return echo.ErrBadRequest + } + + tmpDir := filepath.Join(appConfig.GeneratedContentDir, "images") + if err := os.MkdirAll(tmpDir, 0750); err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, "failed to prepare storage") + } + + // Write input image to a temp file + srcTmp, err := os.CreateTemp(tmpDir, "upscale_src_") + if err != nil { + return err + } + if _, err := srcTmp.Write(imgBytes); err != nil { + _ = srcTmp.Close() + _ = os.Remove(srcTmp.Name()) + return err + } + if err := srcTmp.Close(); err != nil { + xlog.Warn("Upscale Endpoint - failed to close src temp file", "error", err) + } + srcPath := srcTmp.Name() + defer os.Remove(srcPath) + + // Prepare output file path + id := uuid.New().String() + dstPath := filepath.Join(tmpDir, fmt.Sprintf("upscale_%s.png", id)) + + 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 + } + + created := int(time.Now().Unix()) + resp := &schema.OpenAIResponse{ + ID: id, + Created: created, + Data: []schema.Item{{URL: imgURL}}, + Usage: &schema.OpenAIUsage{ + InputTokensDetails: &schema.InputTokensDetails{}, + }, + } + + return c.JSON(http.StatusOK, resp) + } +} diff --git a/core/http/endpoints/openai/upscale_test.go b/core/http/endpoints/openai/upscale_test.go new file mode 100644 index 000000000..2e13c8fc0 --- /dev/null +++ b/core/http/endpoints/openai/upscale_test.go @@ -0,0 +1,89 @@ +package openai + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + model "github.com/mudler/LocalAI/pkg/model" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Image upscaling", func() { + var ( + appConfig *config.ApplicationConfig + tmpDir string + ) + + BeforeEach(func() { + var err error + tmpDir, err = os.MkdirTemp("", "upscale") + Expect(err).ToNot(HaveOccurred()) + appConfig = config.NewApplicationConfig(config.WithGeneratedContentDir(tmpDir)) + }) + + AfterEach(func() { + Expect(os.RemoveAll(tmpDir)).To(Succeed()) + }) + + 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)) + return func() error { + return os.WriteFile(dst, []byte("PNGDATA"), 0o644) + }, nil + } + DeferCleanup(func() { backend.ImageUpscaleFunc = original }) + + req, _ := makeMultipartRequest( + map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "4"}, + map[string][]byte{"image": []byte("IMAGEDATA")}, + ) + rec := httptest.NewRecorder() + ctx := echo.New().NewContext(req, rec) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Backend: "diffusers"}) + + Expect(UpscaleEndpoint(nil, nil, appConfig)(ctx)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusOK)) + + var response schema.OpenAIResponse + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + Expect(response.Data).To(HaveLen(1)) + Expect(response.Data[0].URL).To(ContainSubstring("/generated-images/upscale_")) + + filename := filepath.Base(response.Data[0].URL) + contents, err := os.ReadFile(filepath.Join(tmpDir, "images", filename)) + Expect(err).ToNot(HaveOccurred()) + Expect(contents).To(Equal([]byte("PNGDATA"))) + }) + + It("rejects unsupported scale factors", func() { + req, _ := makeMultipartRequest( + map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "3"}, + map[string][]byte{"image": []byte("IMAGEDATA")}, + ) + rec := httptest.NewRecorder() + ctx := echo.New().NewContext(req, rec) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Backend: "diffusers"}) + + 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()) + }) +}) diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index 9c118bdbc..b1671077a 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -254,6 +254,11 @@ 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 + upscaleHandler := openai.UpscaleEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) + app.POST("/v1/images/upscale", upscaleHandler, imageMiddleware...) + app.POST("/images/upscale", upscaleHandler, imageMiddleware...) + // List models app.GET("/v1/models", openai.ListModelsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig(), application.AuthDB())) app.GET("/models", openai.ListModelsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig(), application.AuthDB())) diff --git a/core/services/nodes/health_mock_test.go b/core/services/nodes/health_mock_test.go index b2583fadd..c52712dab 100644 --- a/core/services/nodes/health_mock_test.go +++ b/core/services/nodes/health_mock_test.go @@ -154,6 +154,9 @@ func (c *fakeBackendClient) Predict(_ context.Context, _ *pb.PredictOptions, _ . func (c *fakeBackendClient) GenerateImage(_ context.Context, _ *pb.GenerateImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { return nil, nil } +func (c *fakeBackendClient) UpscaleImage(_ context.Context, _ *pb.UpscaleImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return nil, nil +} func (c *fakeBackendClient) GenerateVideo(_ context.Context, _ *pb.GenerateVideoRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { return nil, nil } diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index de851121c..3102a3254 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -138,6 +138,12 @@ func (c *InFlightTrackingClient) GenerateImage(ctx context.Context, in *pb.Gener return res, c.reconcile(err) } +func (c *InFlightTrackingClient) UpscaleImage(ctx context.Context, in *pb.UpscaleImageRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { + defer c.track(ctx)() + res, err := c.inner.UpscaleImage(ctx, in, opts...) + return res, c.reconcile(err) +} + func (c *InFlightTrackingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { defer c.track(ctx)() res, err := c.inner.GenerateVideo(ctx, in, opts...) diff --git a/core/services/nodes/inflight_test.go b/core/services/nodes/inflight_test.go index 98cc03557..7726e720e 100644 --- a/core/services/nodes/inflight_test.go +++ b/core/services/nodes/inflight_test.go @@ -83,6 +83,10 @@ func (f *fakeGRPCBackend) GenerateImage(_ context.Context, _ *pb.GenerateImageRe return &pb.Result{}, nil } +func (f *fakeGRPCBackend) UpscaleImage(_ context.Context, _ *pb.UpscaleImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return &pb.Result{}, nil +} + func (f *fakeGRPCBackend) GenerateVideo(_ context.Context, _ *pb.GenerateVideoRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { return &pb.Result{}, nil } diff --git a/docs/content/features/image-generation.md b/docs/content/features/image-generation.md index aca97b236..d94c0f9b4 100644 --- a/docs/content/features/image-generation.md +++ b/docs/content/features/image-generation.md @@ -164,6 +164,30 @@ By default the RPC devices join the pool and participate in placement; combine w ![anime_girl](https://github.com/go-skynet/LocalAI/assets/2420543/8aaca62a-e864-4011-98ae-dcc708103928) (Generated with [AnimagineXL](https://huggingface.co/Linaqruf/animagine-xl)) +#### 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 +gallery model first: + +```bash +local-ai models install stable-diffusion-x4-upscaler +``` + +Then send the model name, scale, and image as multipart form fields: + +```bash +curl http://localhost:8080/v1/images/upscale \ + -F model=stable-diffusion-x4-upscaler \ + -F scale=4 \ + -F image=@input.png +``` + +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. + #### Model setup The models will be downloaded the first time you use the backend from `huggingface` automatically. diff --git a/gallery/index.yaml b/gallery/index.yaml index 1825f4f7b..32ad2a436 100644 --- a/gallery/index.yaml +++ b/gallery/index.yaml @@ -9015,6 +9015,34 @@ - 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 +- name: stable-diffusion-x4-upscaler + url: github:mudler/LocalAI/gallery/virtual.yaml@master + urls: + - https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler + description: | + Stable Diffusion x4 Upscaler is Stability AI's diffusion-based super-resolution model. It enlarges low-resolution images by four times while reconstructing image detail. + license: openrail++ + tags: + - image-upscaling + - super-resolution + - image-to-image + - diffusers + last_checked: "2026-07-30" + overrides: + backend: diffusers + f16: true + diffusers: + pipeline_type: StableDiffusionUpscalePipeline + known_usecases: + - image + parameters: + model: stabilityai/stable-diffusion-x4-upscaler + artifacts: + - name: model + target: model + source: + type: huggingface + repo: stabilityai/stable-diffusion-x4-upscaler - name: z-image-diffusers url: github:mudler/LocalAI/gallery/virtual.yaml@master urls: diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 32564655f..93dde0099 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -72,6 +72,7 @@ type InferenceBackend interface { PredictStream(ctx context.Context, in *pb.PredictOptions, f func(reply *pb.Reply), opts ...grpc.CallOption) error Predict(ctx context.Context, in *pb.PredictOptions, opts ...grpc.CallOption) (*pb.Reply, error) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, opts ...grpc.CallOption) (*pb.Result, error) + UpscaleImage(ctx context.Context, in *pb.UpscaleImageRequest, opts ...grpc.CallOption) (*pb.Result, error) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...grpc.CallOption) (*pb.Result, error) Generate3D(ctx context.Context, in *pb.Generate3DRequest, opts ...grpc.CallOption) (*pb.Result, error) TTS(ctx context.Context, in *pb.TTSRequest, opts ...grpc.CallOption) (*pb.Result, error) diff --git a/pkg/grpc/base/base.go b/pkg/grpc/base/base.go index 94e4e7462..46cf367a8 100644 --- a/pkg/grpc/base/base.go +++ b/pkg/grpc/base/base.go @@ -63,6 +63,10 @@ func (llm *Base) Generate3D(*pb.Generate3DRequest) error { return fmt.Errorf("unimplemented") } +func (llm *Base) UpscaleImage(*pb.UpscaleImageRequest) error { + return fmt.Errorf("unimplemented") +} + func (llm *Base) AudioTranscription(context.Context, *pb.TranscriptRequest) (pb.TranscriptResult, error) { return pb.TranscriptResult{}, fmt.Errorf("unimplemented") } diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index d42978c4e..506cee725 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -230,6 +230,23 @@ func (c *Client) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, return client.GenerateImage(ctx, in, opts...) } +func (c *Client) UpscaleImage(ctx context.Context, in *pb.UpscaleImageRequest, opts ...grpc.CallOption) (*pb.Result, error) { + if !c.parallel { + c.opMutex.Lock() + defer c.opMutex.Unlock() + } + c.setBusy(true) + defer c.setBusy(false) + defer c.wdMark()() + conn, err := c.dial() + if err != nil { + return nil, err + } + defer conn.Close() + client := pb.NewBackendClient(conn) + return client.UpscaleImage(ctx, in, opts...) +} + func (c *Client) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...grpc.CallOption) (*pb.Result, error) { if !c.parallel { c.opMutex.Lock() diff --git a/pkg/grpc/embed.go b/pkg/grpc/embed.go index a819d3a53..6b616dfa0 100644 --- a/pkg/grpc/embed.go +++ b/pkg/grpc/embed.go @@ -49,6 +49,10 @@ func (e *embedBackend) GenerateImage(ctx context.Context, in *pb.GenerateImageRe return e.s.GenerateImage(ctx, in) } +func (e *embedBackend) UpscaleImage(ctx context.Context, in *pb.UpscaleImageRequest, opts ...grpc.CallOption) (*pb.Result, error) { + return e.s.UpscaleImage(ctx, in) +} + func (e *embedBackend) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...grpc.CallOption) (*pb.Result, error) { return e.s.GenerateVideo(ctx, in) } diff --git a/pkg/grpc/interface.go b/pkg/grpc/interface.go index 974474546..37a9d52a3 100644 --- a/pkg/grpc/interface.go +++ b/pkg/grpc/interface.go @@ -17,6 +17,7 @@ type AIModel interface { Free() error Embeddings(*pb.PredictOptions) ([]float32, error) GenerateImage(*pb.GenerateImageRequest) error + UpscaleImage(*pb.UpscaleImageRequest) error GenerateVideo(*pb.GenerateVideoRequest) error Generate3D(*pb.Generate3DRequest) error Detect(*pb.DetectOptions) (pb.DetectResponse, error) diff --git a/pkg/grpc/server.go b/pkg/grpc/server.go index 903532e4b..528ba0a36 100644 --- a/pkg/grpc/server.go +++ b/pkg/grpc/server.go @@ -150,6 +150,18 @@ func (s *server) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest) return &pb.Result{Message: "Image generated", Success: true}, nil } +func (s *server) UpscaleImage(ctx context.Context, in *pb.UpscaleImageRequest) (*pb.Result, error) { + if s.llm.Locking() { + s.llm.Lock() + defer s.llm.Unlock() + } + err := s.llm.UpscaleImage(in) + if err != nil { + return &pb.Result{Message: fmt.Sprintf("Error upscaling image: %s", err.Error()), Success: false}, err + } + return &pb.Result{Message: "Image upscaled", Success: true}, nil +} + func (s *server) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest) (*pb.Result, error) { if err := s.checkModelIdentity(in); err != nil { return nil, err diff --git a/swagger/docs.go b/swagger/docs.go index f9b79654a..b69860b01 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -3122,6 +3122,69 @@ const docTemplate = `{ } } }, + "/v1/images/upscale": { + "post": { + "description": "Upscale an image using a specified model (e.g. stable-diffusion-x4-upscaler). Accepts multipart/form-data.", + "consumes": [ + "multipart/form-data" + ], + "produces": [ + "application/json" + ], + "tags": [ + "images" + ], + "summary": "Image upscaling", + "parameters": [ + { + "type": "string", + "description": "Upscaler model identifier (e.g. stable-diffusion-x4-upscaler)", + "name": "model", + "in": "formData", + "required": true + }, + { + "type": "file", + "description": "Input image file", + "name": "image", + "in": "formData", + "required": true + }, + { + "type": "integer", + "description": "Upscale factor: 2 or 4 (default 2)", + "name": "scale", + "in": "formData" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.OpenAIResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, "/v1/mcp/chat/completions": { "post": { "tags": [ diff --git a/swagger/swagger.json b/swagger/swagger.json index fc6b61921..8ce062051 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -3119,6 +3119,69 @@ } } }, + "/v1/images/upscale": { + "post": { + "description": "Upscale an image using a specified model (e.g. stable-diffusion-x4-upscaler). Accepts multipart/form-data.", + "consumes": [ + "multipart/form-data" + ], + "produces": [ + "application/json" + ], + "tags": [ + "images" + ], + "summary": "Image upscaling", + "parameters": [ + { + "type": "string", + "description": "Upscaler model identifier (e.g. stable-diffusion-x4-upscaler)", + "name": "model", + "in": "formData", + "required": true + }, + { + "type": "file", + "description": "Input image file", + "name": "image", + "in": "formData", + "required": true + }, + { + "type": "integer", + "description": "Upscale factor: 2 or 4 (default 2)", + "name": "scale", + "in": "formData" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.OpenAIResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, "/v1/mcp/chat/completions": { "post": { "tags": [ diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index 9bbdc66c7..2963dfd2c 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -4857,6 +4857,49 @@ paths: summary: Image inpainting tags: - images + /v1/images/upscale: + post: + consumes: + - multipart/form-data + description: Upscale an image using a specified model (e.g. stable-diffusion-x4-upscaler). + Accepts multipart/form-data. + parameters: + - description: Upscaler model identifier (e.g. stable-diffusion-x4-upscaler) + in: formData + name: model + required: true + type: string + - description: Input image file + in: formData + name: image + required: true + type: file + - description: 'Upscale factor: 2 or 4 (default 2)' + in: formData + name: scale + type: integer + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/schema.OpenAIResponse' + "400": + description: Bad Request + schema: + additionalProperties: + type: string + type: object + "500": + description: Internal Server Error + schema: + additionalProperties: + type: string + type: object + summary: Image upscaling + tags: + - images /v1/mcp/chat/completions: post: parameters: