ModularPipelineBlocks

ModularPipelineBlocks es el bloque básico para construir un ModularPipeline. Define qué componentes, entradas/salidas y cómputo debe realizar un bloque para un paso específico en un pipeline. Un ModularPipelineBlocks se conecta con otros bloques, utilizando state, para permitir la construcción modular de flujos de trabajo.

Un ModularPipelineBlocks por sí solo no se puede ejecutar. Es un plano de lo que un paso debe hacer en un pipeline. Para ejecutar y procesar un pipeline, los ModularPipelineBlocks deben convertirse en un ModularPipeline.

Esta guía te mostrará cómo crear un ModularPipelineBlocks.

Entradas y salidas

Consulte la guía de States si no está familiarizado con cómo funciona el estado en Modular Diffusers.

Un ModularPipelineBlocks requiere inputs y intermediate_outputs.

inputs son los valores que un bloque lee del PipelineState para realizar su cálculo. Estos pueden ser valores proporcionados por un usuario (como un prompt o una imagen) o valores generados por un bloque 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 son nuevos valores creados por un bloque y añadidos al PipelineState. Los intermediate_outputs están disponibles como inputs para bloques posteriores o como la salida final de la ejecución del 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"),
        ]

    ...

Las entradas y salidas intermedias comparten datos para conectar bloques. Son accesibles en cualquier punto, lo que permite seguir el progreso del flujo de trabajo.

Componentes y configuraciones

Los componentes y las configuraciones a nivel de pipeline que necesita un bloque se especifican en ComponentSpec y ConfigSpec.

ComponentSpec contiene los componentes esperados que utiliza un bloque. Necesitas el name del componente e, idealmente, un type_hint que especifique exactamente qué es el componente.
ConfigSpec contiene configuraciones a nivel de pipeline que controlan el comportamiento en todos los bloques.

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

    ...

Cuando los bloques se convierten en un pipeline, los componentes quedan disponibles para el bloque como primer argumento en __call__.

Lógica de cálculo

El cálculo que realiza un bloque se define en el método __call__ y sigue una estructura específica.

Obtenga el BlockState para obtener una vista local de los inputs.
Implemente la lógica de cálculo en los inputs.
Actualice PipelineState para enviar los cambios del BlockState local de vuelta al PipelineState global.
Devuelve los componentes y el estado que quedan disponibles para el siguiente bloque.

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

Poniéndolo todo en práctica

Aquí está el bloque completo con todas las piezas 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

Cada bloque tiene una propiedad doc que se genera automáticamente a partir de las propiedades que definiste anteriormente. Proporciona un resumen de la descripción, los componentes, las entradas y las salidas del bloque.

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