Skip to content

Commit 68bfe00

Browse files
gmaOCRgmaOCR
committed
fix(inpainting): write outputs to GeneratedContentDir/images and construct generated-images URL via url.JoinPath
Signed-off-by: Greg <[email protected]>
1 parent a3e93b6 commit 68bfe00

File tree

2 files changed

+167
-158
lines changed

2 files changed

+167
-158
lines changed

‎core/http/endpoints/openai/inpainting.go‎

Lines changed: 166 additions & 157 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,26 @@
11
package openai
22

33
import (
4-
"encoding/base64"
5-
"encoding/json"
6-
"fmt"
7-
"io"
8-
"net/http"
9-
"os"
10-
"path/filepath"
11-
"strconv"
12-
"time"
13-
14-
"github.com/google/uuid"
15-
"github.com/labstack/echo/v4"
16-
"github.com/rs/zerolog/log"
17-
18-
"github.com/mudler/LocalAI/core/config"
19-
"github.com/mudler/LocalAI/core/http/middleware"
20-
"github.com/mudler/LocalAI/core/schema"
21-
"github.com/mudler/LocalAI/core/backend"
22-
model "github.com/mudler/LocalAI/pkg/model"
4+
"encoding/base64"
5+
"encoding/json"
6+
"fmt"
7+
"io"
8+
"net/http"
9+
"net/url"
10+
"os"
11+
"path/filepath"
12+
"strconv"
13+
"time"
14+
15+
"github.com/google/uuid"
16+
"github.com/labstack/echo/v4"
17+
"github.com/rs/zerolog/log"
18+
19+
"github.com/mudler/LocalAI/core/backend"
20+
"github.com/mudler/LocalAI/core/config"
21+
"github.com/mudler/LocalAI/core/http/middleware"
22+
"github.com/mudler/LocalAI/core/schema"
23+
model "github.com/mudler/LocalAI/pkg/model"
2324
)
2425

