Stable Diffusion 3 local support

griptape/drivers/__init__.py

@@ -53,6 +53,18 @@
     BedrockStableDiffusionImageGenerationModelDriver,
 )
 from .image_generation_model.bedrock_titan_image_generation_model_driver import BedrockTitanImageGenerationModelDriver
+from .image_generation_model.stable_diffusion_3_pipeline_image_generation_model_driver import (
+    StableDiffusion3PipelineImageGenerationModelDriver,
+)
+from .image_generation_model.base_diffusion_pipeline_image_generation_model_driver import (
+    BaseDiffusionPipelineImageGenerationModelDriver,
+)
+from .image_generation_model.stable_diffusion_3_img_2_img_pipeline_image_generation_model_driver import (
+    StableDiffusion3Img2ImgPipelineImageGenerationModelDriver,
+)
+from .image_generation_model.stable_diffusion_3_controlnet_pipeline_image_generation_model_driver import (
+    StableDiffusion3ControlNetPipelineImageGenerationModelDriver,
+)
 
 from .image_generation.base_image_generation_driver import BaseImageGenerationDriver
 from .image_generation.base_multi_model_image_generation_driver import BaseMultiModelImageGenerationDriver
@@ -61,6 +73,9 @@
 from .image_generation.amazon_bedrock_image_generation_driver import AmazonBedrockImageGenerationDriver
 from .image_generation.azure_openai_image_generation_driver import AzureOpenAiImageGenerationDriver
 from .image_generation.dummy_image_generation_driver import DummyImageGenerationDriver
+from .image_generation.huggingface_diffusion_pipeline_image_generation_driver import (
+    HuggingFaceDiffusionPipelineImageGenerationDriver,
+)
 
 from .image_query_model.base_image_query_model_driver import BaseImageQueryModelDriver
 from .image_query_model.bedrock_claude_image_query_model_driver import BedrockClaudeImageQueryModelDriver
@@ -164,13 +179,18 @@
     "BaseImageGenerationModelDriver",
     "BedrockStableDiffusionImageGenerationModelDriver",
     "BedrockTitanImageGenerationModelDriver",
+    "BaseDiffusionPipelineImageGenerationModelDriver",
+    "StableDiffusion3PipelineImageGenerationModelDriver",
+    "StableDiffusion3Img2ImgPipelineImageGenerationModelDriver",
+    "StableDiffusion3ControlNetPipelineImageGenerationModelDriver",
     "BaseImageGenerationDriver",
     "BaseMultiModelImageGenerationDriver",
     "OpenAiImageGenerationDriver",
     "LeonardoImageGenerationDriver",
     "AmazonBedrockImageGenerationDriver",
     "AzureOpenAiImageGenerationDriver",
     "DummyImageGenerationDriver",
+    "HuggingFaceDiffusionPipelineImageGenerationDriver",
     "BaseImageQueryModelDriver",
     "BedrockClaudeImageQueryModelDriver",
     "BaseImageQueryDriver",

