ModularPipelineBlocks

ModularPipelineBlocks — это базовый блок для построения ModularPipeline. Он определяет, какие компоненты, входные/выходные данные и вычисления блок должен выполнять на конкретном этапе конвейера. Блок ModularPipelineBlocks взаимодействует с другими блоками с помощью state, что позволяет модульно конструировать рабочие процессы.

Сам по себе ModularPipelineBlocks не может быть выполнен. Это чертёж того, что должен делать шаг в конвейере. Чтобы фактически запустить и выполнить конвейер, ModularPipelineBlocks необходимо преобразовать в ModularPipeline.

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

Входы и выходы

Если вам не знакомо, как работает состояние в Modular Diffusers, обратитесь к руководству States.

Для ModularPipelineBlocks необходимы inputs и intermediate_outputs.

inputs — это значения, которые блок считывает из PipelineState для выполнения вычислений. Это могут быть значения, предоставленные пользователем (например, промпт или изображение), или значения, полученные от предыдущего блока (например, закодированные image_latents).

Используйте InputParam для определения inputs.

class ImageEncodeStep(ModularPipelineBlocks):
    ...

    @property
    def inputs(self):
        return [
            InputParam(name="image", type_hint="PIL.Image", required=True, description="raw input image to process"),
        ]
    ...

intermediate_outputs — это новые значения, создаваемые блоком и добавляемые в PipelineState. Значения intermediate_outputs доступны в качестве inputs для последующих блоков или в качестве итогового результата выполнения конвейера.

Используйте OutputParam для определения intermediate_outputs.

class ImageEncodeStep(ModularPipelineBlocks):
    ...

    @property
    def intermediate_outputs(self):
        return [
            OutputParam(name="image_latents", description="latents representing the image"),
        ]

    ...

Промежуточные входные и выходные данные разделяют данные для соединения блоков. Они доступны в любой момент, что позволяет отслеживать ход выполнения рабочего процесса.

Компоненты и конфигурации

Компоненты и конфигурации на уровне пайплайна, необходимые для блока, задаются в ComponentSpec и ConfigSpec.

ComponentSpec содержит ожидаемые компоненты, используемые блоком. Вам необходимо указать name компонента и, в идеале, type_hint, который точно определяет, что представляет собой компонент.
ConfigSpec содержит настройки на уровне конвейера, которые управляют поведением всех блоков.

class ImageEncodeStep(ModularPipelineBlocks):
    ...

    @property
    def expected_components(self):
        return [
            ComponentSpec(name="vae", type_hint=AutoencoderKL),
        ]

    @property
    def expected_configs(self):
        return [
            ConfigSpec("force_zeros_for_empty_prompt", True),
        ]

    ...

Когда блоки преобразуются в конвейер, компоненты становятся доступными для блока в качестве первого аргумента в __call__.

Логика вычислений

Вычисления, выполняемые блоком, определяются в методе __call__ и следуют определённой структуре.

Получите BlockState, чтобы получить локальный вид inputs.
Реализуйте логику вычислений для inputs.
Обновите PipelineState, чтобы отправить изменения из локального BlockState обратно в глобальный PipelineState.
Возвращает компоненты и состояние, которые становятся доступными для следующего блока.

class ImageEncodeStep(ModularPipelineBlocks):

    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
        # 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

Собираем всё вместе

Вот полный блок со всеми соединёнными частями.

from diffusers import ComponentSpec, AutoencoderKL
from diffusers.modular_pipelines import InputParam, ModularPipelineBlocks, OutputParam


class ImageEncodeStep(ModularPipelineBlocks):

    @property
    def description(self):
        return "Encode an image into latent space."

    @property
    def expected_components(self):
        return [
            ComponentSpec(name="vae", type_hint=AutoencoderKL),
        ]

    @property
    def inputs(self):
        return [
            InputParam(name="image", type_hint="PIL.Image", required=True, description="raw input image to process"),
        ]

    @property
    def intermediate_outputs(self):
        return [
            OutputParam(name="image_latents", type_hint="torch.Tensor", description="latents representing the image"),
        ]

    def __call__(self, components, state):
        block_state = self.get_block_state(state)
        block_state.image_latents = components.vae.encode(block_state.image)
        self.set_block_state(state, block_state)
        return components, state

У каждого блока есть свойство doc, которое автоматически генерируется на основе свойств, определённых выше. Оно содержит краткое описание блока, его компонентов, входов и выходов.

block = ImageEncoderStep()
print(block.doc)
class ImageEncodeStep

  Encode an image into latent space.

  Components:
      vae (`AutoencoderKL`)

  Inputs:
      image (`PIL.Image`):
          raw input image to process

  Outputs:
      image_latents (`torch.Tensor`):
          latents representing the image

Update on GitHub