Nova Canvas complete image generation tasks (#9177) (#9525)

* Nova Canvas complete image generation tasks (#9177) * add initial support for Amazon Nova Canvas model Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * adjust name to AmazonNovaCanvas and map function variables to config Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * tighten model name check Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * fix quality mapping Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * add premium quality in config Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * support all Amazon Nova Canvas tasks * remove unused import Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * add tests for image generation tasks and fix payload Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * add missing util file Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * update model prices backup file Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * remove image tasks other than text->image Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * add color guided generation task for Nova Canvas Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * fix merge Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * add nova canvas image generation documentation Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> * add nova canvas unit tests Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> --------- Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com> * ci(config.yml): bump ci config * test: fix test --------- Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com> Co-authored-by: omrishiv <327609+omrishiv@users.noreply.github.com>
2025-04-25 02:34:29 +00:00 · 2025-03-26 11:28:20 -07:00 · 2025-03-26 11:28:20 -07:00 · 801ecb6517
commit 801ecb6517
parent 8c845847cd
6 changed files with 216 additions and 2 deletions
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@ -640,6 +640,7 @@ jobs:

    steps:
      - checkout
+      - setup_google_dns
      - run:
          name: Install Dependencies
          command: |
--- a/docs/my-website/docs/providers/bedrock.md
+++ b/docs/my-website/docs/providers/bedrock.md
@ -1917,12 +1917,49 @@ model_list:

 </Tabs>

+Text to Image : 
+```bash
+curl -L -X POST 'http://0.0.0.0:4000/v1/images/generations' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer $LITELLM_VIRTUAL_KEY' \
+-d '{
+    "model": "amazon.nova-canvas-v1:0",
+    "prompt": "A cute baby sea otter"
+}'
+```

+Color Guided Generation:
+```bash
+curl -L -X POST 'http://0.0.0.0:4000/v1/images/generations' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer $LITELLM_VIRTUAL_KEY' \
+-d '{
+    "model": "amazon.nova-canvas-v1:0",
+    "prompt": "A cute baby sea otter",
+    "taskType": "COLOR_GUIDED_GENERATION",
+    "colorGuidedGenerationParams":{"colors":["#FFFFFF"]}
+}'
+```
+</TabItem>
+</Tabs>
+
+
+| Model Name              | Function Call                               |
+|-------------------------|---------------------------------------------|
+| Stable Diffusion 3 - v0 | `image_generation(model="bedrock/stability.stability.sd3-large-v1:0", prompt=prompt)` |
+| Stable Diffusion - v0   | `image_generation(model="bedrock/stability.stable-diffusion-xl-v0", prompt=prompt)` |
+| Stable Diffusion - v1   | `image_generation(model="bedrock/stability.stable-diffusion-xl-v1", prompt=prompt)` |
+| Amazon Nova Canvas - v0 | `image_generation(model="bedrock/amazon.nova-canvas-v1:0", prompt=prompt)` |
+  
+  
 ### Passing an external BedrockRuntime.Client as a parameter - Completion()
+  
+This is a deprecated flow. Boto3 is not async. And boto3.client does not let us make the http call through httpx. Pass in your aws params through the method above 👆. [See Auth Code](https://github.com/BerriAI/litellm/blob/55a20c7cce99a93d36a82bf3ae90ba3baf9a7f89/litellm/llms/bedrock_httpx.py#L284) [Add new auth flow](https://github.com/BerriAI/litellm/issues)

 :::warning

-This is a deprecated flow. Boto3 is not async. And boto3.client does not let us make the http call through httpx. Pass in your aws params through the method above 👆. [See Auth Code](https://github.com/BerriAI/litellm/blob/55a20c7cce99a93d36a82bf3ae90ba3baf9a7f89/litellm/llms/bedrock_httpx.py#L284) [Add new auth flow](https://github.com/BerriAI/litellm/issues)
+
+


 Experimental - 2024-Jun-23:
--- a/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py
+++ b/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py
@ -5,7 +5,8 @@ from openai.types.image import Image

 from litellm.types.llms.bedrock import (
    AmazonNovaCanvasTextToImageRequest, AmazonNovaCanvasTextToImageResponse,
-    AmazonNovaCanvasTextToImageParams, AmazonNovaCanvasRequestBase,
+    AmazonNovaCanvasTextToImageParams, AmazonNovaCanvasRequestBase, AmazonNovaCanvasColorGuidedGenerationParams,
+    AmazonNovaCanvasColorGuidedRequest,
 )
 from litellm.types.utils import ImageResponse

@ -69,6 +70,13 @@ class AmazonNovaCanvasConfig:
            text_to_image_params = AmazonNovaCanvasTextToImageParams(**text_to_image_params)
            return AmazonNovaCanvasTextToImageRequest(textToImageParams=text_to_image_params, taskType=task_type,
                                                      imageGenerationConfig=image_generation_config)
+        if task_type == "COLOR_GUIDED_GENERATION":
+            color_guided_generation_params = image_generation_config.pop("colorGuidedGenerationParams", {})
+            color_guided_generation_params = {"text": text, **color_guided_generation_params}
+            color_guided_generation_params = AmazonNovaCanvasColorGuidedGenerationParams(**color_guided_generation_params)
+            return AmazonNovaCanvasColorGuidedRequest(taskType=task_type,
+                                                      colorGuidedGenerationParams=color_guided_generation_params,
+                                                      imageGenerationConfig=image_generation_config)
        raise NotImplementedError(f"Task type {task_type} is not supported")

    @classmethod
--- a/litellm/types/llms/bedrock.py
+++ b/litellm/types/llms/bedrock.py
@ -418,6 +418,29 @@ class AmazonNovaCanvasTextToImageRequest(
    imageGenerationConfig: AmazonNovaCanvasImageGenerationConfig


+class AmazonNovaCanvasColorGuidedGenerationParams(TypedDict, total=False):
+    """
+    Params for Amazon Nova Canvas Color Guided Generation API
+    """
+
+    colors: List[str]
+    referenceImage: str
+    text: str
+    negativeText: str
+
+
+class AmazonNovaCanvasColorGuidedRequest(AmazonNovaCanvasRequestBase, TypedDict, total=False):
+    """
+    Request for Amazon Nova Canvas Color Guided Generation API
+
+    Ref: https://docs.aws.amazon.com/nova/latest/userguide/image-gen-req-resp-structure.html
+    """
+
+    taskType: Literal["COLOR_GUIDED_GENERATION"]
+    colorGuidedGenerationParams: AmazonNovaCanvasColorGuidedGenerationParams
+    imageGenerationConfig: AmazonNovaCanvasImageGenerationConfig
+
+
 class AmazonNovaCanvasTextToImageResponse(TypedDict, total=False):
    """
    Response for Amazon Nova Canvas Text to Image API
--- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py
+++ b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py
@ -6,6 +6,14 @@ import traceback
 from dotenv import load_dotenv
 from openai.types.image import Image

+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+from litellm.llms.bedrock.image.amazon_nova_canvas_transformation import (
+    AmazonNovaCanvasConfig,
+)
+
 logging.basicConfig(level=logging.DEBUG)
 load_dotenv()
 import asyncio
@ -54,6 +62,25 @@ def test_is_stability_3_model(model, expected):
    assert result == expected


+@pytest.mark.parametrize(
+    "model,expected",
+    [
+        ("amazon.nova-canvas", True),
+        ("sd3-large", False),
+        ("sd3-large-turbo", False),
+        ("sd3-medium", False),
+        ("sd3.5-large", False),
+        ("sd3.5-large-turbo", False),
+        ("gpt-4", False),
+        (None, False),
+        ("other-model", False),
+    ],
+)
+def test_is_nova_canvas_model(model, expected):
+    result = AmazonNovaCanvasConfig._is_nova_model(model)
+    assert result == expected
+
+
 def test_transform_request_body():
    prompt = "A beautiful sunset"
    optional_params = {"size": "1024x1024"}
@ -161,6 +188,110 @@ def test_get_request_body_stability():
    assert result["cfg_scale"] == 7


+def test_transform_request_body_nova_canvas():
+    prompt = "A beautiful sunset"
+    optional_params = {"size": "1024x1024"}
+
+    result = AmazonNovaCanvasConfig.transform_request_body(prompt, optional_params)
+
+    assert result["taskType"] == "TEXT_IMAGE"
+    assert result["textToImageParams"]["text"] == prompt
+    assert result["imageGenerationConfig"]["size"] == "1024x1024"
+
+
+def test_map_openai_params_nova_canvas():
+    non_default_params = {"n": 2, "size": "1024x1024"}
+    optional_params = {"cfg_scale": 7}
+
+    result = AmazonNovaCanvasConfig.map_openai_params(
+        non_default_params, optional_params
+    )
+
+    assert result == optional_params
+    assert "n" not in result  # OpenAI params should not be included
+
+
+def test_transform_response_dict_to_openai_response_nova_canvas():
+    # Create a mock response
+    response_dict = {"images": ["base64_encoded_image_1", "base64_encoded_image_2"]}
+    model_response = ImageResponse()
+
+    result = AmazonNovaCanvasConfig.transform_response_dict_to_openai_response(
+        model_response, response_dict
+    )
+
+    assert isinstance(result, ImageResponse)
+    assert len(result.data) == 2
+    assert all(hasattr(img, "b64_json") for img in result.data)
+    assert [img.b64_json for img in result.data] == response_dict["images"]
+
+
+def test_amazon_nova_canvas_get_supported_openai_params():
+    result = AmazonNovaCanvasConfig.get_supported_openai_params()
+    assert result == ["n", "size", "quality"]
+
+
+def test_get_request_body_nova_canvas_default():
+    handler = BedrockImageGeneration()
+    prompt = "A beautiful sunset"
+    optional_params = {"cfg_scale": 7}
+    model = "amazon.nova-canvas-v1"
+
+    result = handler._get_request_body(
+        model=model, prompt=prompt, optional_params=optional_params
+    )
+
+    assert result["taskType"] == "TEXT_IMAGE"
+    assert result["textToImageParams"]["text"] == prompt
+    assert result["imageGenerationConfig"]["cfg_scale"] == 7
+
+
+def test_get_request_body_nova_canvas_text_image():
+    handler = BedrockImageGeneration()
+    prompt = "A beautiful sunset"
+    optional_params = {"cfg_scale": 7, "taskType": "TEXT_IMAGE"}
+    model = "amazon.nova-canvas-v1"
+
+    result = handler._get_request_body(
+        model=model, prompt=prompt, optional_params=optional_params
+    )
+
+    assert result["taskType"] == "TEXT_IMAGE"
+    assert result["textToImageParams"]["text"] == prompt
+    assert result["imageGenerationConfig"]["cfg_scale"] == 7
+
+
+def test_get_request_body_nova_canvas_color_guided_generation():
+    handler = BedrockImageGeneration()
+    prompt = "A beautiful sunset"
+    optional_params = {
+        "cfg_scale": 7,
+        "taskType": "COLOR_GUIDED_GENERATION",
+        "colorGuidedGenerationParams": {"colors": ["#FF0000"]},
+    }
+    model = "amazon.nova-canvas-v1"
+
+    result = handler._get_request_body(
+        model=model, prompt=prompt, optional_params=optional_params
+    )
+
+    assert result["taskType"] == "COLOR_GUIDED_GENERATION"
+    assert result["colorGuidedGenerationParams"]["text"] == prompt
+    assert result["colorGuidedGenerationParams"]["colors"] == ["#FF0000"]
+    assert result["imageGenerationConfig"]["cfg_scale"] == 7
+
+
+def test_transform_request_body_with_invalid_task_type():
+    text = "An image of a otter"
+    optional_params = {"taskType": "INVALID_TASK"}
+
+    with pytest.raises(NotImplementedError) as exc_info:
+        AmazonNovaCanvasConfig.transform_request_body(
+            text=text, optional_params=optional_params
+        )
+    assert "Task type INVALID_TASK is not supported" in str(exc_info.value)
+
+
 def test_transform_response_dict_to_openai_response_stability3():
    handler = BedrockImageGeneration()
    model_response = ImageResponse()
--- a/tests/image_gen_tests/test_image_generation.py
+++ b/tests/image_gen_tests/test_image_generation.py
@ -143,6 +143,20 @@ class TestBedrockNovaCanvasTextToImage(BaseImageGenTest):
        }


+class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest):
+    def get_base_image_generation_call_args(self) -> dict:
+        litellm.in_memory_llm_clients_cache = InMemoryCache()
+        return {
+                "model": "bedrock/amazon.nova-canvas-v1:0",
+                "n": 1,
+                "size": "320x320",
+                "imageGenerationConfig": {"cfgScale":6.5,"seed":12},
+                "taskType": "COLOR_GUIDED_GENERATION",
+                "colorGuidedGenerationParams":{"colors":["#FFFFFF"]},
+                "aws_region_name": "us-east-1",
+        }
+
+
 class TestOpenAIDalle3(BaseImageGenTest):
    def get_base_image_generation_call_args(self) -> dict:
        return {"model": "dall-e-3"}