2526
// InpaintingEndpoint handles POST /v1/images/inpainting
@@ -40,142 +41,150 @@ import (
4041
// @Failure 500 {object} map[string]string
4142
// @Router /v1/images/inpainting [post]
4243
func InpaintingEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc {
43-
return func(c echo.Context) error {
44-
// Parse basic form values
45-
modelName := c.FormValue("model")
46-
prompt := c.FormValue("prompt")
47-
stepsStr := c.FormValue("steps")
48-
49-
if modelName == "" || prompt == "" {
50-
log.Error().Msg("Inpainting Endpoint - missing model or prompt")
51-
return echo.ErrBadRequest
52-
}
53-
54-
// steps default
55-
steps := 25
56-
if stepsStr != "" {
57-
if v, err := strconv.Atoi(stepsStr); err == nil {
58-
steps = v
59-
}
60-
}
61-
62-
// Get uploaded files
63-
imageFile, err := c.FormFile("image")
64-
if err != nil {
65-
log.Error().Err(err).Msg("Inpainting Endpoint - missing image file")
66-
return echo.NewHTTPError(http.StatusBadRequest, "missing image file")
67-
}
68-
maskFile, err := c.FormFile("mask")
69-
if err != nil {
70-
log.Error().Err(err).Msg("Inpainting Endpoint - missing mask file")
71-
return echo.NewHTTPError(http.StatusBadRequest, "missing mask file")
72-
}
73-
74-
// Read files into memory (small files expected)
75-
imgSrc, err := imageFile.Open()
76-
if err != nil {
77-
return err
78-
}
79-
defer imgSrc.Close()
80-
imgBytes, err := io.ReadAll(imgSrc)
81-
if err != nil {
82-
return err
83-
}
84-
85-
maskSrc, err := maskFile.Open()
86-
if err != nil {
87-
return err
88-
}
89-
defer maskSrc.Close()
90-
maskBytes, err := io.ReadAll(maskSrc)
91-
if err != nil {
92-
return err
93-
}
94-
95-
// Create JSON with base64 fields expected by backend
96-
b64Image := base64.StdEncoding.EncodeToString(imgBytes)
97-
b64Mask := base64.StdEncoding.EncodeToString(maskBytes)
98-
99-
// get model config from context (middleware set it)
100-
cfg, ok := c.Get("MODEL_CONFIG").(*config.ModelConfig)
101-
if !ok || cfg == nil {
102-
log.Error().Msg("Inpainting Endpoint - model config not found in context")
103-
return echo.ErrBadRequest
104-
}
105-
106-
tmpDir := appConfig.GeneratedContentDir
107-
id := uuid.New().String()
108-
jsonName := fmt.Sprintf("inpaint_%s.json", id)
109-
jsonPath := filepath.Join(tmpDir, jsonName)
110-
jsonFile := map[string]string{
111-
"image": b64Image,
112-
"mask_image": b64Mask,
113-
}
114-
jf, err := os.CreateTemp(tmpDir, "inpaint_")
115-
if err != nil {
116-
return err
117-
}
118-
// write JSON
119-
enc := json.NewEncoder(jf)
120-
if err := enc.Encode(jsonFile); err != nil {
121-
jf.Close()
122-
os.Remove(jf.Name())
123-
return err
124-
}
125-
jf.Close()
126-
// rename to desired name
127-
if err := os.Rename(jf.Name(), jsonPath); err != nil {
128-
os.Remove(jf.Name())
129-
return err
130-
}
131-
// prepare dst
132-
outTmp, err := os.CreateTemp(tmpDir, "out_")
133-
if err != nil {
134-
os.Remove(jsonPath)
135-
return err
136-
}
137-
outTmp.Close()
138-
dst := outTmp.Name() + ".png"
139-
if err := os.Rename(outTmp.Name(), dst); err != nil {
140-
os.Remove(jsonPath)
141-
return err
142-
}
143-
144-
// Determine width/height default
145-
width := 512
146-
height := 512
147-
148-
// Call backend image generation via indirection so tests can stub it
149-
// Note: ImageGenerationFunc will call into the loaded model's GenerateImage which expects src JSON
150-
fn, err := backend.ImageGenerationFunc(height, width, 0, steps, 0, prompt, "", jsonPath, dst, ml, *cfg, appConfig, nil)
151-
if err != nil {
152-
os.Remove(jsonPath)
153-
return err
154-
}
155-
156-
// Execute generation function (blocking)
157-
if err := fn(); err != nil {
158-
os.Remove(jsonPath)
159-
os.Remove(dst)
160-
return err
161-
}
162-
163-
// On success, build response URL using BaseURL middleware helper
164-
baseURL := middleware.BaseURL(c)
165-
166-
// Return response
167-
created := int(time.Now().Unix())
168-
resp := &schema.OpenAIResponse{
169-
ID: id,
170-
Created: created,
171-
Data: []schema.Item{{
172-
URL: fmt.Sprintf("%sgenerated-images/%s", baseURL, filepath.Base(dst)),
173-
}},
174-
}
175-
176-
// cleanup json
177-
defer os.Remove(jsonPath)
178-
179-
return c.JSON(http.StatusOK, resp)
180-
}
44+
return func(c echo.Context) error {
45+
// Parse basic form values
46+
modelName := c.FormValue("model")
47+
prompt := c.FormValue("prompt")
48+
stepsStr := c.FormValue("steps")
49+
50+
if modelName == "" || prompt == "" {
51+
log.Error().Msg("Inpainting Endpoint - missing model or prompt")
52+
return echo.ErrBadRequest
53+
}
54+
55+
// steps default
56+
steps := 25
57+
if stepsStr != "" {
58+
if v, err := strconv.Atoi(stepsStr); err == nil {
59+
steps = v
60+
}
61+
}
62+
63+
// Get uploaded files
64+
imageFile, err := c.FormFile("image")
65+
if err != nil {
66+
log.Error().Err(err).Msg("Inpainting Endpoint - missing image file")
67+
return echo.NewHTTPError(http.StatusBadRequest, "missing image file")
68+
}
69+
maskFile, err := c.FormFile("mask")
70+
if err != nil {
71+
log.Error().Err(err).Msg("Inpainting Endpoint - missing mask file")
72+
return echo.NewHTTPError(http.StatusBadRequest, "missing mask file")
73+
}
74+
75+
// Read files into memory (small files expected)
76+
imgSrc, err := imageFile.Open()
77+
if err != nil {
78+
return err
79+
}
80+
defer imgSrc.Close()
81+
imgBytes, err := io.ReadAll(imgSrc)
82+
if err != nil {
83+
return err
84+
}
85+
86+
maskSrc, err := maskFile.Open()
87+
if err != nil {
88+
return err
89+
}
90+
defer maskSrc.Close()
91+
maskBytes, err := io.ReadAll(maskSrc)
92+
if err != nil {
93+
return err
94+
}
95+
96+
// Create JSON with base64 fields expected by backend
97+
b64Image := base64.StdEncoding.EncodeToString(imgBytes)
98+
b64Mask := base64.StdEncoding.EncodeToString(maskBytes)
99+
100+
// get model config from context (middleware set it)
101+
cfg, ok := c.Get("MODEL_CONFIG").(*config.ModelConfig)
102+
if !ok || cfg == nil {
103+
log.Error().Msg("Inpainting Endpoint - model config not found in context")
104+
return echo.ErrBadRequest
105+
}
106+
107+
// Use the images subdirectory under GeneratedContentDir so the generated
108+
// PNG is placed where the HTTP static handler serves `/generated-images`.
109+
tmpDir := filepath.Join(appConfig.GeneratedContentDir, "images")
110+
id := uuid.New().String()
111+
jsonName := fmt.Sprintf("inpaint_%s.json", id)
112+
jsonPath := filepath.Join(tmpDir, jsonName)
113+
jsonFile := map[string]string{
114+
"image": b64Image,
115+
"mask_image": b64Mask,
116+
}
117+
jf, err := os.CreateTemp(tmpDir, "inpaint_")
118+
if err != nil {
119+
return err
120+
}
121+
// write JSON
122+
enc := json.NewEncoder(jf)
123+
if err := enc.Encode(jsonFile); err != nil {
124+
jf.Close()
125+
os.Remove(jf.Name())
126+
return err
127+
}
128+
jf.Close()
129+
// rename to desired name
130+
if err := os.Rename(jf.Name(), jsonPath); err != nil {
131+
os.Remove(jf.Name())
132+
return err
133+
}
134+
// prepare dst
135+
outTmp, err := os.CreateTemp(tmpDir, "out_")
136+
if err != nil {
137+
os.Remove(jsonPath)
138+
return err
139+
}
140+
outTmp.Close()
141+
dst := outTmp.Name() + ".png"
142+
if err := os.Rename(outTmp.Name(), dst); err != nil {
143+
os.Remove(jsonPath)
144+
return err
145+
}
146+
147+
// Determine width/height default
148+
width := 512
149+
height := 512
150+
151+
// Call backend image generation via indirection so tests can stub it
152+
// Note: ImageGenerationFunc will call into the loaded model's GenerateImage which expects src JSON
153+
fn, err := backend.ImageGenerationFunc(height, width, 0, steps, 0, prompt, "", jsonPath, dst, ml, *cfg, appConfig, nil)
154+
if err != nil {
155+
os.Remove(jsonPath)
156+
return err
157+
}
158+
159+
// Execute generation function (blocking)
160+
if err := fn(); err != nil {
161+
os.Remove(jsonPath)
162+
os.Remove(dst)
163+
return err
164+
}
165+
166+
// On success, build response URL using BaseURL middleware helper and
167+
// the same `generated-images` prefix used by the server static mount.
168+
baseURL := middleware.BaseURL(c)
169+
170+
// Build response using url.JoinPath for correct URL escaping
171+
imgPath, err := url.JoinPath(baseURL, "generated-images", filepath.Base(dst))
172+
if err != nil {
173+
return err
174+
}
175+
176+
created := int(time.Now().Unix())
177+
resp := &schema.OpenAIResponse{
178+
ID: id,
179+
Created: created,
180+
Data: []schema.Item{{
181+
URL: imgPath,
182+
}},
183+
}
184+
185+
// cleanup json
186+
defer os.Remove(jsonPath)
187+
188+
return c.JSON(http.StatusOK, resp)
189+
}
181190
}

‎core/http/routes/openai.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
140140
// images
141141
imageHandler := openai.ImageEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig())
142142
imageMiddleware := []echo.MiddlewareFunc{
143-
// Par défaut, utiliser le modèle d'inpainting souhaité pour les endpoints images/inpainting
143+
// Default: use the desired inpainting model for image endpoints
144144
re.BuildConstantDefaultModelNameMiddleware("dreamshaper-8-inpainting"),
145145
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
146146
func(next echo.HandlerFunc) echo.HandlerFunc {

0 commit comments

Comments
 (0)