PipelineBlock

In Modular Diffusers, you build your workflow using ModularPipelineBlocks. We support 4 different types of blocks: PipelineBlock, SequentialPipelineBlocks, LoopSequentialPipelineBlocks, and AutoPipelineBlocks. Among them, PipelineBlock is the most fundamental building block of the whole system - it’s like a brick in a Lego system. These blocks are designed to easily connect with each other, allowing for modular construction of creative and potentially very complex workflows.

In this tutorial, we will focus on how to write a basic PipelineBlock and how it interacts with the pipeline state.

PipelineState

Before we dive into creating PipelineBlocks, make sure you have a basic understanding of PipelineState. It acts as the global state container that all blocks operate on - each block gets a local view (BlockState) of the relevant variables it needs from PipelineState, performs its operations, and then updates PipelineState with any changes. See the PipelineState and BlockState guide for more details.

Define a PipelineBlock

To write a PipelineBlock class, you need to define a few properties that determine how your block interacts with the pipeline state. Understanding these properties is crucial - they define what data your block can access and what it can produce.

The three main properties you need to define are:

• inputs: Immutable values from the user that cannot be modified
• intermediate_inputs: Mutable values from previous blocks that can be read and modified
• intermediate_outputs: New values your block creates for subsequent blocks and user access

Let’s explore each one and understand how they work with the pipeline state.

Inputs: Immutable User Values

Inputs are variables your block needs from the immutable pipeline state - these are user-provided values that cannot be modified by any block. You define them using InputParam:

user_inputs = [
    InputParam(name="image", type_hint="PIL.Image", description="raw input image to process")
]

When you list something as an input, you’re saying “I need this value directly from the end user, and I will talk to them directly, telling them what I need in the ‘description’ field. They will provide it and it will come to me unchanged.”

This is especially useful for raw values that serve as the “source of truth” in your workflow. For example, with a raw image, many workflows require preprocessing steps like resizing that a previous block might have performed. But in many cases, you also want the raw PIL image. In some inpainting workflows, you need the original image to overlay with the generated result for better control and consistency.

Intermediate Inputs: Mutable Values from Previous Blocks, or Users

Intermediate inputs are variables your block needs from the mutable pipeline state - these are values that can be read and modified. They’re typically created by previous blocks, but could also be directly provided by the user if not the case:

user_intermediate_inputs = [
    InputParam(name="processed_image", type_hint="torch.Tensor", description="image that has been preprocessed and normalized"),
]

When you list something as an intermediate input, you’re saying “I need this value, but I want to work with a different block that has already created it. I already know for sure that I can get it from this other block, but it’s okay if other developers want use something different.”

Intermediate Outputs: New Values for Subsequent Blocks and User Access

Intermediate outputs are new variables your block creates and adds to the mutable pipeline state. They serve two purposes:

1. For subsequent blocks: They can be used as intermediate inputs by other blocks in the pipeline
2. For users: They become available as final outputs that users can access when running the pipeline

user_intermediate_outputs = [
    OutputParam(name="image_latents", description="latents representing the image")
]

Intermediate inputs and intermediate outputs work together like Lego studs and anti-studs - they’re the connection points that make blocks modular. When one block produces an intermediate output, it becomes available as an intermediate input for subsequent blocks. This is where the “modular” nature of the system really shines - blocks can be connected and reconnected in different ways as long as their inputs and outputs match.

Additionally, all intermediate outputs are accessible to users when they run the pipeline, typically you would only need the final images, but they are also able to access intermediate results like latents, embeddings, or other processing steps.

The __call__ Method Structure

Your PipelineBlock’s __call__ method should follow this structure:

def __call__(self, components, state):
    # Get a local view of the state variables this block needs
    block_state = self.get_block_state(state)
    
    # Your computation logic here
    # block_state contains all your inputs and intermediate_inputs
    # You can access them like: block_state.image, block_state.processed_image
    
    # Update the pipeline state with your updated block_states
    self.set_block_state(state, block_state)
    return components, state

The block_state object contains all the variables you defined in inputs and intermediate_inputs, making them easily accessible for your computation.

Components and Configs

You can define the components and pipeline-level configs your block needs using ComponentSpec and ConfigSpec:

from diffusers import ComponentSpec, ConfigSpec

# Define components your block needs
expected_components = [
    ComponentSpec(name="unet", type_hint=UNet2DConditionModel),
    ComponentSpec(name="scheduler", type_hint=EulerDiscreteScheduler)
]

# Define pipeline-level configs
expected_config = [
    ConfigSpec("force_zeros_for_empty_prompt", True)
]

Components: In the ComponentSpec, you must provide a name and ideally a type_hint. You can also specify a default_creation_method to indicate whether the component should be loaded from a pretrained model or created with default configurations. The actual loading details (repo, subfolder, variant and revision fields) are typically specified when creating the pipeline, as we covered in the Modular Pipeline Guide.

Configs: Pipeline-level settings that control behavior across all blocks.

When you convert your blocks into a pipeline using blocks.init_pipeline(), the pipeline collects all component requirements from the blocks and fetches the loading specs from the modular repository. The components are then made available to your block as the first argument of the __call__ method. You can access any component you need using dot notation:

def __call__(self, components, state):
    # Access components using dot notation
    unet = components.unet
    vae = components.vae
    scheduler = components.scheduler

