Buckets:

mishig's picture
|
download
raw
7.13 kB

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 是由某个 block 创建并添加到 PipelineState 中的新值。intermediate_outputs 可作为后续 block 的 inputs 使用,也可作为运行 pipeline 后的最终输出。

    使用 OutputParam 来定义 intermediate_outputs。

class ImageEncodeStep(ModularPipelineBlocks):
    ...

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

    ...

中间输入和输出通过共享数据来连接各个块。它们在任何时间点都可访问,从而允许您跟踪工作流的进度。

组件和配置

一个 block 所需的组件和流水线级别的配置在 ComponentSpec 和 ConfigSpec 中指定。

  • ComponentSpec 包含一个 block 所使用的预期组件。你需要该组件的 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__ 方法中定义,并遵循特定的结构。

  1. 获取 BlockState 以获取 inputs 的本地视图。
  2. 在 inputs 上实现计算逻辑。
  3. 更新 PipelineState,将本地 BlockState 中的更改推送回全局 PipelineState。
  4. 返回可供下一个块使用的组件和状态。
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

Xet Storage Details

Size:
7.13 kB
·
Xet hash:
16dd2ea9437487afc8fc900a1b8437dbb855b599363c8318e5e743913206d92b

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.