LoopSequentialPipelineBlocks

LoopSequentialPipelineBlocks — это многоблочный тип, который объединяет другие ModularPipelineBlocks в цикле. Данные циркулируют, используя inputs и intermediate_outputs, при этом каждый блок выполняется итеративно. Обычно это используется для создания цикла шумоподавления, который по умолчанию является итеративным.

В этом руководстве показано, как создавать LoopSequentialPipelineBlocks.

Обёртка цикла

LoopSequentialPipelineBlocks, также известный как обёртка цикла, поскольку он определяет структуру цикла, переменные итерации и конфигурацию. Внутри обёртки цикла вам потребуются следующие переменные.

loop_inputs — это значения, предоставленные пользователем, и они эквивалентны inputs.
loop_intermediate_outputs — это новые промежуточные переменные, создаваемые блоком и добавляемые в PipelineState. Это эквивалентно intermediate_outputs.
Метод __call__ определяет структуру цикла и логику итераций.

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

Обёртка цикла может передавать дополнительные аргументы, такие как индекс текущей итерации, в блоки цикла.

Блоки циклов

Блок цикла представляет собой ModularPipelineBlocks, но метод __call__ работает иначе.

Он получает переменную итерации от обёртки цикла.
Работает непосредственно с BlockState, а не с PipelineState.
Для этого не требуется получать или обновлять BlockState.

Блоки цикла разделяют один и тот же BlockState, что позволяет значениям накапливаться и изменяться на каждой итерации цикла.

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