Skip to content

Openai

OpenAIImageGenerationToolSpec #

Bases: BaseToolSpec

OpenAI Image Generation tool spec.

Source code in llama-index-integrations/tools/llama-index-tools-openai/llama_index/tools/openai/image_generation/base.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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 io import BytesIO

            from PIL import Image
        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"

        return self.save_base64_image(image_bytes, filename)

image_generation #

image_generation(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.

Parameters:

Name Type Description Default
text str

The text to generate an image from.

 required
size str

The size of the image to generate (1024x1024, 256x256, 512x512).

 required
model str

The model to use to generate the image (dall-e-3, dall-e-2).

 'dall-e-3'
quality str

The quality of the image to generate (standard, hd).

 'standard'
num_images int

The number of images to generate.

 1
Source code in llama-index-integrations/tools/llama-index-tools-openai/llama_index/tools/openai/image_generation/base.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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"

    return self.save_base64_image(image_bytes, filename)