griptape/drivers/image_generation/huggingface_diffusion_pipeline_image_generation_driver.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import io
+from abc import ABC
+from typing import Optional
+
+from attrs import define, field
+
+from griptape.artifacts import ImageArtifact
+from griptape.drivers import BaseDiffusionPipelineImageGenerationModelDriver, BaseImageGenerationDriver
+from griptape.utils import import_optional_dependency
+
+
+@define
+class HuggingFaceDiffusionPipelineImageGenerationDriver(BaseImageGenerationDriver, ABC):
+    """Image generation driver for models hosted by Hugging Face's Diffusion Pipeline.
+
+    For more information, see the HuggingFace documentation for Diffusers:
+            https://huggingface.co/docs/diffusers/en/index
+
+    Attributes:
+        model_driver: A pipeline image generation model driver typed for the specific pipeline required by the model.
+        device: The hardware device used for inference. For example, "cpu", "cuda", or "mps".
+    """
+
+    model_driver: BaseDiffusionPipelineImageGenerationModelDriver = field(kw_only=True, metadata={"serializable": True})
+    device: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+
+    def try_text_to_image(self, prompts: list[str], negative_prompts: Optional[list[str]] = None) -> ImageArtifact:
+        pipeline = self.model_driver.prepare_pipeline(self.model, self.device)
+
+        prompt = ", ".join(prompts)
+        output_image = pipeline(
+            prompt, **self.model_driver.make_additional_params(negative_prompts, self.device)
+        ).images[0]
+
+        buffer = io.BytesIO()
+        output_image.save(buffer, format="PNG")
+
+        return ImageArtifact(
+            value=buffer.getvalue(), format="png", height=output_image.height, width=output_image.width, prompt=prompt
+        )
+
+    def try_image_variation(
+        self, prompts: list[str], image: ImageArtifact, negative_prompts: Optional[list[str]] = None
+    ) -> ImageArtifact:
+        pipeline = self.model_driver.prepare_pipeline(self.model, self.device)
+
+        prompt = ", ".join(prompts)
+        input_image = import_optional_dependency("PIL.Image").open(io.BytesIO(image.value))
+        # The size of the input image drives the size of the output image.
+        # Resize the input image to the configured dimensions.
+        requested_dimensions = self.model_driver.get_output_image_dimensions()
+        if requested_dimensions is not None and (
+            input_image.height != requested_dimensions[0] or input_image.width != requested_dimensions[1]
+        ):
+            input_image = input_image.resize(requested_dimensions)
+
+        output_image = pipeline(
+            prompt,
+            **self.model_driver.make_image_param(input_image),
+            **self.model_driver.make_additional_params(negative_prompts, self.device),
+        ).images[0]
+
+        buffer = io.BytesIO()
+        output_image.save(buffer, format="PNG")
+
+        return ImageArtifact(
+            value=buffer.getvalue(), format="png", height=output_image.height, width=output_image.width, prompt=prompt
+        )
+
+    def try_image_inpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: Optional[list[str]] = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError("Inpainting is not supported by this driver.")
+
+    def try_image_outpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: Optional[list[str]] = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError("Outpainting is not supported by this driver.")

griptape/drivers/image_generation_model/base_diffusion_pipeline_image_generation_model_driver.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional
+
+from attrs import define
+
+if TYPE_CHECKING:
+    from PIL.Image import Image
+
+
+@define
+class BaseDiffusionPipelineImageGenerationModelDriver(ABC):
+    @abstractmethod
+    def prepare_pipeline(self, model: str, device: Optional[str]) -> Any: ...
+
+    @abstractmethod
+    def make_image_param(self, image: Optional[Image]) -> Optional[dict[str, Image]]: ...
+
+    @abstractmethod
+    def make_additional_params(self, negative_prompts: Optional[list[str]], device: Optional[str]) -> dict: ...
+
+    @abstractmethod
+    def get_output_image_dimensions(self) -> Optional[tuple[int, int]]: ...

griptape/drivers/image_generation_model/stable_diffusion_3_controlnet_pipeline_image_generation_model_driver.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any, Optional
+
+from attrs import define, field
+
+from griptape.drivers import StableDiffusion3PipelineImageGenerationModelDriver
+from griptape.utils import import_optional_dependency
+
+if TYPE_CHECKING:
+    from PIL.Image import Image
+else:
+    StableDiffusion3ControlNetPipeline = import_optional_dependency(
+        "diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet"
+    ).StableDiffusion3ControlNetPipeline
+
+
+@define
+class StableDiffusion3ControlNetPipelineImageGenerationModelDriver(StableDiffusion3PipelineImageGenerationModelDriver):
+    """Image generation model driver for Stable Diffusion 3 models with ControlNet.
+
+    For more information, see the HuggingFace documentation for the StableDiffusion3ControlNetPipeline:
+        https://huggingface.co/docs/diffusers/en/api/pipelines/controlnet_sd3
+
+    Attributes:
+        controlnet_model: The ControlNet model to use for image generation.
+        controlnet_conditioning_scale: The conditioning scale for the ControlNet model. Defaults to None.
+    """
+
+    controlnet_model: str = field(kw_only=True)
+    controlnet_conditioning_scale: Optional[float] = field(default=None, kw_only=True, metadata={"serializable": True})
+
+    def prepare_pipeline(self, model: str, device: Optional[str]) -> Any:
+        pipeline_params = {}
+        controlnet_pipeline_params = {}
+        if self.torch_dtype is not None:
+            pipeline_params["torch_dtype"] = self.torch_dtype
+            controlnet_pipeline_params["torch_dtype"] = self.torch_dtype
+
+        # For both Stable Diffusion and ControlNet, models can be provided either
+        # as a path to a local file or as a HuggingFace model repo name.
+        # We use the from_single_file method if the model is a local file and the
+        # from_pretrained method if the model is a local directory or hosted on HuggingFace.
+        sd3_controlnet_model = import_optional_dependency("diffusers.models.controlnet_sd3").SD3ControlNetModel
+        if os.path.isfile(self.controlnet_model):
+            pipeline_params["controlnet"] = sd3_controlnet_model.from_single_file(
+                self.controlnet_model, **controlnet_pipeline_params
+            )
+
+        else:
+            pipeline_params["controlnet"] = sd3_controlnet_model.from_pretrained(
+                self.controlnet_model, **controlnet_pipeline_params
+            )
+
+        sd3_controlnet_pipeline = import_optional_dependency(
+            "diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet"
+        ).StableDiffusion3ControlNetPipeline
+        if os.path.isfile(model):
+            pipeline = sd3_controlnet_pipeline.from_single_file(model, **pipeline_params)
+
+        else:
+            pipeline = sd3_controlnet_pipeline.from_pretrained(model, **pipeline_params)
+
+        if device is not None:
+            pipeline.to(device)
+
+        return pipeline
+
+    def make_image_param(self, image: Optional[Image]) -> Optional[dict[str, Image]]:
+        if image is None:
+            raise ValueError("Input image is required for ControlNet pipelines.")
+
+        return {"control_image": image}
+
+    def make_additional_params(self, negative_prompts: Optional[list[str]], device: Optional[str]) -> dict[str, Any]:
+        additional_params = super().make_additional_params(negative_prompts, device)
+
+        del additional_params["height"]
+        del additional_params["width"]
+
+        if self.controlnet_conditioning_scale is not None:
+            additional_params["controlnet_conditioning_scale"] = self.controlnet_conditioning_scale
+
+        return additional_params

