update

Author: DN6

src/diffusers/modular_pipelines/flux/before_denoise.py

@@ -21,7 +21,7 @@
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 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 FluxModularPipeline
 
@@ -125,7 +125,7 @@ def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
     return latent_image_ids.to(device=device, dtype=dtype)
 
 
-class FluxInputStep(PipelineBlock):
+class FluxInputStep(ModularPipelineBlocks):
     model_name = "flux"
 
     @property
@@ -143,11 +143,6 @@ def description(self) -> str:
     def inputs(self) -> List[InputParam]:
         return [
             InputParam("num_images_per_prompt", default=1),
-        ]
-
-    @property
-    def intermediate_inputs(self) -> List[str]:
-        return [
             InputParam(
                 "prompt_embeds",
                 required=True,
@@ -216,7 +211,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip
         return components, state
 
 
-class FluxSetTimestepsStep(PipelineBlock):
+class FluxSetTimestepsStep(ModularPipelineBlocks):
     model_name = "flux"
 
     @property
@@ -235,17 +230,12 @@ def inputs(self) -> List[InputParam]:
             InputParam("sigmas"),
             InputParam("guidance_scale", default=3.5),
             InputParam("latents", type_hint=torch.Tensor),
-        ]
-
-    @property
-    def intermediate_inputs(self) -> List[str]:
-        return [
             InputParam(
                 "latents",
                 required=True,
                 type_hint=torch.Tensor,
                 description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
-            )
+            ),
         ]
 
     @property
@@ -296,7 +286,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip
         return components, state
 
 
-class FluxPrepareLatentsStep(PipelineBlock):
+class FluxPrepareLatentsStep(ModularPipelineBlocks):
     model_name = "flux"
 
     @property
@@ -314,11 +304,6 @@ def inputs(self) -> List[InputParam]:
             InputParam("width", type_hint=int),
             InputParam("latents", type_hint=Optional[torch.Tensor]),
             InputParam("num_images_per_prompt", type_hint=int, default=1),
-        ]
-
-    @property
-    def intermediate_inputs(self) -> List[InputParam]:
-        return [
             InputParam("generator"),
             InputParam(
                 "batch_size",

src/diffusers/modular_pipelines/flux/decoders.py

@@ -22,7 +22,7 @@
 from ...models import AutoencoderKL
 from ...utils import logging
 from ...video_processor import VaeImageProcessor
-from ..modular_pipeline import PipelineBlock, PipelineState
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
 
 
@@ -45,7 +45,7 @@ def _unpack_latents(latents, height, width, vae_scale_factor):
     return latents
 
 
-class FluxDecodeStep(PipelineBlock):
+class FluxDecodeStep(ModularPipelineBlocks):
     model_name = "flux"
 
     @property
@@ -70,17 +70,12 @@ def inputs(self) -> List[Tuple[str, Any]]:
             InputParam("output_type", default="pil"),
             InputParam("height", default=1024),
             InputParam("width", default=1024),
-        ]
-
-    @property
-    def intermediate_inputs(self) -> List[str]:
-        return [
             InputParam(
                 "latents",
                 required=True,
                 type_hint=torch.Tensor,
                 description="The denoised latents from the denoising step",
-            )
+            ),
         ]
 
     @property

src/diffusers/modular_pipelines/flux/denoise.py

@@ -22,7 +22,7 @@
 from ..modular_pipeline import (
     BlockState,
     LoopSequentialPipelineBlocks,
-    PipelineBlock,
+    ModularPipelineBlocks,
     PipelineState,
 )
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
@@ -32,7 +32,7 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-class FluxLoopDenoiser(PipelineBlock):
+class FluxLoopDenoiser(ModularPipelineBlocks):
     model_name = "flux"
 
     @property
@@ -49,11 +49,8 @@ def description(self) -> str:
 
     @property
     def inputs(self) -> List[Tuple[str, Any]]:
-        return [InputParam("joint_attention_kwargs")]
-
-    @property
-    def intermediate_inputs(self) -> List[str]:
         return [
+            InputParam("joint_attention_kwargs"),
             InputParam(
                 "latents",
                 required=True,
@@ -113,7 +110,7 @@ def __call__(
         return components, block_state
 
 
-class FluxLoopAfterDenoiser(PipelineBlock):
+class FluxLoopAfterDenoiser(ModularPipelineBlocks):
     model_name = "flux"
 
     @property
@@ -175,7 +172,7 @@ def loop_expected_components(self) -> List[ComponentSpec]:
         ]
 
     @property
-    def loop_intermediate_inputs(self) -> List[InputParam]:
+    def loop_inputs(self) -> List[InputParam]:
         return [
             InputParam(
                 "timesteps",

src/diffusers/modular_pipelines/flux/encoders.py

@@ -21,7 +21,7 @@
 
 from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin
 from ...utils import USE_PEFT_BACKEND, is_ftfy_available, logging, scale_lora_layers, unscale_lora_layers
-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 FluxModularPipeline
 
@@ -50,7 +50,7 @@ def prompt_clean(text):
     return text
 
 
-class FluxTextEncoderStep(PipelineBlock):
+class FluxTextEncoderStep(ModularPipelineBlocks):
     model_name = "flux"
 
     @property