tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 06:26:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(custom_llm): add image_edit and aimage_edit support (#17999)

Browse Source 
* feat(custom_llm): add image_edit and aimage_edit support

Add support for image_edit and aimage_edit methods in CustomLLM class,
allowing users to implement custom image editing providers.

Changes:
- Add image_edit() and aimage_edit() methods to CustomLLM base class
- Add custom provider detection in litellm.image_edit() function
- Add tests for sync and async image_edit with custom handlers

* docs: add image_edit to CustomLLM documentation

- Add /v1/images/edits to supported routes
- Add Image Edit section with example
- Update Custom Handler Spec with image_edit methods
This commit is contained in:
Cesar Garcia
2025-12-15 14:16:36 -08:00
committed by GitHub
parent b57b1beb61
commit df19a747a2
4 changed files with 292 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/custom_llm_server.md
+108
View File
@@ -17,6 +17,7 @@ Supported Routes:
- `/v1/completions` -> `litellm.atext_completion`
- `/v1/embeddings` -> `litellm.aembedding`
- `/v1/images/generations` -> `litellm.aimage_generation`
- `/v1/images/edits` -> `litellm.aimage_edit`
- `/v1/messages` -> `litellm.acompletion`
@@ -263,6 +264,83 @@ Expected Response
}
```
## Image Edit
1. Setup your `custom_handler.py` file
```python
import litellm
from litellm import CustomLLM
from litellm.types.utils import ImageResponse, ImageObject
import time
class MyCustomLLM(CustomLLM):
async def aimage_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
# Your custom image edit logic here
# e.g., call Stability AI, Black Forest Labs, etc.
return ImageResponse(
created=int(time.time()),
data=[ImageObject(url="https://example.com/edited-image.png")],
)
my_custom_llm = MyCustomLLM()
```
2. Add to `config.yaml`
In the config below, we pass
python_filename: `custom_handler.py`
custom_handler_instance_name: `my_custom_llm`. This is defined in Step 1
custom_handler: `custom_handler.my_custom_llm`
```yaml
model_list:
- model_name: "my-custom-image-edit-model"
litellm_params:
model: "my-custom-llm/my-model"
litellm_settings:
custom_provider_map:
- {"provider": "my-custom-llm", "custom_handler": custom_handler.my_custom_llm}
```
```bash
litellm --config /path/to/config.yaml
```
3. Test it!
```bash
curl -X POST 'http://0.0.0.0:4000/v1/images/edits' \
-H 'Authorization: Bearer sk-1234' \
-F 'model=my-custom-image-edit-model' \
-F 'image=@/path/to/image.png' \
-F 'prompt=Make the sky blue'
```
Expected Response
```
{
"created": 1721955063,
"data": [{"url": "https://example.com/edited-image.png"}],
}
```
## Anthropic `/v1/messages`
- Write the integration for .acompletion
@@ -517,4 +595,34 @@ class CustomLLM(BaseLLM):
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
raise CustomLLMError(status_code=500, message="Not implemented yet!")
def image_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[HTTPHandler] = None,
) -> ImageResponse:
raise CustomLLMError(status_code=500, message="Not implemented yet!")
async def aimage_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
raise CustomLLMError(status_code=500, message="Not implemented yet!")
```
litellm/images/main.py
+53
View File
@@ -702,6 +702,59 @@ def image_edit(
custom_llm_provider=custom_llm_provider,
)
# Check for custom provider
if custom_llm_provider in litellm._custom_providers:
custom_handler: Optional[CustomLLM] = None
for item in litellm.custom_provider_map:
if item["provider"] == custom_llm_provider:
custom_handler = item["custom_handler"]
if custom_handler is None:
raise LiteLLMUnknownProvider(
model=model, custom_llm_provider=custom_llm_provider
)
model_response = ImageResponse()
if _is_async:
async_custom_client: Optional[AsyncHTTPHandler] = None
if kwargs.get("client") is not None and isinstance(
kwargs.get("client"), AsyncHTTPHandler
):
async_custom_client = kwargs.get("client")
return custom_handler.aimage_edit(
model=model,
image=images,
prompt=prompt,
model_response=model_response,
api_key=kwargs.get("api_key"),
api_base=kwargs.get("api_base"),
optional_params=kwargs,
logging_obj=litellm_logging_obj,
timeout=timeout,
client=async_custom_client,
)
else:
custom_client: Optional[HTTPHandler] = None
if kwargs.get("client") is not None and isinstance(
kwargs.get("client"), HTTPHandler
):
custom_client = kwargs.get("client")
return custom_handler.image_edit(
model=model,
image=images,
prompt=prompt,
model_response=model_response,
api_key=kwargs.get("api_key"),
api_base=kwargs.get("api_base"),
optional_params=kwargs,
logging_obj=litellm_logging_obj,
timeout=timeout,
client=custom_client,
)
# get provider config
image_edit_provider_config: Optional[BaseImageEditConfig] = (
ProviderConfigManager.get_provider_image_edit_config(
litellm/llms/custom_llm.py
+30
View File
@@ -197,6 +197,36 @@ class CustomLLM(BaseLLM):
) -> EmbeddingResponse:
raise CustomLLMError(status_code=500, message="Not implemented yet!")
def image_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[HTTPHandler] = None,
) -> ImageResponse:
raise CustomLLMError(status_code=500, message="Not implemented yet!")
async def aimage_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
raise CustomLLMError(status_code=500, message="Not implemented yet!")
def custom_chat_llm_router(
async_fn: bool, stream: Optional[bool], custom_llm: CustomLLM
tests/local_testing/test_custom_llm.py
+101
View File
@@ -309,6 +309,44 @@ class MyCustomLLM(CustomLLM):
return model_response
def image_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout=None,
client: Optional[HTTPHandler] = None,
) -> ImageResponse:
return ImageResponse(
created=int(time.time()),
data=[ImageObject(url="https://example.com/edited-image.png")],
response_ms=1000,
)
async def aimage_edit(
self,
model: str,
image: Any,
prompt: str,
model_response: ImageResponse,
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict,
logging_obj: Any,
timeout=None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
return ImageResponse(
created=int(time.time()),
data=[ImageObject(url="https://example.com/edited-image.png")],
response_ms=1000,
)
def test_get_llm_provider():
""""""
@@ -451,6 +489,69 @@ async def test_image_generation_async_additional_params():
}
def test_simple_image_edit():
"""Test sync image_edit with custom handler"""
my_custom_llm = MyCustomLLM()
litellm.custom_provider_map = [
{"provider": "custom_llm", "custom_handler": my_custom_llm}
]
resp = litellm.image_edit(
model="custom_llm/my-fake-model",
image=b"fake_image_bytes",
prompt="Edit this image",
)
print(resp)
assert resp.data[0].url == "https://example.com/edited-image.png"
@pytest.mark.asyncio
async def test_simple_image_edit_async():
"""Test async image_edit with custom handler"""
my_custom_llm = MyCustomLLM()
litellm.custom_provider_map = [
{"provider": "custom_llm", "custom_handler": my_custom_llm}
]
resp = await litellm.aimage_edit(
model="custom_llm/my-fake-model",
image=b"fake_image_bytes",
prompt="Edit this image",
)
print(resp)
assert resp.data[0].url == "https://example.com/edited-image.png"
@pytest.mark.asyncio
async def test_image_edit_async_additional_params():
"""Test that additional params are passed to custom handler"""
my_custom_llm = MyCustomLLM()
litellm.custom_provider_map = [
{"provider": "custom_llm", "custom_handler": my_custom_llm}
]
with patch.object(
my_custom_llm, "aimage_edit", new=AsyncMock(return_value=ImageResponse(
created=int(time.time()),
data=[ImageObject(url="https://example.com/edited-image.png")],
))
) as mock_client:
resp = await litellm.aimage_edit(
model="custom_llm/my-fake-model",
image=b"fake_image_bytes",
prompt="Edit this image",
api_key="my-api-key",
api_base="my-api-base",
my_custom_param="my-custom-param",
)
print(resp)
mock_client.assert_awaited_once()
assert mock_client.call_args.kwargs["api_key"] == "my-api-key"
assert mock_client.call_args.kwargs["api_base"] == "my-api-base"
def test_get_supported_openai_params():
class MyCustomLLM(CustomLLM):