Skip to content

Module: stable_diffusion_workflow

Inference workflow for Stable Diffusion pipelines.

This module contains the implementation of the Stable Diffusion Inference Workflow. The workflow is responsible for setting up the Stable Diffusion pipeline and performing inference on the input data.

Additional Installation

To use the Stable Diffusion pipeline, some additional packages need to be installed. You can install the required packages using the following command:

uv pip install "infernet-ml[diffusion_inference]"

pip install "infernet-ml[diffusion_inference]"

Example Usage

from infernet_ml.workflows.inference.stable_diffusion_workflow import (
    StableDiffusionWorkflow,
    SupportedPipelines,
)

def main():
    # Initialize the Stable Diffusion Inference Workflow
    workflow = StableDiffusionWorkflow(
        pipeline=SupportedPipelines.STABLE_DIFFUSION,
        model="stabilityai/stable-diffusion-2-1",
    )

    # Setup the workflow
    workflow.setup()

    # Perform inference on the input data
    input_data = HFDiffusionInferenceInput(prompt="A photo of a cat")

    output_data = workflow.run_model(input_data)

StableDiffusionWorkflow

Bases: BaseInferenceWorkflow

Inference workflow for Stable Diffusion pipelines.

Source code in src/infernet_ml/workflows/inference/stable_diffusion_workflow.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
class StableDiffusionWorkflow(BaseInferenceWorkflow):
    """
    Inference workflow for Stable Diffusion pipelines.
    """

    def __init__(
        self,
        pipeline: SupportedPipelines,
        model: Optional[str] = None,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        """
        Initialize the Stable Diffusion Inference Workflow object

        Args:
            pipeline (SupportedPipelines): pipeline to be used for inference. Supported
                pipelines are "StableDiffusion"
            model (str): Stable Diffusion model to be used for inference.
                Defaults to "runwayml/stable-diffusion-v1-5"
            vae_type (str): Type of VAE to be used for inference:
                `mse`, `ema`. Defaults to "ema"
            torch_dtype (str): Type of torch dtype to be used for inference.
                Defaults to "float32"
            enable_xformers (bool): Enable xformers memory efficient attention.
                NOTE: requires `xformers` to be installed. Defaults to False
        Raises:
            ValueError: if pipeline is not supported
        """
        super().__init__(*args, **kwargs)
        self.pipeline_id = pipeline.value
        # Check to ensure that the pipeline is supported
        # Else raise a ValueError
        if self.pipeline_id not in SUPPORTED_DIFFUSION_PIPELINES:
            raise ValueError(
                f"Pipeline {pipeline} is not supported.  \
                Supported pipelines are {SUPPORTED_DIFFUSION_PIPELINES}"
            )
        self.model_id = model or "runwayml/stable-diffusion-v1-5"
        self.vae_type = kwargs.get("vae_type", VaeType.ema)
        self.torch_dtype = kwargs.get("torch_dtype", "float32")
        self.enable_xformers = kwargs.get("enable_xformers", False)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.pipeline_options = StableDiffusionPipelineOptions(
            model=self.model_id,
            vae_type=self.vae_type,
            torch_dtype=self.torch_dtype,
            enable_xformers=self.enable_xformers,
        )
        self.inference_params = {
            "model": model,
        }

    def create_pipeline(
        self, pipeline_id: str, pipeline_options: StableDiffusionPipelineOptions
    ) -> StableDiffusionPipeline:
        # Validate the pipeline options
        try:
            model = pipeline_options.model
            torch_dtype = get_torch_dtype(pipeline_options.torch_dtype)
            vae = get_vae_kl(
                VaeConfig(
                    type=pipeline_options.vae_type,
                    model=model,
                    torch_dtype=torch_dtype,
                )
            )
            pipeline = StableDiffusionPipeline.from_pretrained(  # type: ignore
                model, torch_dtype=torch_dtype, vae=vae
            )
            pipeline.to(self.device)
            return pipeline  # type: ignore
        except Exception as e:
            logger.error(f"Error creating pipeline: {e}")
            raise e

    def get_pipeline(
        self, pipeline_id: str, pipeline_options: StableDiffusionPipelineOptions
    ) -> StableDiffusionPipeline:
        pipeline = self.create_pipeline(pipeline_id, pipeline_options=pipeline_options)
        return pipeline

    def do_setup(self) -> bool:
        """
        Setup the inference client
        Returns:
            bool: True if setup is successful, False otherwise
        """
        done = False
        self.pipeline = self.get_pipeline(self.pipeline_id, self.pipeline_options)
        if self.pipeline_options.enable_xformers:
            self.pipeline.enable_xformers_memory_efficient_attention()  # type: ignore
        done = self.pipeline is not None
        logger.debug(f"Setup done: {done}")
        return done

    def do_preprocessing(self, input_data: HFDiffusionInferenceInput) -> dict[str, Any]:
        try:
            input_data_dict = input_data.model_dump()
        except ValidationError as e:
            raise ValueError(f"Invalid input data: {e}")
        return input_data_dict

    def do_run_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Perform inference on the input data

        Args:
            input_data (dict): input data supported by the task to perform inference on

        Returns:
            dict: output data from the inference call
        """
        output: dict[str, Any] = {}
        prompt = input_data.get("prompt", None)
        assert prompt is not None, "Prompt is required for inference"
        negative_prompt = input_data.get("negative_prompt", None)
        images = self.pipeline(  # type: ignore
            prompt, negative_prompt=negative_prompt
        ).images
        output["images"] = images[0]
        return {"output": output}

    def do_stream(self, preprocessed_input: Any) -> Iterator[Any]:
        raise NotImplementedError

    def do_postprocessing(
        self, input_data: Any, output: dict[str, Any]
    ) -> dict[str, Any]:
        # Postprocessing logic here
        return output

    def do_generate_proof(self) -> Any:
        raise NotImplementedError

__init__(pipeline, model=None, *args, **kwargs)

Initialize the Stable Diffusion Inference Workflow object

Parameters:

Name Type Description Default
pipeline SupportedPipelines

pipeline to be used for inference. Supported pipelines are "StableDiffusion"

 required
model str

Stable Diffusion model to be used for inference. Defaults to "runwayml/stable-diffusion-v1-5"

 None
vae_type str

Type of VAE to be used for inference: mse, ema. Defaults to "ema"

 required
torch_dtype str

Type of torch dtype to be used for inference. Defaults to "float32"

 required
enable_xformers bool

Enable xformers memory efficient attention. NOTE: requires xformers to be installed. Defaults to False

 required
Source code in src/infernet_ml/workflows/inference/stable_diffusion_workflow.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def __init__(
    self,
    pipeline: SupportedPipelines,
    model: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
) -> None:
    """
    Initialize the Stable Diffusion Inference Workflow object

    Args:
        pipeline (SupportedPipelines): pipeline to be used for inference. Supported
            pipelines are "StableDiffusion"
        model (str): Stable Diffusion model to be used for inference.
            Defaults to "runwayml/stable-diffusion-v1-5"
        vae_type (str): Type of VAE to be used for inference:
            `mse`, `ema`. Defaults to "ema"
        torch_dtype (str): Type of torch dtype to be used for inference.
            Defaults to "float32"
        enable_xformers (bool): Enable xformers memory efficient attention.
            NOTE: requires `xformers` to be installed. Defaults to False
    Raises:
        ValueError: if pipeline is not supported
    """
    super().__init__(*args, **kwargs)
    self.pipeline_id = pipeline.value
    # Check to ensure that the pipeline is supported
    # Else raise a ValueError
    if self.pipeline_id not in SUPPORTED_DIFFUSION_PIPELINES:
        raise ValueError(
            f"Pipeline {pipeline} is not supported.  \
            Supported pipelines are {SUPPORTED_DIFFUSION_PIPELINES}"
        )
    self.model_id = model or "runwayml/stable-diffusion-v1-5"
    self.vae_type = kwargs.get("vae_type", VaeType.ema)
    self.torch_dtype = kwargs.get("torch_dtype", "float32")
    self.enable_xformers = kwargs.get("enable_xformers", False)
    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    self.pipeline_options = StableDiffusionPipelineOptions(
        model=self.model_id,
        vae_type=self.vae_type,
        torch_dtype=self.torch_dtype,
        enable_xformers=self.enable_xformers,
    )
    self.inference_params = {
        "model": model,
    }

do_run_model(input_data)

Perform inference on the input data

Parameters:

Name Type Description Default
input_data dict

input data supported by the task to perform inference on

 required

Returns:

Name Type Description
dict dict[str, Any]

output data from the inference call
Source code in src/infernet_ml/workflows/inference/stable_diffusion_workflow.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
def do_run_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
    """
    Perform inference on the input data

    Args:
        input_data (dict): input data supported by the task to perform inference on

    Returns:
        dict: output data from the inference call
    """
    output: dict[str, Any] = {}
    prompt = input_data.get("prompt", None)
    assert prompt is not None, "Prompt is required for inference"
    negative_prompt = input_data.get("negative_prompt", None)
    images = self.pipeline(  # type: ignore
        prompt, negative_prompt=negative_prompt
    ).images
    output["images"] = images[0]
    return {"output": output}

do_setup()

Setup the inference client Returns: bool: True if setup is successful, False otherwise

Source code in src/infernet_ml/workflows/inference/stable_diffusion_workflow.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
def do_setup(self) -> bool:
    """
    Setup the inference client
    Returns:
        bool: True if setup is successful, False otherwise
    """
    done = False
    self.pipeline = self.get_pipeline(self.pipeline_id, self.pipeline_options)
    if self.pipeline_options.enable_xformers:
        self.pipeline.enable_xformers_memory_efficient_attention()  # type: ignore
    done = self.pipeline is not None
    logger.debug(f"Setup done: {done}")
    return done