griptape/drivers/image_generation_model/stable_diffusion_3_img_2_img_pipeline_image_generation_model_driver.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any, Optional
+
+from attrs import define, field
+
+from griptape.drivers import StableDiffusion3PipelineImageGenerationModelDriver
+from griptape.utils import import_optional_dependency
+
+if TYPE_CHECKING:
+    from PIL.Image import Image
+
+
+@define
+class StableDiffusion3Img2ImgPipelineImageGenerationModelDriver(StableDiffusion3PipelineImageGenerationModelDriver):
+    """Image generation model driver for Stable Diffusion 3 model image to image pipelines.
+
+    For more information, see the HuggingFace documentation for the StableDiffusion3Img2ImgPipeline:
+        https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
+
+    Attributes:
+        strength: A value [0.0, 1.0] that determines the strength of the initial image in the output.
+    """
+
+    strength: Optional[float] = field(default=None, kw_only=True, metadata={"serializable": True})
+
+    def prepare_pipeline(self, model: str, device: Optional[str]) -> Any:
+        pipeline_params = {}
+        if self.torch_dtype is not None:
+            pipeline_params["torch_dtype"] = self.torch_dtype
+
+        # A model can be provided either as a path to a local file
+        # or as a HuggingFace model repo name.
+        sd3_img2img_pipeline = import_optional_dependency(
+            "diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img"
+        ).StableDiffusion3Img2ImgPipeline
+        if os.path.isfile(model):
+            # If the model provided is a local file (not a directory),
+            # we load it using the from_single_file method.
+
+            raise NotImplementedError(
+                "StableDiffusion3Img2ImgPipeline does not yet support loading from a single file."
+            )
+        else:
+            # If the model is a local directory or hosted on HuggingFace,
+            # we load it using the from_pretrained method.
+            pipeline = sd3_img2img_pipeline.from_pretrained(model, **pipeline_params)
+
+        # Move inference to particular device if requested.
+        if device is not None:
+            pipeline.to(device)
+
+        return pipeline
+
+    def make_image_param(self, image: Optional[Image]) -> Optional[dict[str, Image]]:
+        if image is None:
+            raise ValueError("Input image is required for image to image pipelines.")
+
+        return {"image": image}
+
+    def make_additional_params(self, negative_prompts: Optional[list[str]], device: Optional[str]) -> dict[str, Any]:
+        additional_params = super().make_additional_params(negative_prompts, device)
+
+        # Explicit height and width params are not supported, but
+        # are instead inferred from input image.
+        del additional_params["height"]
+        del additional_params["width"]
+
+        if self.strength is not None:
+            additional_params["strength"] = self.strength
+
+        return additional_params