Add Black Forest Labs Image Generation Driver

README.md

@@ -1,142 +1,32 @@
-# Griptape Extension Template
-
-A Github template repository for creating Griptape extensions.
-
-## Getting Started
-
-Via github web page:
-
-Click on `Use this template`
-
-![](https://docs.github.com/assets/cb-36544/images/help/repository/use-this-template-button.png)
-
-
-Via `gh`:
-
-```
-$ gh repo create griptape-extension-name -p griptape/griptape-extension-template
-```
-
-## What is a Griptape Extension?
-
-Griptape Extensions can add new functionality to the [Griptape framework](https://github.com/griptape-ai/griptape), such as new Tools, Drivers, Tasks, or Structures.
-
-This repository provides a recommended structure for organizing your extension code, as well as helpful tools for testing and development.
-
-## Extension Structure
-
-The template repository is structured as follows:
-
-```bash
-tree -I __init__.py -I __pycache__
-
-├── griptape
-│   └── extension_name # Name whatever you want
-│       └── tools
-│           └── reverse_string
-│               └── tool.py
-            ...more directories for other interfaces (drivers, tasks, structures, etc)...
-└── tests
-    └── unit
-        └── tools
-            └── test_reverse_string_tool.py
-├── examples
-    └── tools
-        └── example_agent.py # Example usage of the extension 
-├── LICENSE # Choose the appropriate license
-├── Makefile # Contains useful commands for development
-├── pyproject.toml # Contains the project's metadata
-├── README.md #  Describes the extension and how to use it
-```
-
-## Development
-
-### Poetry
-
-This project uses [Poetry](https://python-poetry.org/) for dependency management.
-It is recommended to configure Poetry to use [in-project](https://python-poetry.org/docs/configuration/#virtualenvsin-project) virtual environments:
-
-```bash
-poetry config virtualenvs.in-project true
-```
-
-This will create a `.venv` directory in the project root, where the virtual environment will be stored.
-This ensures that the virtual environment is always in the same location, regardless of where the project is cloned.
-
-### Useful Commands
-
-#### Installing Dependencies
-
-```bash
-make install
-```
-
-#### Running Tests
-
-```bash
-make test
-```
-
-#### Running Checks (linting, formatting, etc)
-
-```bash
-make check
-```
-
-#### Running Formatter
-
-```bash
-make format
-```
-
-#### Running Example
-
-This template includes an [example](https://github.com/griptape-ai/tool-template/blob/main/examples/tools/example_agent.py) demonstrating how to use the extension. It shows how to import the `ReverseStringTool`, provide it to an Agent, and run it.
-
-1. Set the required environment variables. The example needs the `OPENAI_API_KEY` environment variable to be set.
-2. Run the example:
-
-```bash
-poetry run python examples/tools/example_agent.py
-```
-
-If successful, you should see:
+# Griptape Black Forest Extension
+
+This extension provides an [Image Generation Driver](https://docs.griptape.ai/stable/griptape-framework/drivers/image-generation-drivers/#amazon-bedrock) for [Black Forest Labs](https://docs.bfl.ml/quick_start/gen_image).
+
+For example:
+
+```python
+from griptape.black_forest.drivers.black_forest_image_generation_driver import (
+    BlackForestImageGenerationDriver,
+)
+from griptape.engines import PromptImageGenerationEngine
+from griptape.structures import Agent
+from griptape.tools import FileManagerTool, PromptImageGenerationTool
+
+agent = Agent(
+    tools=[
+        PromptImageGenerationTool(
+            engine=PromptImageGenerationEngine(
+                image_generation_driver=BlackForestImageGenerationDriver(
+                    model="flux-pro-1.1"
+                )
+            ),
+            off_prompt=True,
+        ),
+        FileManagerTool(),
+    ]
+)
+
+agent.run(
+    "Save a picture of a watercolor painting of a dog riding a skateboard to the desktop."
+)
 ```
-[11/18/24 14:55:14] INFO     ToolkitTask 6bb7fa5581d147b2a39e801631c98005
-                             Input: Use the ReverseStringTool to reverse 'Griptape'
-[11/18/24 14:55:15] INFO     Subtask c3036471831144529b8d5300c6849203
-                             Actions: [
-                               {
-                                 "tag": "call_VE4tGBFL7iB7VDbkKaIFIkwY",
-                                 "name": "ReverseStringTool",
-                                 "path": "reverse_string",
-                                 "input": {
-                                   "values": {
-                                     "input": "Griptape"
-                                   }
-                                 }
-                               }
-                             ]
-                    INFO     Subtask c3036471831144529b8d5300c6849203
-                             Response: epatpirG
-[11/18/24 14:55:16] INFO     ToolkitTask 6bb7fa5581d147b2a39e801631c98005
-                             Output: The reversed string of "Griptape" is "epatpirG".
-```
-
-### Installing in Other Projects
-
-Extensions are designed to be shared. Extensions are made to easily install into existing Python projects.
-
-The easiest way to include your extension into an existing project is to install directly from the repository, like so:
-```bash
-poetry add git+https://github.com/{your-org}/{your-extension-name}.git
-```
-
-To install a local copy of the extension for development, run:
-```bash
-poetry add -e /path/to/your/extension
-```
-
-Any changes made to the extension will be automatically reflected in the project without needing to reinstall it.
-
-Advanced customers may seek to publish their extensions to PyPi. Those instructions are beyond the scope of this README.

examples/drivers/example_agent.py

@@ -0,0 +1,24 @@
+from griptape.black_forest.drivers.black_forest_image_generation_driver import (
+    BlackForestImageGenerationDriver,
+)
+from griptape.engines import PromptImageGenerationEngine
+from griptape.structures import Agent
+from griptape.tools import FileManagerTool, PromptImageGenerationTool
+
+agent = Agent(
+    tools=[
+        PromptImageGenerationTool(
+            engine=PromptImageGenerationEngine(
+                image_generation_driver=BlackForestImageGenerationDriver(
+                    model="flux-pro-1.1"
+                )
+            ),
+            off_prompt=True,
+        ),
+        FileManagerTool(),
+    ]
+)
+
+agent.run(
+    "Save a picture of a watercolor painting of a dog riding a skateboard to the desktop."
+)

griptape/black_forest/drivers/black_forest_image_generation_driver.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+import os
+import time
+
+from urllib.parse import urljoin
+from griptape.artifacts import ImageArtifact
+import requests
+
+from attrs import define, field, Factory
+
+from griptape.drivers import BaseImageGenerationDriver
+
+
+@define
+class BlackForestImageGenerationDriver(BaseImageGenerationDriver):
+    base_url: str = field(
+        default="https://api.bfl.ml",
+        kw_only=True,
+        metadata={"serializable": False},
+    )
+    api_key: str | None = field(
+        default=Factory(lambda: os.environ["BFL_API_KEY"]),
+        kw_only=True,
+        metadata={"serializable": False},
+    )
+    width: int = field(default=1024, kw_only=True)
+    height: int = field(default=768, kw_only=True)
+    sleep_interval: float = field(default=0.5, kw_only=True)
+
+    def try_text_to_image(
+        self, prompts: list[str], negative_prompts: list[str] | None = None
+    ) -> ImageArtifact:
+        prompt = " ".join(prompts)
+
+        request = requests.post(
+            urljoin(self.base_url, f"v1/{self.model}"),
+            headers={
+                "accept": "application/json",
+                "x-key": self.api_key,
+                "Content-Type": "application/json",
+            },
+            json={
+                "prompt": prompt,
+                "width": self.width,
+                "height": self.height,
+            },
+        ).json()
+
+        request_id = request["id"]
+
+        image_url = None
+        while True:
+            time.sleep(self.sleep_interval)
+            result = requests.get(
+                urljoin(self.base_url, "v1/get_result"),
+                headers={
+                    "accept": "application/json",
+                    "x-key": self.api_key,
+                },
+                params={
+                    "id": request_id,
+                },
+            ).json()
+            if result["status"] == "Ready":
+                image_url = result["result"]["sample"]
+                break
+
+        image_response = requests.get(image_url)
+        image_bytes = image_response.content
+
+        return ImageArtifact(
+            value=image_bytes, format="jpeg", width=self.width, height=self.height
+        )
+
+    def try_image_variation(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        negative_prompts: list[str] | None = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support variation"
+        )
+
+    def try_image_inpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: list[str] | None = 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: list[str] | None = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support outpainting"
+        )

pyproject.toml

@@ -1,8 +1,8 @@
 [tool.poetry]
-name = "griptape_plugin_template"
+name = "griptape_black_forest"
 version = "0.1.0"
-description = "A Griptape plugin."
-authors = ["Author Name <[email protected]>"]
+description = "Griptape extension for Black Forest Labs"
+authors = ["Collin Dutter <[email protected]>"]
 readme = "README.md"
 packages = [
     {include = "griptape"}