Upload modular_blocks.py with huggingface_hub

modular_blocks.py

@@ -0,0 +1,57 @@
+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())