ModularPipelineBlocks

ModularPipelineBlocks é o bloco básico para a construção de um ModularPipeline. Ele define quais componentes, entradas/saídas e computação um bloco deve executar para uma etapa específica em um pipeline. Um ModularPipelineBlocks se conecta a outros blocos, usando state, para permitir a construção modular de fluxos de trabalho.

Um ModularPipelineBlocks isolado não pode ser executado. Ele é um modelo do que uma etapa deve fazer em um pipeline. Para realmente executar e rodar um pipeline, o ModularPipelineBlocks precisa ser convertido em um ModularPipeline.

Este guia mostrará como criar um ModularPipelineBlocks.

Entradas e saídas

Consulte o guia States caso você não esteja familiarizado com o funcionamento do estado no Modular Diffusers.

A ModularPipelineBlocks requer inputs e intermediate_outputs.

inputs são os valores que um bloco lê do PipelineState para executar seu cálculo. Esses valores podem ser fornecidos por um usuário (como um prompt ou imagem) ou produzidos por um bloco anterior (como image_latents codificados).

Use InputParam para definir 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 são novos valores criados por um bloco e adicionados ao PipelineState. Os intermediate_outputs estão disponíveis como inputs para blocos subsequentes ou como a saída final da execução do pipeline.

Use OutputParam para definir intermediate_outputs.

class ImageEncodeStep(ModularPipelineBlocks):
    ...

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

    ...

As entradas e saídas intermediárias compartilham dados para conectar os blocos. Elas são acessíveis a qualquer momento, permitindo que você acompanhe o progresso do fluxo de trabalho.

Componentes e configurações

Os componentes e as configurações de nível de pipeline de que um bloco precisa são especificados em ComponentSpec e ConfigSpec.

ComponentSpec contém os componentes esperados usados por um bloco. Você precisa do name do componente e, idealmente, de um type_hint que especifique exatamente o que é o componente.
ConfigSpec contém configurações de nível de pipeline que controlam o comportamento em todos os blocos.

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),
        ]

    ...

Quando os blocos são convertidos em um pipeline, os componentes ficam disponíveis para o bloco como o primeiro argumento em __call__.

Lógica de computação

O cálculo que um bloco executa é definido no método __call__ e segue uma estrutura específica.

Obtenha o BlockState para obter uma visão local dos inputs.
Implemente a lógica de computação nos inputs.
Atualize o PipelineState para enviar as alterações do BlockState local de volta para o PipelineState global.
Retorna os componentes e o estado que ficam disponíveis para o próximo bloco.

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

Juntando tudo

Aqui está o bloco completo com todas as partes conectadas.

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

Todo bloco possui uma propriedade doc que é gerada automaticamente a partir das propriedades definidas acima. Ela fornece um resumo da descrição, dos componentes, das entradas e das saídas do bloco.

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