update

Author: DN6

src/diffusers/modular_pipelines/wan/before_denoise.py

@@ -20,7 +20,7 @@
 from ...schedulers import UniPCMultistepScheduler
 from ...utils import logging
 from ...utils.torch_utils import randn_tensor
-from ..modular_pipeline import PipelineBlock, PipelineState
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
 from .modular_pipeline import WanModularPipeline
 
@@ -94,7 +94,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class WanInputStep(PipelineBlock):
+class WanInputStep(ModularPipelineBlocks):
     model_name = "wan"
 
     @property
@@ -194,7 +194,7 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe
         return components, state
 
 
-class WanSetTimestepsStep(PipelineBlock):
+class WanSetTimestepsStep(ModularPipelineBlocks):
     model_name = "wan"
 
     @property
@@ -243,7 +243,7 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe
         return components, state
 
 
-class WanPrepareLatentsStep(PipelineBlock):
+class WanPrepareLatentsStep(ModularPipelineBlocks):
     model_name = "wan"
 
     @property

src/diffusers/modular_pipelines/wan/decoders.py

@@ -22,14 +22,14 @@
 from ...models import AutoencoderKLWan
 from ...utils import logging
 from ...video_processor import VideoProcessor
-from ..modular_pipeline import PipelineBlock, PipelineState
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-class WanDecodeStep(PipelineBlock):
+class WanDecodeStep(ModularPipelineBlocks):
     model_name = "wan"
 
     @property

src/diffusers/modular_pipelines/wan/denoise.py

@@ -23,8 +23,8 @@
 from ...utils import logging
 from ..modular_pipeline import (
     BlockState,
-    LoopSequentialPipelineBlocks,
-    PipelineBlock,
+    LoopSequentialModularPipelineBlockss,
+    ModularPipelineBlocks,
     PipelineState,
 )
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
@@ -34,7 +34,7 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-class WanLoopDenoiser(PipelineBlock):
+class WanLoopDenoiser(ModularPipelineBlocks):
     model_name = "wan"
 
     @property
@@ -53,7 +53,7 @@ def expected_components(self) -> List[ComponentSpec]:
     def description(self) -> str:
         return (
             "Step within the denoising loop that denoise the latents with guidance. "
-            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialModularPipelineBlockss` "
             "object (e.g. `WanDenoiseLoopWrapper`)"
         )
 
@@ -132,7 +132,7 @@ def __call__(
         return components, block_state
 
 
-class WanLoopAfterDenoiser(PipelineBlock):
+class WanLoopAfterDenoiser(ModularPipelineBlocks):
     model_name = "wan"
 
     @property
@@ -145,7 +145,7 @@ def expected_components(self) -> List[ComponentSpec]:
     def description(self) -> str:
         return (
             "step within the denoising loop that update the latents. "
-            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialModularPipelineBlockss` "
             "object (e.g. `WanDenoiseLoopWrapper`)"
         )
 
@@ -181,7 +181,7 @@ def __call__(self, components: WanModularPipeline, block_state: BlockState, i: i
         return components, block_state
 
 
-class WanDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+class WanDenoiseLoopWrapper(LoopSequentialModularPipelineBlockss):
     model_name = "wan"
 
     @property

src/diffusers/modular_pipelines/wan/encoders.py

@@ -22,7 +22,7 @@
 from ...configuration_utils import FrozenDict
 from ...guiders import ClassifierFreeGuidance
 from ...utils import is_ftfy_available, logging
-from ..modular_pipeline import PipelineBlock, PipelineState
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
 from .modular_pipeline import WanModularPipeline
 
@@ -51,7 +51,7 @@ def prompt_clean(text):
     return text
 
 
-class WanTextEncoderStep(PipelineBlock):
+class WanTextEncoderStep(ModularPipelineBlocks):
     model_name = "wan"
 
     @property

src/diffusers/modular_pipelines/wan/modular_blocks.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ...utils import logging
-from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline import AutoModularPipelineBlockss, SequentialModularPipelineBlockss
 from ..modular_pipeline_utils import InsertableDict
 from .before_denoise import (
     WanInputStep,
@@ -29,7 +29,7 @@
 
 
 # before_denoise: text2vid
-class WanBeforeDenoiseStep(SequentialPipelineBlocks):
+class WanBeforeDenoiseStep(SequentialModularPipelineBlockss):
     block_classes = [
         WanInputStep,
         WanSetTimestepsStep,
@@ -49,7 +49,7 @@ def description(self):
 
 
 # before_denoise: all task (text2vid,)
-class WanAutoBeforeDenoiseStep(AutoPipelineBlocks):
+class WanAutoBeforeDenoiseStep(AutoModularPipelineBlockss):
     block_classes = [
         WanBeforeDenoiseStep,
     ]
@@ -66,7 +66,7 @@ def description(self):
 
 
 # denoise: text2vid
-class WanAutoDenoiseStep(AutoPipelineBlocks):
+class WanAutoDenoiseStep(AutoModularPipelineBlockss):
     block_classes = [
         WanDenoiseStep,
     ]
@@ -83,7 +83,7 @@ def description(self) -> str:
 
 
 # decode: all task (text2img, img2img, inpainting)
-class WanAutoDecodeStep(AutoPipelineBlocks):
+class WanAutoDecodeStep(AutoModularPipelineBlockss):
     block_classes = [WanDecodeStep]
     block_names = ["non-inpaint"]
     block_trigger_inputs = [None]
@@ -94,7 +94,7 @@ def description(self):
 
 
 # text2vid
-class WanAutoBlocks(SequentialPipelineBlocks):
+class WanAutoBlocks(SequentialModularPipelineBlockss):
     block_classes = [
         WanTextEncoderStep,
         WanAutoBeforeDenoiseStep,