feat(drivers): Add GriptapeCloudImageGenerationDriver

docs/griptape-framework/drivers/image-generation-drivers.md

@@ -23,6 +23,16 @@ Provide a Driver to a [Tool](../tools/index.md) for use by an [Agent](../structu
 
 ## Image Generation Drivers
 
+### Griptape Cloud
+
+The [Griptape Cloud Image Generation Driver](../../reference/griptape/drivers/image_generation/griptape_cloud_image_generation_driver.md) provides access to image generation models hosted by Griptape Cloud.
+
+Today, the only accessible model is `dall-e-3`.
+
+```python
+--8<-- "docs/griptape-framework/drivers/src/image_generation_drivers_griptape_cloud.py"
+```
+
 ### Amazon Bedrock
 
 The [Amazon Bedrock Image Generation Driver](../../reference/griptape/drivers/image_generation/amazon_bedrock_image_generation_driver.md) provides multi-model access to image generation models hosted by Amazon Bedrock. This Driver manages API calls to the Bedrock API, while the specific Model Drivers below format the API requests and parse the responses.

docs/griptape-framework/drivers/src/image_generation_drivers_griptape_cloud.py

@@ -0,0 +1,16 @@
+import os
+from io import BytesIO
+
+from PIL import Image
+
+from griptape.drivers.image_generation.griptape_cloud import GriptapeCloudImageGenerationDriver
+
+driver = GriptapeCloudImageGenerationDriver(
+    api_key=os.environ["GT_CLOUD_API_KEY"],
+    model="dall-e-3",
+)
+
+
+image = driver.run_text_to_image(["A capybara sitting on a rock in the sun."])
+
+Image.open(BytesIO(image.value)).show()

griptape/drivers/__init__.py

@@ -73,6 +73,7 @@
 from .image_generation.amazon_bedrock import AmazonBedrockImageGenerationDriver
 from .image_generation.dummy import DummyImageGenerationDriver
 from .image_generation.huggingface_pipeline import HuggingFacePipelineImageGenerationDriver
+from .image_generation.griptape_cloud import GriptapeCloudImageGenerationDriver
 
 from .web_scraper import BaseWebScraperDriver
 from .web_scraper.trafilatura import TrafilaturaWebScraperDriver
@@ -237,6 +238,24 @@
     "StableDiffusion3ControlNetImageGenerationPipelineDriver",
     "StableDiffusion3ImageGenerationPipelineDriver",
     "StableDiffusion3Img2ImgImageGenerationPipelineDriver",
+    "StableDiffusion3ControlNetImageGenerationPipelineDriver",
+    "BaseImageGenerationDriver",
+    "BaseMultiModelImageGenerationDriver",
+    "OpenAiImageGenerationDriver",
+    "LeonardoImageGenerationDriver",
+    "AmazonBedrockImageGenerationDriver",
+    "AzureOpenAiImageGenerationDriver",
+    "DummyImageGenerationDriver",
+    "HuggingFacePipelineImageGenerationDriver",
+    "GriptapeCloudImageGenerationDriver",
+    "BaseWebScraperDriver",
+    "TrafilaturaWebScraperDriver",
+    "MarkdownifyWebScraperDriver",
+    "ProxyWebScraperDriver",
+    "BaseWebSearchDriver",
+    "GoogleWebSearchDriver",
+    "DuckDuckGoWebSearchDriver",
+    "ExaWebSearchDriver",
     "TavilyWebSearchDriver",
     "TrafilaturaWebScraperDriver",
     "VoyageAiEmbeddingDriver",

griptape/drivers/image_generation/griptape_cloud/__init__.py

@@ -0,0 +1,7 @@
+from griptape.drivers.image_generation.griptape_cloud_image_generation_driver import (
+    GriptapeCloudImageGenerationDriver,
+)
+
+__all__ = [
+    "GriptapeCloudImageGenerationDriver",
+]

griptape/drivers/image_generation/griptape_cloud_image_generation_driver.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import os
+from typing import Literal, Optional
+from urllib.parse import urljoin
+
+import requests
+from attrs import Factory, define, field
+
+from griptape.artifacts import ImageArtifact
+from griptape.drivers.image_generation import BaseImageGenerationDriver
+
+
+@define
+class GriptapeCloudImageGenerationDriver(BaseImageGenerationDriver):
+    model: Optional[str] = field(default=None, kw_only=True)
+    base_url: str = field(
+        default=Factory(lambda: os.getenv("GT_CLOUD_BASE_URL", "https://cloud.griptape.ai")),
+    )
+    api_key: str = field(default=Factory(lambda: os.environ["GT_CLOUD_API_KEY"]))
+    headers: dict = field(
+        default=Factory(lambda self: {"Authorization": f"Bearer {self.api_key}"}, takes_self=True), kw_only=True
+    )
+    style: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    quality: Literal["standard", "hd"] = field(default="standard", kw_only=True, metadata={"serializable": True})
+    image_size: Literal["1024x1024", "1024x1792", "1792x1024"] = field(
+        default="1024x1024", kw_only=True, metadata={"serializable": True}
+    )
+
+    def try_text_to_image(self, prompts: list[str], negative_prompts: Optional[list[str]] = None) -> ImageArtifact:
+        url = urljoin(self.base_url.strip("/"), "/api/images/generations")
+
+        response = requests.post(
+            url,
+            headers=self.headers,
+            json={
+                "prompts": prompts,
+                "driver_configuration": {
+                    "model": self.model,
+                    "image_size": self.image_size,
+                    "quality": self.quality,
+                    "style": self.style,
+                },
+            },
+        )
+        response.raise_for_status()
+        response = response.json()
+
+        return ImageArtifact.from_dict(response["artifact"])
+
+    def try_image_variation(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        negative_prompts: Optional[list[str]] = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support image variation")
+
+    def try_image_inpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: Optional[list[str]] = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support inpainting")
+
+    def try_image_outpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: Optional[list[str]] = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support outpainting")

tests/unit/drivers/image_generation/test_griptape_cloud_image_generation_driver.py

@@ -0,0 +1,64 @@
+from unittest.mock import Mock
+
+import pytest
+
+from griptape.drivers.image_generation.griptape_cloud import GriptapeCloudImageGenerationDriver
+
+
+class TestGriptapeCloudImageGenerationDriver:
+    @pytest.fixture(autouse=True)
+    def mock_post(self, mocker):
+        def request(*args, **kwargs):
+            mock_response = mocker.Mock()
+            if "images/generations" in args[0]:
+                mock_response.json.return_value = {
+                    "artifact": {
+                        "type": "ImageArtifact",
+                        "width": 512,
+                        "height": 512,
+                        "format": "png",
+                        "value": "aW1hZ2UgZGF0YQ==",
+                        "meta": {"model": "dall-e-2", "prompt": "test prompt"},
+                    },
+                }
+                return mock_response
+            return mocker.Mock(
+                raise_for_status=lambda: None,
+            )
+
+        return mocker.patch("requests.post", side_effect=request)
+
+    @pytest.fixture()
+    def driver(self):
+        return GriptapeCloudImageGenerationDriver(model="dall-e-3", api_key="foo", quality="hd")
+
+    def test_init(self, driver):
+        assert driver
+
+    def test_try_text_to_image(self, driver):
+        image_artifact = driver.try_text_to_image(prompts=["test prompt"])
+
+        assert image_artifact.value == b"image data"
+        assert image_artifact.mime_type == "image/png"
+        assert image_artifact.width == 512
+        assert image_artifact.height == 512
+        assert image_artifact.meta["model"] == "dall-e-2"
+        assert image_artifact.meta["prompt"] == "test prompt"
+
+    def test_try_image_variation(self, driver):
+        with pytest.raises(NotImplementedError):
+            driver.try_image_variation(prompts=[], image=Mock(value=b"image data"))
+
+    def test_try_image_variation_invalid_size(self, driver):
+        with pytest.raises(NotImplementedError):
+            driver.try_image_variation(prompts=[], image=Mock(value=b"image data"))
+
+    def test_try_image_variation_invalid_model(self, driver):
+        with pytest.raises(NotImplementedError):
+            driver.try_image_variation(prompts=[], image=Mock(value=b"image data"))
+
+    def test_try_image_inpainting(self, driver):
+        with pytest.raises(NotImplementedError):
+            driver.try_image_inpainting(
+                prompts=["test prompt"], image=Mock(value=b"image data"), mask=Mock(value=b"mask data")
+            )