Diffusers documentation
LoopSequentialPipelineBlocks
LoopSequentialPipelineBlocks
LoopSequentialPipelineBlocks are a multi-block type that composes other ModularPipelineBlocks together in a loop. Data flows circularly, using intermediate_inputs
and intermediate_outputs
, and each block is run iteratively. This is typically used to create a denoising loop which is iterative by default.
This guide shows you how to create LoopSequentialPipelineBlocks.
Loop wrapper
LoopSequentialPipelineBlocks, is also known as the loop wrapper because it defines the loop structure, iteration variables, and configuration. Within the loop wrapper, you need the following variables.
loop_inputs
are user provided values and equivalent toinputs
.loop_intermediate_inputs
are intermediate variables from the PipelineState and equivalent to~modular_pipelines.ModularPipelineBlocks.intermediate_inputs
.loop_intermediate_outputs
are new intermediate variables created by the block and added to the PipelineState. It is equivalent tointermediate_outputs
.__call__
method defines the loop structure and iteration logic.
import torch
from diffusers.modular_pipelines import LoopSequentialPipelineBlocks, ModularPipelineBlocks, InputParam, OutputParam
class LoopWrapper(LoopSequentialPipelineBlocks):
model_name = "test"
@property
def description(self):
return "I'm a loop!!"
@property
def loop_inputs(self):
return [InputParam(name="num_steps")]
@torch.no_grad()
def __call__(self, components, state):
block_state = self.get_block_state(state)
# Loop structure - can be customized to your needs
for i in range(block_state.num_steps):
# loop_step executes all registered blocks in sequence
components, block_state = self.loop_step(components, block_state, i=i)
self.set_block_state(state, block_state)
return components, state
The loop wrapper can pass additional arguments, like current iteration index, to the loop blocks.
Loop blocks
A loop block is a ModularPipelineBlocks, but the __call__
method behaves differently.
- It recieves the iteration variable from the loop wrapper.
- It works directly with the BlockState instead of the PipelineState.
- It doesn’t require retrieving or updating the BlockState.
Loop blocks share the same BlockState to allow values to accumulate and change for each iteration in the loop.
class LoopBlock(ModularPipelineBlocks):
model_name = "test"
@property
def inputs(self):
return [InputParam(name="x")]
@property
def intermediate_outputs(self):
# outputs produced by this block
return [OutputParam(name="x")]
@property
def description(self):
return "I'm a block used inside the `LoopWrapper` class"
def __call__(self, components, block_state, i: int):
block_state.x += 1
return components, block_state
LoopSequentialPipelineBlocks
Use the from_blocks_dict() method to add the loop block to the loop wrapper to create LoopSequentialPipelineBlocks.
loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock})
Add more loop blocks to run within each iteration with from_blocks_dict(). This allows you to modify the blocks without changing the loop logic itself.
loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock(), "block2": LoopBlock})