That’s all you need to define in order to create a PipelineBlock. There is no hidden complexity. In fact we are going to create a helper function that take exactly these variables as input and return a pipeline block. We will use this helper function through out the tutorial to create test blocks

Note that for __call__ method, the only part you should implement differently is the part between self.get_block_state() and self.set_block_state(), which can be abstracted into a simple function that takes block_state and returns the updated state. Our helper function accepts a block_fn that does exactly that.

Helper Function

from diffusers.modular_pipelines import PipelineBlock, InputParam, OutputParam
import torch

def make_block(inputs=[], intermediate_inputs=[], intermediate_outputs=[], block_fn=None, description=None):
    class TestBlock(PipelineBlock):
        model_name = "test"
        
        @property
        def inputs(self):
            return inputs
            
        @property
        def intermediate_inputs(self):
            return intermediate_inputs
            
        @property
        def intermediate_outputs(self):
            return intermediate_outputs
            
        @property
        def description(self):
            return description if description is not None else ""
            
        def __call__(self, components, state):
            block_state = self.get_block_state(state)
            if block_fn is not None:
                block_state = block_fn(block_state, state)
            self.set_block_state(state, block_state)
            return components, state
    
    return TestBlock

Example: Creating a Simple Pipeline Block

Let’s create a simple block to see how these definitions interact with the pipeline state. To better understand what’s happening, we’ll print out the states before and after updates to inspect them:

inputs = [
    InputParam(name="image", type_hint="PIL.Image", description="raw input image to process")
]

intermediate_inputs = [InputParam(name="batch_size", type_hint=int)]

intermediate_outputs = [
    OutputParam(name="image_latents", description="latents representing the image")
]

def image_encoder_block_fn(block_state, pipeline_state):
    print(f"pipeline_state (before update): {pipeline_state}")
    print(f"block_state (before update): {block_state}")
    
    # Simulate processing the image
    block_state.image = torch.randn(1, 3, 512, 512)
    block_state.batch_size = block_state.batch_size * 2
    block_state.processed_image = [torch.randn(1, 3, 512, 512)] * block_state.batch_size
    block_state.image_latents = torch.randn(1, 4, 64, 64)
    
    print(f"block_state (after update): {block_state}")
    return block_state

# Create a block with our definitions
image_encoder_block_cls = make_block(
    inputs=inputs, 
    intermediate_inputs=intermediate_inputs,
    intermediate_outputs=intermediate_outputs, 
    block_fn=image_encoder_block_fn,
    description="Encode raw image into its latent presentation"
)
image_encoder_block = image_encoder_block_cls()
pipe = image_encoder_block.init_pipeline()

Let’s check the pipeline’s docstring to see what inputs it expects:

>>> print(pipe.doc)
class TestBlock

  Encode raw image into its latent presentation

  Inputs:

      image (`PIL.Image`, *optional*):
          raw input image to process

      batch_size (`int`, *optional*):

  Outputs:

      image_latents (`None`):
          latents representing the image

Notice that batch_size appears as an input even though we defined it as an intermediate input. This happens because no previous block provided it, so the pipeline makes it available as a user input. However, unlike regular inputs, this value goes directly into the mutable intermediate state.

Now let’s run the pipeline:

from diffusers.utils import load_image

image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/image_of_squirrel_painting.png")
state = pipe(image=image, batch_size=2)
print(f"pipeline_state (after update): {state}")

pipeline_state (before update): PipelineState(
  inputs={
    image: <PIL.Image.Image image mode=RGB size=512x512 at 0x7F3ECC494550>
  },
  intermediates={
    batch_size: 2
  },
)
block_state (before update): BlockState(
    image: <PIL.Image.Image image mode=RGB size=512x512 at 0x7F3ECC494640>
    batch_size: 2
)

block_state (after update): BlockState(
    image: Tensor(dtype=torch.float32, shape=torch.Size([1, 3, 512, 512]))
    batch_size: 4
    processed_image: List[4] of Tensors with shapes [torch.Size([1, 3, 512, 512]), torch.Size([1, 3, 512, 512]), torch.Size([1, 3, 512, 512]), torch.Size([1, 3, 512, 512])]
    image_latents: Tensor(dtype=torch.float32, shape=torch.Size([1, 4, 64, 64]))
)
pipeline_state (after update): PipelineState(
  inputs={
    image: <PIL.Image.Image image mode=RGB size=512x512 at 0x7F3ECC494550>
  },
  intermediates={
    batch_size: 4
    image_latents: Tensor(dtype=torch.float32, shape=torch.Size([1, 4, 64, 64]))
  },
)

Key Observations:

1. Before the update: image (the input) goes to the immutable inputs dict, while batch_size (the intermediate_input) goes to the mutable intermediates dict, and both are available in block_state.

2. After the update:

   • image (inputs) changed in block_state but not in pipeline_state - this change is local to the block only.
   • batch_size (intermediate_inputs) was updated in both block_state and pipeline_state - this change affects subsequent blocks (we didn’t need to declare it as an intermediate output since it was already in the intermediates dict)
   • image_latents (intermediate_outputs) was added to pipeline_state because it was declared as an intermediate output
   • processed_image was not added to pipeline_state because it wasn’t declared as an intermediate output