Skip to content

Commit 5a67b5d

Browse files
Fix image upload processing and img2img pipeline in diffusers backend (mudler#8879)
* fix: add missing bufio.Flush in processImageFile The processImageFile function writes decoded image data (from base64 or URL download) through a bufio.NewWriter but never calls Flush() before closing the underlying file. Since bufio's default buffer is 4096 bytes, small images produce 0-byte files and large images are truncated — causing PIL to fail with "cannot identify image file". This breaks all image input paths: file, files, and ref_images parameters in /v1/images/generations, making img2img, inpainting, and reference image features non-functional. Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> * fix: merge options into kwargs in diffusers GenerateImage The GenerateImage method builds a local `options` dict containing the source image (PIL), negative_prompt, and num_inference_steps, but never merges it into `kwargs` before calling self.pipe(**kwargs). This causes img2img to fail with "Input is in incorrect format" because the pipeline never receives the image parameter. Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> * test: add unit test for processImageFile base64 decoding Verifies that a base64-encoded PNG survives the write path (encode → decode → bufio.Write → Flush → file on disk) with byte-for-byte fidelity. The test image is small enough to fit entirely in bufio's 4096-byte buffer, which is the exact scenario where the missing Flush() produced a 0-byte file. Also tests that invalid base64 input is handled gracefully. Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> * test: verify GenerateImage merges options into pipeline kwargs Mocks the diffusers pipeline and calls GenerateImage with a source image and negative prompt. Asserts that the pipeline receives the image, negative_prompt, and num_inference_steps via kwargs — the exact parameters that were silently dropped before the fix. Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> * fix: move kwargs.update(options) earlier in GenerateImage Move the options merge right after self.options merge (L742) so that image, negative_prompt, and num_inference_steps are available to all downstream code paths including img2vid and txt2vid. Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> * test: convert processImageFile tests to ginkgo Replace standard testing with ginkgo/gomega to be consistent with the rest of the test suites in the project. Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> --------- Signed-off-by: Attila Györffy <attila+git@attilagyorffy.com> Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
1 parent 270eb95 commit 5a67b5d

5 files changed

Lines changed: 137 additions & 1 deletion

File tree

backend/python/diffusers/backend.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -741,6 +741,8 @@ def GenerateImage(self, request, context):
741741
# populate kwargs from self.options.
742742
kwargs.update(self.options)
743743

744+
kwargs.update(options)
745+
744746
# Set seed
745747
if request.seed > 0:
746748
kwargs["generator"] = torch.Generator(device=self.device).manual_seed(

backend/python/diffusers/test.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -312,3 +312,64 @@ def test_resolve_with_model_id_uses_diffusion_pipeline_fallback(self):
312312
# or fail depending on network, but the fallback path should work.
313313
cls = loader.resolve_pipeline_class(model_id="some/nonexistent/model")
314314
self.assertEqual(cls, DiffusionPipeline)
315+
316+
317+
@unittest.skipUnless(GRPC_AVAILABLE, "gRPC modules not available")
318+
class TestGenerateImageOptionsKwargsMerge(unittest.TestCase):
319+
"""Test that GenerateImage merges the options dict into pipeline kwargs.
320+
321+
The options dict holds image (PIL), negative_prompt, and
322+
num_inference_steps. Without the merge, img2img pipelines never
323+
receive the source image and fail with 'Input is in incorrect format'.
324+
"""
325+
326+
def test_options_merged_into_pipeline_kwargs(self):
327+
from backend import BackendServicer
328+
from PIL import Image
329+
import tempfile, os
330+
331+
svc = BackendServicer.__new__(BackendServicer)
332+
# Minimal attributes the method reads
333+
svc.pipe = MagicMock()
334+
svc.pipe.return_value.images = [Image.new("RGB", (4, 4))]
335+
svc.cfg_scale = 7.5
336+
svc.controlnet = None
337+
svc.img2vid = False
338+
svc.txt2vid = False
339+
svc.clip_skip = 0
340+
svc.PipelineType = "StableDiffusionImg2ImgPipeline"
341+
svc.options = {}
342+
343+
# Create a tiny source image for the request's src field
344+
src_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
345+
Image.new("RGB", (4, 4), color="red").save(src_file, format="PNG")
346+
src_file.close()
347+
348+
dst_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
349+
dst_file.close()
350+
351+
try:
352+
request = MagicMock()
353+
request.positive_prompt = "a test prompt"
354+
request.negative_prompt = "bad quality"
355+
request.step = 10
356+
request.seed = 0
357+
request.width = 0
358+
request.height = 0
359+
request.src = src_file.name
360+
request.ref_images = []
361+
request.dst = dst_file.name
362+
363+
svc.GenerateImage(request, context=None)
364+
365+
# The pipeline must have been called with the image kwarg
366+
svc.pipe.assert_called_once()
367+
_, call_kwargs = svc.pipe.call_args
368+
self.assertIn("image", call_kwargs,
369+
"source image must be passed to pipeline via kwargs")
370+
self.assertIn("negative_prompt", call_kwargs,
371+
"negative_prompt must be passed to pipeline via kwargs")
372+
self.assertEqual(call_kwargs["num_inference_steps"], 10)
373+
finally:
374+
os.unlink(src_file.name)
375+
os.unlink(dst_file.name)

core/http/endpoints/openai/image.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -288,14 +288,19 @@ func processImageFile(file string, generatedContentDir string) string {
288288
return ""
289289
}
290290

291-
// write the base64 result
291+
// write the decoded result
292292
writer := bufio.NewWriter(outputFile)
293293
_, err = writer.Write(fileData)
294294
if err != nil {
295295
outputFile.Close()
296296
xlog.Error("Failed writing to temporary file", "error", err)
297297
return ""
298298
}
299+
if err := writer.Flush(); err != nil {
300+
outputFile.Close()
301+
xlog.Error("Failed flushing to temporary file", "error", err)
302+
return ""
303+
}
299304
outputFile.Close()
300305

301306
return outputFile.Name()
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
package openai
2+
3+
import (
4+
"encoding/base64"
5+
"os"
6+
7+
. "github.com/onsi/ginkgo/v2"
8+
. "github.com/onsi/gomega"
9+
)
10+
11+
var _ = Describe("processImageFile", func() {
12+
var tmpDir string
13+
14+
BeforeEach(func() {
15+
var err error
16+
tmpDir, err = os.MkdirTemp("", "processimage")
17+
Expect(err).ToNot(HaveOccurred())
18+
})
19+
20+
AfterEach(func() {
21+
os.RemoveAll(tmpDir)
22+
})
23+
24+
It("should decode base64 and write all bytes to disk", func() {
25+
// 4x4 red pixel PNG (68 bytes raw) — small enough to fit in bufio's
26+
// default 4096-byte buffer, which is exactly the scenario where a
27+
// missing Flush() produces a 0-byte file.
28+
pngBytes := []byte{
29+
0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, // PNG signature
30+
0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52, // IHDR chunk
31+
0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04,
32+
0x08, 0x02, 0x00, 0x00, 0x00, 0x26, 0x93, 0x09,
33+
0x29, 0x00, 0x00, 0x00, 0x1c, 0x49, 0x44, 0x41, // IDAT chunk
34+
0x54, 0x78, 0x9c, 0x62, 0xf8, 0xcf, 0xc0, 0xc0,
35+
0xc0, 0xc0, 0xc0, 0xc0, 0xc0, 0xc0, 0xc0, 0xc0,
36+
0xc0, 0xc0, 0xc0, 0xc0, 0xc0, 0xc0, 0x00, 0x00,
37+
0x00, 0x31, 0x00, 0x01, 0x2e, 0xa8, 0xd1, 0xe5,
38+
0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, // IEND chunk
39+
0xae, 0x42, 0x60, 0x82,
40+
}
41+
b64 := base64.StdEncoding.EncodeToString(pngBytes)
42+
43+
outPath := processImageFile(b64, tmpDir)
44+
Expect(outPath).ToNot(BeEmpty(), "processImageFile should return a file path")
45+
46+
written, err := os.ReadFile(outPath)
47+
Expect(err).ToNot(HaveOccurred())
48+
Expect(written).To(Equal(pngBytes), "file on disk must match the original bytes")
49+
})
50+
51+
It("should return empty string for invalid base64", func() {
52+
outPath := processImageFile("not-valid-base64!!!", tmpDir)
53+
Expect(outPath).To(BeEmpty(), "should return empty string for invalid base64")
54+
})
55+
})
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package openai
2+
3+
import (
4+
"testing"
5+
6+
. "github.com/onsi/ginkgo/v2"
7+
. "github.com/onsi/gomega"
8+
)
9+
10+
func TestOpenAI(t *testing.T) {
11+
RegisterFailHandler(Fail)
12+
RunSpecs(t, "OpenAI Endpoints Suite")
13+
}

0 commit comments

Comments
 (0)