diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 8178cad9038..d63221c2b22 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -232,6 +232,7 @@ async def image_edit_api( image_array: Optional[List[UploadFile]] = File(None, alias="image[]"), mask: Optional[List[UploadFile]] = File(None), mask_array: Optional[List[UploadFile]] = File(None, alias="mask[]"), + style_image: Optional[List[UploadFile]] = File(None), model: Optional[str] = None, ): """ @@ -282,10 +283,13 @@ async def image_edit_api( data = await _read_request_body(request=request) image_files = await batch_to_bytesio(image) mask_files = await batch_to_bytesio(mask) + style_image_files = await batch_to_bytesio(style_image) if image_files: data["image"] = image_files if mask_files: data["mask"] = mask_files + if style_image_files: + data["style_image"] = style_image_files for _field in ("image", "mask"): if _field in data and isinstance(data[_field], str): diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index 91a011a8234..05848968fa4 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -1,14 +1,18 @@ import asyncio import copy +import io +import tempfile from types import SimpleNamespace from typing import Any, Dict import orjson import pytest +from starlette.datastructures import UploadFile from starlette.requests import Request from starlette.responses import Response from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.image_endpoints import endpoints @@ -115,3 +119,67 @@ async def test_image_generation_prompt_rerouting(monkeypatch): assert captured_route_request_data["prompt"] == "sanitized prompt" assert "messages" not in captured_route_request_data assert response.headers.get("x-callback-test") == "value" + + +@pytest.mark.asyncio +async def test_image_edit_converts_style_image_upload_to_bytesio(monkeypatch): + """Bedrock Stability style-transfer needs a second image (``style_image``). + + It must be converted from a Starlette ``UploadFile`` to a synchronously + readable file-like object before routing; a raw ``UploadFile`` has an async + ``.read()`` that the downstream transform reads synchronously, stringifying + the resulting coroutine into the base64 payload sent to Bedrock. + """ + style_bytes = b"\x89PNG\r\n\x1a\nstyle-reference-bytes" + spool = tempfile.SpooledTemporaryFile() + spool.write(style_bytes) + spool.seek(0) + style_upload = UploadFile(file=spool, filename="style.png") + + captured: Dict[str, Any] = {} + + async def fake_base_process(self, **kwargs): + captured["data"] = self.data + return {"result": "ok"} + + monkeypatch.setattr( + ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + fake_base_process, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + + scope = { + "type": "http", + "method": "POST", + "path": "/v1/images/edits", + "headers": [(b"content-type", b"application/json")], + } + body = orjson.dumps( + {"prompt": "an apple", "model": "stability.stable-style-transfer"} + ) + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request(scope, receive) + + result = await endpoints.image_edit_api( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(), + image=None, + image_array=None, + mask=None, + mask_array=None, + style_image=[style_upload], + model="stability.stable-style-transfer", + ) + + assert result == {"result": "ok"} + routed_style_image = captured["data"]["style_image"] + assert isinstance(routed_style_image, list) + buffer = routed_style_image[0] + assert isinstance(buffer, io.BytesIO) + assert buffer.read() == style_bytes