This repository has been archived by the owner on Mar 1, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 737
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: OpenAI Image Generation Tool (#628)
* feat: dall-e-3 * chore: remove checkpoints * lint * cr * chore: use multi-modal as an example * chore: delete checkpoint * chore: fix tests and lint * cr * lint
- v0.0.79
- v0.0.79.post1
- v0.0.78
- v0.0.77
- v0.0.76
- v0.0.75
- v0.0.74
- v0.0.73
- v0.0.72
- v0.0.71
- v0.0.70
- v0.0.69
- v0.0.68
- v0.0.68.post1
- v0.0.67
- v0.0.66
- v0.0.65
- v0.0.64
- v0.0.63
- v0.0.62
- v0.0.61
- v0.0.60
- v0.0.59
- v0.0.58
- v0.0.58.post1
- v0.0.57
- v0.0.56
- v0.0.56.post1
- v0.0.55
- v0.0.55.post1
- v0.0.54
- v0.0.53
- v0.0.52
- v0.0.51
- v0.0.50
- v0.0.49
- v0.0.48
- v0.0.47
- v0.0.47.post1
- v0.0.46
1 parent
c128839
commit a96e996
Showing
13 changed files
with
426 additions
and
7 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -8,3 +8,4 @@ | |
.idea/ | ||
llama-hub.iml | ||
llamahub/ | ||
img_cache/ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
138 changes: 138 additions & 0 deletions
138
llama_hub/tools/notebooks/multimodal_openai_image.ipynb
Large diffs are not rendered by default.
Oops, something went wrong.
125 changes: 125 additions & 0 deletions
125
llama_hub/tools/notebooks/openai_image_generation_agent.ipynb
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,45 @@ | ||
# OpenAI Image Generation Tool | ||
|
||
This tool allows Agents to generate images using OpenAI's DALL-E model. To see more and get started, visit https://openai.com/blog/dall-e/ | ||
|
||
## Usage | ||
|
||
This tool has a more extensive example usage documented in a Jupyter notebook [here](https://github.com/emptycrown/llama-hub/tree/main/llama_hub/tools/notebooks/openai_image_generation.ipynb) | ||
|
||
### Usage with Agent | ||
```python | ||
from llama_hub.tools.openai.image_generation import OpenAIImageGenerationToolSpec | ||
|
||
image_generation_tool = OpenAIImageGenerationToolSpec(api_key=os.environ["OPENAI_API_KEY"]) | ||
|
||
agent = OpenAIAgent.from_tools( | ||
[*image_generation_tool.to_tool_list()], | ||
verbose=True, | ||
) | ||
|
||
response = agent.query('A pink and blue llama in a black background with the output') | ||
|
||
print(response) | ||
``` | ||
|
||
### Usage directly | ||
```python | ||
from llama_hub.tools.openai.image_generation import OpenAIImageGenerationToolSpec | ||
|
||
image_generation_tool = OpenAIImageGenerationToolSpec(api_key=os.environ["OPENAI_API_KEY"]) | ||
|
||
image_data = image_generation_tool.image_generation( | ||
text="A pink and blue llama with a black background", | ||
response_format="b64_json" | ||
) | ||
|
||
image_bytes = base64.b64decode(image_data) | ||
|
||
img = Image.open(BytesIO(image_bytes)) | ||
|
||
display(img) | ||
``` | ||
|
||
`image_generation`: Takes an text input and generates an image | ||
|
||
This loader is designed to be used as a way to load data as a Tool in a Agent. See [here](https://github.com/emptycrown/llama-hub/tree/main) for examples. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
## init file | ||
from llama_hub.tools.openai.image_generation.base import ( | ||
OpenAIImageGenerationToolSpec, | ||
) | ||
|
||
__all__ = ["OpenAIImageGenerationToolSpec"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,90 @@ | ||
"""OpenAI Image Generation tool sppec..""" | ||
|
||
import os | ||
import base64 | ||
import time | ||
|
||
from typing import Optional | ||
from llama_index.tools.tool_spec.base import BaseToolSpec | ||
|
||
DEFAULT_CACHE_DIR = "../../../img_cache" | ||
DEFAULT_SIZE = "1024x1024" # Dall-e-3 only supports 1024x1024 | ||
|
||
|
||
class OpenAIImageGenerationToolSpec(BaseToolSpec): | ||
"""OpenAI Image Generation tool spec.""" | ||
|
||
spec_functions = ["image_generation"] | ||
|
||
def __init__(self, api_key: str, cache_dir: Optional[str] = None) -> None: | ||
try: | ||
from openai import OpenAI | ||
except ImportError: | ||
raise ImportError( | ||
"Please install openai with `pip install openai` to use this tool" | ||
) | ||
|
||
"""Initialize with parameters.""" | ||
self.client = OpenAI(api_key=api_key) | ||
self.cache_dir = cache_dir or DEFAULT_CACHE_DIR | ||
|
||
def get_cache_dir(self): | ||
return self.cache_dir | ||
|
||
def save_base64_image(self, base64_str, image_name): | ||
try: | ||
from PIL import Image | ||
from io import BytesIO | ||
except ImportError: | ||
raise ImportError( | ||
"Please install Pillow with `pip install Pillow` to use this tool" | ||
) | ||
cache_dir = self.cache_dir | ||
|
||
# Create cache directory if it doesn't exist | ||
if not os.path.exists(cache_dir): | ||
os.makedirs(cache_dir) | ||
|
||
# Decode the base64 string | ||
image_data = base64.b64decode(base64_str) | ||
|
||
# Create an image from the decoded bytes and save it | ||
image_path = os.path.join(cache_dir, image_name) | ||
with Image.open(BytesIO(image_data)) as img: | ||
img.save(image_path) | ||
|
||
return image_path | ||
|
||
def image_generation( | ||
self, | ||
text: str, | ||
model: Optional[str] = "dall-e-3", | ||
quality: Optional[str] = "standard", | ||
num_images: Optional[int] = 1, | ||
) -> str: | ||
""" | ||
This tool accepts a natural language string and will use OpenAI's DALL-E model to generate an image. | ||
args: | ||
text (str): The text to generate an image from. | ||
size (str): The size of the image to generate (1024x1024, 256x256, 512x512). | ||
model (str): The model to use to generate the image (dall-e-3, dall-e-2). | ||
quality (str): The quality of the image to generate (standard, hd). | ||
num_images (int): The number of images to generate. | ||
""" | ||
response = self.client.images.generate( | ||
model=model, | ||
prompt=text, | ||
size=DEFAULT_SIZE, | ||
quality=quality, | ||
n=num_images, | ||
response_format="b64_json", | ||
) | ||
|
||
image_bytes = response.data[0].b64_json | ||
|
||
filename = f"{time.time()}.jpg" | ||
|
||
saved_image_path = self.save_base64_image(image_bytes, filename) | ||
|
||
return saved_image_path |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters