modular_blocks.py

from diffusers.modular_pipelines import SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict
from diffusers.modular_pipelines.wan.decoders import WanDecodeStep

from before_denoise import ChronoEditPrepareLatentStep, ChronoEditSetTimestepsStep
from denoise import ChronoEditDenoiseStep
from inputs import ChronoEditImageInputStep, ChronoEditInputStep
from encoders import ChronoEditImageEncoderStep, ChronoEditProcessImageStep, ChronoEditTextEncoderStep
from decoders import ChronoEditDecodeStep


class ChronoEditBeforeDenoiseStep(SequentialPipelineBlocks):
    block_classes = [
        ChronoEditInputStep,
        ChronoEditSetTimestepsStep,
        ChronoEditImageInputStep,
        ChronoEditImageEncoderStep,
        ChronoEditProcessImageStep,
        ChronoEditPrepareLatentStep,
    ]
    block_names = [
        "input",
        "set_timesteps",
        "image_input_resolution",
        "encode_image",
        "process_image_for_latent_prep",
        "prepare_latent",
    ]

    @property
    def description(self):
        return (
            "Before denoise step that prepare the inputs for the denoise step.\n"
            + "This is a sequential pipeline blocks:\n"
            + " - `ChronoEditSetTimestepsStep` is used to set timesteps\n"
            + " - `ChronoEditImageInputStep` is used to set the image resolution\n"
            + " - `ChronoEditImageEncoderStep` is used to derive image embeddings for conditioning\n"
            + " - `ChronoEditProcessImageStep` is used to process the image for latent preparation\n"
            + " - `ChronoEditPrepareLatentStep` is used to prepare the latents"
        )


EDIT_BLOCKS = InsertableDict(
    [
        ("text_encoder", ChronoEditTextEncoderStep),
        ("before_denoise", ChronoEditBeforeDenoiseStep),
        ("denoise", ChronoEditDenoiseStep),
        ("decode", ChronoEditDecodeStep),
    ]
)

ALL_BLOCKS = {"img2img": EDIT_BLOCKS}


class ChronoEditBlocks(SequentialPipelineBlocks):
    block_classes = list(EDIT_BLOCKS.copy().values())
    block_names = list(EDIT_BLOCKS.copy().keys())