Source code for langchain.tools.steamship_image_generation.tool
"""This tool allows agents to generate images using Steamship.
Steamship offers access to different third party image generation APIs
using a single API key.
Today the following models are supported:
- Dall-E
- Stable Diffusion
To use this tool, you must first set as environment variables:
    STEAMSHIP_API_KEY
```
"""
from __future__ import annotations
from enum import Enum
from typing import TYPE_CHECKING, Dict, Optional
from langchain.callbacks.manager import CallbackManagerForToolRun
from langchain.pydantic_v1 import root_validator
from langchain.tools import BaseTool
from langchain.tools.steamship_image_generation.utils import make_image_public
from langchain.utils import get_from_dict_or_env
if TYPE_CHECKING:
    from steamship import Steamship
[docs]class ModelName(str, Enum):
    """Supported Image Models for generation."""
    DALL_E = "dall-e"
    STABLE_DIFFUSION = "stable-diffusion"
SUPPORTED_IMAGE_SIZES = {
    ModelName.DALL_E: ("256x256", "512x512", "1024x1024"),
    ModelName.STABLE_DIFFUSION: ("512x512", "768x768"),
}
[docs]class SteamshipImageGenerationTool(BaseTool):
    """Tool used to generate images from a text-prompt."""
    model_name: ModelName
    size: Optional[str] = "512x512"
    steamship: Steamship
    return_urls: Optional[bool] = False
    name: str = "GenerateImage"
    description: str = (
        "Useful for when you need to generate an image."
        "Input: A detailed text-2-image prompt describing an image"