# Copyright 2025 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Any, List, Tuple import torch from ...models import FluxTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ..modular_pipeline import ( BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState, ) from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import FluxModularPipeline logger = logging.get_logger(__name__) # pylint: disable=invalid-name class FluxLoopDenoiser(ModularPipelineBlocks): model_name = "flux" @property def expected_components(self) -> List[ComponentSpec]: return [ComponentSpec("transformer", FluxTransformer2DModel)] @property def description(self) -> str: return ( "Step within the denoising loop that denoise the latents. " "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " "object (e.g. `FluxDenoiseLoopWrapper`)" ) @property def inputs(self) -> List[Tuple[str, Any]]: return [ InputParam("joint_attention_kwargs"), 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.", ), InputParam( "guidance", required=True, type_hint=torch.Tensor, description="Guidance scale as a tensor", ), InputParam( "prompt_embeds", required=True, type_hint=torch.Tensor, description="Prompt embeddings", ), InputParam( "pooled_prompt_embeds", required=True, type_hint=torch.Tensor, description="Pooled prompt embeddings", ), InputParam( "text_ids", required=True, type_hint=torch.Tensor, description="IDs computed from text sequence needed for RoPE", ), InputParam( "latent_image_ids", required=True, type_hint=torch.Tensor, description="IDs computed from image sequence needed for RoPE", ), # TODO: guidance ] @torch.no_grad() def __call__( self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor ) -> PipelineState: noise_pred = components.transformer( hidden_states=block_state.latents, timestep=t.flatten() / 1000, guidance=block_state.guidance, encoder_hidden_states=block_state.prompt_embeds, pooled_projections=block_state.pooled_prompt_embeds, joint_attention_kwargs=block_state.joint_attention_kwargs, txt_ids=block_state.text_ids, img_ids=block_state.latent_image_ids, return_dict=False, )[0] block_state.noise_pred = noise_pred return components, block_state class FluxLoopAfterDenoiser(ModularPipelineBlocks): model_name = "flux" @property def expected_components(self) -> List[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property 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` " "object (e.g. `FluxDenoiseLoopWrapper`)" ) @property def inputs(self) -> List[Tuple[str, Any]]: return [] @property def intermediate_inputs(self) -> List[str]: return [InputParam("generator")] @property def intermediate_outputs(self) -> List[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] @torch.no_grad() def __call__(self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): # Perform scheduler step using the predicted output latents_dtype = block_state.latents.dtype block_state.latents = components.scheduler.step( block_state.noise_pred, t, block_state.latents, return_dict=False, )[0] if block_state.latents.dtype != latents_dtype: block_state.latents = block_state.latents.to(latents_dtype) return components, block_state class FluxDenoiseLoopWrapper(LoopSequentialPipelineBlocks): model_name = "flux" @property def description(self) -> str: return ( "Pipeline block that iteratively denoise the latents over `timesteps`. " "The specific steps with each iteration can be customized with `sub_blocks` attributes" ) @property def loop_expected_components(self) -> List[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ComponentSpec("transformer", FluxTransformer2DModel), ] @property def loop_inputs(self) -> List[InputParam]: return [ InputParam( "timesteps", required=True, type_hint=torch.Tensor, description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", ), InputParam( "num_inference_steps", required=True, type_hint=int, description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", ), ] @torch.no_grad() def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) block_state.num_warmup_steps = max( len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 ) # We set the index here to remove DtoH sync, helpful especially during compilation. # Check out more details here: https://github.com/huggingface/diffusers/pull/11696 components.scheduler.set_begin_index(0) with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: for i, t in enumerate(block_state.timesteps): components, block_state = self.loop_step(components, block_state, i=i, t=t) if i == len(block_state.timesteps) - 1 or ( (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 ): progress_bar.update() self.set_block_state(state, block_state) return components, state class FluxDenoiseStep(FluxDenoiseLoopWrapper): block_classes = [FluxLoopDenoiser, FluxLoopAfterDenoiser] block_names = ["denoiser", "after_denoiser"] @property def description(self) -> str: return ( "Denoise step that iteratively denoise the latents. \n" "Its loop logic is defined in `FluxDenoiseLoopWrapper.__call__` method \n" "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n" " - `FluxLoopDenoiser`\n" " - `FluxLoopAfterDenoiser`\n" "This block supports both text2image and img2img tasks." )