Skip to content

Base image generation driver

BaseImageGenerationDriver

Bases: ExponentialBackoffMixin, ABC

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
@define
class BaseImageGenerationDriver(ExponentialBackoffMixin, ABC):
    model: str = field(kw_only=True)
    structure: Structure | None = field(default=None, kw_only=True)

    def before_run(self, prompts: list[str], negative_prompts: list[str] | None = None) -> None:
        if self.structure:
            self.structure.publish_event(StartImageGenerationEvent(prompts=prompts, negative_prompts=negative_prompts))

    def after_run(self) -> None:
        if self.structure:
            self.structure.publish_event(FinishImageGenerationEvent())

    def run_text_to_image(self, prompts: list[str], negative_prompts: list[str] | None = None) -> ImageArtifact:
        for attempt in self.retrying():
            with attempt:
                self.before_run(prompts, negative_prompts)
                result = self.try_text_to_image(prompts, negative_prompts)
                self.after_run()

                return result

        else:
            raise Exception("Failed to run text to image generation")

    def run_image_variation(
        self, prompts: list[str], image: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        for attempt in self.retrying():
            with attempt:
                self.before_run(prompts, negative_prompts)
                result = self.try_image_variation(prompts, image, negative_prompts)
                self.after_run()

                return result

        else:
            raise Exception("Failed to generate image variations")

    def run_image_inpainting(
        self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        for attempt in self.retrying():
            with attempt:
                self.before_run(prompts, negative_prompts)
                result = self.try_image_inpainting(prompts, image, mask, negative_prompts)
                self.after_run()

                return result

        else:
            raise Exception("Failed to run image inpainting")

    def run_image_outpainting(
        self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        for attempt in self.retrying():
            with attempt:
                self.before_run(prompts, negative_prompts)
                result = self.try_image_outpainting(prompts, image, mask, negative_prompts)
                self.after_run()

                return result

        else:
            raise Exception("Failed to run image outpainting")

    @abstractmethod
    def try_text_to_image(self, prompts: list[str], negative_prompts: list[str] | None = None) -> ImageArtifact:
        ...

    @abstractmethod
    def try_image_variation(
        self, prompts: list[str], image: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        ...

    @abstractmethod
    def try_image_inpainting(
        self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        ...

    @abstractmethod
    def try_image_outpainting(
        self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        ...

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

structure: Structure | None = field(default=None, kw_only=True) class-attribute instance-attribute

after_run()

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
def after_run(self) -> None:
    if self.structure:
        self.structure.publish_event(FinishImageGenerationEvent())

before_run(prompts, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
def before_run(self, prompts: list[str], negative_prompts: list[str] | None = None) -> None:
    if self.structure:
        self.structure.publish_event(StartImageGenerationEvent(prompts=prompts, negative_prompts=negative_prompts))

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

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
def run_image_inpainting(
    self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
) -> ImageArtifact:
    for attempt in self.retrying():
        with attempt:
            self.before_run(prompts, negative_prompts)
            result = self.try_image_inpainting(prompts, image, mask, negative_prompts)
            self.after_run()

            return result

    else:
        raise Exception("Failed to run image inpainting")

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

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
def run_image_outpainting(
    self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
) -> ImageArtifact:
    for attempt in self.retrying():
        with attempt:
            self.before_run(prompts, negative_prompts)
            result = self.try_image_outpainting(prompts, image, mask, negative_prompts)
            self.after_run()

            return result

    else:
        raise Exception("Failed to run image outpainting")

run_image_variation(prompts, image, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
def run_image_variation(
    self, prompts: list[str], image: ImageArtifact, negative_prompts: list[str] | None = None
) -> ImageArtifact:
    for attempt in self.retrying():
        with attempt:
            self.before_run(prompts, negative_prompts)
            result = self.try_image_variation(prompts, image, negative_prompts)
            self.after_run()

            return result

    else:
        raise Exception("Failed to generate image variations")

run_text_to_image(prompts, negative_prompts=None)

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
def run_text_to_image(self, prompts: list[str], negative_prompts: list[str] | None = None) -> ImageArtifact:
    for attempt in self.retrying():
        with attempt:
            self.before_run(prompts, negative_prompts)
            result = self.try_text_to_image(prompts, negative_prompts)
            self.after_run()

            return result

    else:
        raise Exception("Failed to run text to image generation")

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

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
@abstractmethod
def try_image_inpainting(
    self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
) -> ImageArtifact:
    ...

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

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
@abstractmethod
def try_image_outpainting(
    self, prompts: list[str], image: ImageArtifact, mask: ImageArtifact, negative_prompts: list[str] | None = None
) -> ImageArtifact:
    ...

try_image_variation(prompts, image, negative_prompts=None) abstractmethod

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
@abstractmethod
def try_image_variation(
    self, prompts: list[str], image: ImageArtifact, negative_prompts: list[str] | None = None
) -> ImageArtifact:
    ...

try_text_to_image(prompts, negative_prompts=None) abstractmethod

Source code in griptape/griptape/drivers/image_generation/base_image_generation_driver.py
@abstractmethod
def try_text_to_image(self, prompts: list[str], negative_prompts: list[str] | None = None) -> ImageArtifact:
    ...