diff --git a/docs/my-website/docs/providers/custom_llm_server.md b/docs/my-website/docs/providers/custom_llm_server.md index 61099d1a35..4fcbf8942c 100644 --- a/docs/my-website/docs/providers/custom_llm_server.md +++ b/docs/my-website/docs/providers/custom_llm_server.md @@ -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!") ``` diff --git a/litellm/images/main.py b/litellm/images/main.py index 4aae96bf71..ca2d4e0b91 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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( diff --git a/litellm/llms/custom_llm.py b/litellm/llms/custom_llm.py index e88e8d5f1e..d235df30f2 100644 --- a/litellm/llms/custom_llm.py +++ b/litellm/llms/custom_llm.py @@ -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 diff --git a/tests/local_testing/test_custom_llm.py b/tests/local_testing/test_custom_llm.py index e61ede755e..d0f3292655 100644 --- a/tests/local_testing/test_custom_llm.py +++ b/tests/local_testing/test_custom_llm.py @@ -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):