diffusers-internal-dev/chronoedit-modular

1from diffusers.modular_pipelines import SequentialPipelineBlocks2from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict3 4from .before_denoise import ChronoEditPrepareLatentStep, ChronoEditSetTimestepsStep5from .denoise import ChronoEditDenoiseStep6from .inputs import ChronoEditImageInputStep, ChronoEditInputStep7from .encoders import ChronoEditImageEncoderStep, ChronoEditProcessImageStep, ChronoEditTextEncoderStep8from .decoders import ChronoEditDecodeStep9 10 11class ChronoEditBeforeDenoiseStep(SequentialPipelineBlocks):12    block_classes = [13        ChronoEditInputStep,14        ChronoEditSetTimestepsStep,15        ChronoEditImageInputStep,16        ChronoEditImageEncoderStep,17        ChronoEditProcessImageStep,18        ChronoEditPrepareLatentStep,19    ]20    block_names = [21        "input",22        "set_timesteps",23        "image_input_resolution",24        "encode_image",25        "process_image_for_latent_prep",26        "prepare_latent",27    ]28 29    @property30    def description(self):31        return (32            "Before denoise step that prepare the inputs for the denoise step.\n"33            + "This is a sequential pipeline blocks:\n"34            + " - `ChronoEditSetTimestepsStep` is used to set timesteps\n"35            + " - `ChronoEditImageInputStep` is used to set the image resolution\n"36            + " - `ChronoEditImageEncoderStep` is used to derive image embeddings for conditioning\n"37            + " - `ChronoEditProcessImageStep` is used to process the image for latent preparation\n"38            + " - `ChronoEditPrepareLatentStep` is used to prepare the latents"39        )40 41 42EDIT_BLOCKS = InsertableDict(43    [44        ("text_encoder", ChronoEditTextEncoderStep),45        ("before_denoise", ChronoEditBeforeDenoiseStep),46        ("denoise", ChronoEditDenoiseStep),47        ("decode", ChronoEditDecodeStep),48    ]49)50 51ALL_BLOCKS = {"img2img": EDIT_BLOCKS}52 53 54class ChronoEditBlocks(SequentialPipelineBlocks):55    block_classes = list(EDIT_BLOCKS.copy().values())56    block_names = list(EDIT_BLOCKS.copy().keys())57