LoopSequentialPipelineBlocks

LoopSequentialPipelineBlocks est un type multi-blocs qui compose d’autres ModularPipelineBlocks dans une boucle. Les données circulent de manière circulaire, en utilisant inputs et intermediate_outputs, et chaque bloc est exécuté itérativement. Ce type est généralement utilisé pour créer une boucle de débruitage, qui est itérative par défaut.

Ce guide vous montre comment créer des LoopSequentialPipelineBlocks.

Boucle d'encapsulation

LoopSequentialPipelineBlocks, également connu sous le nom de loop wrapper, car il définit la structure de la boucle, les variables d’itération et la configuration. À l’intérieur du loop wrapper, vous avez besoin des variables suivantes.

• loop_inputs sont des valeurs fournies par l'utilisateur et équivalentes à inputs.
• loop_intermediate_outputs sont de nouvelles variables intermédiaires créées par le bloc et ajoutées à l'PipelineState. C'est équivalent à intermediate_outputs.
• La méthode __call__ définit la structure de la boucle et la logique d'itération.

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

Le wrapper de boucle peut transmettre des arguments supplémentaires, comme l'index de l'itération actuelle, aux blocs de la boucle.

Blocs de boucle

Un bloc de boucle est un ModularPipelineBlocks, mais la méthode __call__ se comporte différemment.

• Il reçoit la variable d'itération de l'enveloppe de boucle.
• Il fonctionne directement avec le BlockState plutôt qu'avec le PipelineState.
• Il ne nécessite pas de récupérer ou de mettre à jour le BlockState.

Les blocs de boucle partagent le même BlockState afin de permettre aux valeurs de s'accumuler et de changer à chaque itération de la boucle.

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

LoopSequentialPipelineBlocks

Utilisez la méthode from_blocks_dict() pour ajouter le bloc de boucle à l'enveloppe de boucle afin de créer LoopSequentialPipelineBlocks.

loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock})

Ajoutez d’autres blocs de boucle à exécuter au sein de chaque itération à l’aide de from_blocks_dict(). Cela vous permet de modifier les blocs sans changer la logique de la boucle elle-même.

loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock(), "block2": LoopBlock})

Pour ajouter un bloc de boucle à la boucle de débruitage d’un pipeline existant, par exemple un bloc qui exécute une fonction callback_on_step_end à chaque étape, consultez Modular pipelines dans le guide des callbacks.