Skip to content

Leonardo image generation driver

LeonardoImageGenerationDriver

Bases: BaseImageGenerationDriver

Driver for the Leonardo image generation API.

Details on Leonardo image generation parameters can be found here: https://docs.leonardo.ai/reference/creategeneration

Attributes:

Name Type Description
model

The ID of the model to use when generating images.

api_key str

The API key to use when making requests to the Leonardo API.

requests_session Session

The requests session to use when making requests to the Leonardo API.

api_base str

The base URL of the Leonardo API.

max_attempts int

The maximum number of times to poll the Leonardo API for a completed image.

image_width int

The width of the generated image in the range [32, 1024] and divisible by 8.

image_height int

The height of the generated image in the range [32, 1024] and divisible by 8.

steps Optional[int]

Optionally specify the number of inference steps to run for each image generation request, [30, 60].

seed Optional[int]

Optionally provide a consistent seed to generation requests, increasing consistency in output.

Source code in griptape/griptape/drivers/image_generation/leonardo_image_generation_driver.py
@define
class LeonardoImageGenerationDriver(BaseImageGenerationDriver):
    """Driver for the Leonardo image generation API.

    Details on Leonardo image generation parameters can be found here:
    https://docs.leonardo.ai/reference/creategeneration

    Attributes:
        model: The ID of the model to use when generating images.
        api_key: The API key to use when making requests to the Leonardo API.
        requests_session: The requests session to use when making requests to the Leonardo API.
        api_base: The base URL of the Leonardo API.
        max_attempts: The maximum number of times to poll the Leonardo API for a completed image.
        image_width: The width of the generated image in the range [32, 1024] and divisible by 8.
        image_height: The height of the generated image in the range [32, 1024] and divisible by 8.
        steps: Optionally specify the number of inference steps to run for each image generation request, [30, 60].
        seed: Optionally provide a consistent seed to generation requests, increasing consistency in output.
    """

    api_key: str = field(kw_only=True)
    requests_session: requests.Session = field(default=Factory(lambda: requests.Session()), kw_only=True)
    api_base: str = "https://cloud.leonardo.ai/api/rest/v1"
    max_attempts: int = field(default=10, kw_only=True)
    image_width: int = field(default=512, kw_only=True)
    image_height: int = field(default=512, kw_only=True)
    steps: Optional[int] = field(default=None, kw_only=True)
    seed: Optional[int] = field(default=None, kw_only=True)

    def try_text_to_image(self, prompts: list[str], negative_prompts: Optional[list[str]] = None) -> ImageArtifact:
        if negative_prompts is None:
            negative_prompts = []

        prompt = ", ".join(prompts)
        negative_prompt = ", ".join(negative_prompts)

        generation_id = self._create_generation(prompt=prompt, negative_prompt=negative_prompt)
        image_url = self._get_image_url(generation_id=generation_id)
        image_data = self._download_image(url=image_url)

        return ImageArtifact(
            value=image_data,
            mime_type="image/png",
            width=self.image_width,
            height=self.image_height,
            model=self.model,
            prompt=prompt,
        )

    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 variation")

    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")

    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 _create_generation(self, prompt: str, negative_prompt: str) -> str:
        request = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": self.image_width,
            "height": self.image_height,
            "num_images": 1,
            "modelId": self.model,
        }

        if self.steps:
            request["num_inference_steps"] = self.steps

        if self.seed is not None:
            request["seed"] = self.seed

        response = self.requests_session.post(
            url=f"{self.api_base}/generations", headers={"Authorization": f"Bearer {self.api_key}"}, json=request
        ).json()

        return response["sdGenerationJob"]["generationId"]

    def _get_image_url(self, generation_id: str) -> str:
        for attempt in range(self.max_attempts):
            response = self.requests_session.get(
                url=f"{self.api_base}/generations/{generation_id}", headers={"Authorization": f"Bearer {self.api_key}"}
            ).json()

            if response["generations_by_pk"]["status"] == "PENDING":
                time.sleep(attempt + 1)
                continue

            return response["generations_by_pk"]["generated_images"][0]["url"]
        else:
            raise Exception("image generation failed to complete")

    def _download_image(self, url: str) -> bytes:
        response = self.requests_session.get(url=url, headers={"Authorization": f"Bearer {self.api_key}"})

        return response.content

api_base: str = 'https://cloud.leonardo.ai/api/rest/v1' class-attribute instance-attribute

api_key: str = field(kw_only=True) class-attribute instance-attribute

image_height: int = field(default=512, kw_only=True) class-attribute instance-attribute

image_width: int = field(default=512, kw_only=True) class-attribute instance-attribute

max_attempts: int = field(default=10, kw_only=True) class-attribute instance-attribute

requests_session: requests.Session = field(default=Factory(lambda : requests.Session()), kw_only=True) class-attribute instance-attribute

seed: Optional[int] = field(default=None, kw_only=True) class-attribute instance-attribute

steps: Optional[int] = field(default=None, kw_only=True) class-attribute instance-attribute

try_image_inpainting(prompts, image, mask, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/leonardo_image_generation_driver.py
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")

try_image_outpainting(prompts, image, mask, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/leonardo_image_generation_driver.py
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")

try_image_variation(prompts, image, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/leonardo_image_generation_driver.py
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 variation")

try_text_to_image(prompts, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/leonardo_image_generation_driver.py
def try_text_to_image(self, prompts: list[str], negative_prompts: Optional[list[str]] = None) -> ImageArtifact:
    if negative_prompts is None:
        negative_prompts = []

    prompt = ", ".join(prompts)
    negative_prompt = ", ".join(negative_prompts)

    generation_id = self._create_generation(prompt=prompt, negative_prompt=negative_prompt)
    image_url = self._get_image_url(generation_id=generation_id)
    image_data = self._download_image(url=image_url)

    return ImageArtifact(
        value=image_data,
        mime_type="image/png",
        width=self.image_width,
        height=self.image_height,
        model=self.model,
        prompt=prompt,
    )