From d1342d7464955c71c9ecacca41a6f1b2a502197e Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 8 Aug 2025 12:10:06 +0200 Subject: [PATCH] update --- .../modular_pipelines/flux/before_denoise.py | 25 ++++--------------- .../modular_pipelines/flux/decoders.py | 11 +++----- .../modular_pipelines/flux/denoise.py | 13 ++++------ .../modular_pipelines/flux/encoders.py | 4 +-- 4 files changed, 15 insertions(+), 38 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index ffc77bb24f..c51061b78f 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -21,7 +21,7 @@ import torch 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 @@ class FluxInputStep(PipelineBlock): 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 @@ class FluxInputStep(PipelineBlock): return components, state -class FluxSetTimestepsStep(PipelineBlock): +class FluxSetTimestepsStep(ModularPipelineBlocks): model_name = "flux" @property @@ -235,17 +230,12 @@ class FluxSetTimestepsStep(PipelineBlock): 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 @@ class FluxSetTimestepsStep(PipelineBlock): return components, state -class FluxPrepareLatentsStep(PipelineBlock): +class FluxPrepareLatentsStep(ModularPipelineBlocks): model_name = "flux" @property @@ -314,11 +304,6 @@ class FluxPrepareLatentsStep(PipelineBlock): 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", diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 8d561d38c6..846549b1a3 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/src/diffusers/modular_pipelines/flux/decoders.py @@ -22,7 +22,7 @@ from ...configuration_utils import FrozenDict 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 @@ class FluxDecodeStep(PipelineBlock): 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 diff --git a/src/diffusers/modular_pipelines/flux/denoise.py b/src/diffusers/modular_pipelines/flux/denoise.py index c4619c17fb..062f1ad49f 100644 --- a/src/diffusers/modular_pipelines/flux/denoise.py +++ b/src/diffusers/modular_pipelines/flux/denoise.py @@ -22,7 +22,7 @@ from ...utils import logging from ..modular_pipeline import ( BlockState, LoopSequentialPipelineBlocks, - PipelineBlock, + ModularPipelineBlocks, PipelineState, ) from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam @@ -32,7 +32,7 @@ from .modular_pipeline import FluxModularPipeline logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class FluxLoopDenoiser(PipelineBlock): +class FluxLoopDenoiser(ModularPipelineBlocks): model_name = "flux" @property @@ -49,11 +49,8 @@ class FluxLoopDenoiser(PipelineBlock): @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 @@ class FluxLoopDenoiser(PipelineBlock): return components, block_state -class FluxLoopAfterDenoiser(PipelineBlock): +class FluxLoopAfterDenoiser(ModularPipelineBlocks): model_name = "flux" @property @@ -175,7 +172,7 @@ class FluxDenoiseLoopWrapper(LoopSequentialPipelineBlocks): ] @property - def loop_intermediate_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> List[InputParam]: return [ InputParam( "timesteps", diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py index 9bf2f54eec..9681c1e9d9 100644 --- a/src/diffusers/modular_pipelines/flux/encoders.py +++ b/src/diffusers/modular_pipelines/flux/encoders.py @@ -21,7 +21,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokeniz 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