diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index f328078ce4..c36c0c31ea 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -133,6 +133,7 @@ def _register_attention_processors_metadata(): skip_processor_output_fn=_skip_proc_output_fn_Attention_WanAttnProcessor2_0, ), ) + # FluxAttnProcessor AttentionProcessorRegistry.register( model_class=FluxAttnProcessor, diff --git a/src/diffusers/hooks/utils.py b/src/diffusers/hooks/utils.py new file mode 100644 index 0000000000..c5260eeebe --- /dev/null +++ b/src/diffusers/hooks/utils.py @@ -0,0 +1,43 @@ +# 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. + +import torch + +from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS, _ATTENTION_CLASSES, _FEEDFORWARD_CLASSES + + +def _get_identifiable_transformer_blocks_in_module(module: torch.nn.Module): + module_list_with_transformer_blocks = [] + for name, submodule in module.named_modules(): + name_endswith_identifier = any(name.endswith(identifier) for identifier in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS) + is_modulelist = isinstance(submodule, torch.nn.ModuleList) + if name_endswith_identifier and is_modulelist: + module_list_with_transformer_blocks.append((name, submodule)) + return module_list_with_transformer_blocks + + +def _get_identifiable_attention_layers_in_module(module: torch.nn.Module): + attention_layers = [] + for name, submodule in module.named_modules(): + if isinstance(submodule, _ATTENTION_CLASSES): + attention_layers.append((name, submodule)) + return attention_layers + + +def _get_identifiable_feedforward_layers_in_module(module: torch.nn.Module): + feedforward_layers = [] + for name, submodule in module.named_modules(): + if isinstance(submodule, _FEEDFORWARD_CLASSES): + feedforward_layers.append((name, submodule)) + return feedforward_layers diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py index 7211fb5693..124e611bd0 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux.py @@ -310,7 +310,7 @@ class FluxPipeline( def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control.py b/src/diffusers/pipelines/flux/pipeline_flux_control.py index 5a057f94cf..51d6ecbe31 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control.py @@ -324,7 +324,7 @@ class FluxControlPipeline( def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py index 8d5439daf6..c61d46daef 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py @@ -335,7 +335,7 @@ class FluxControlImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSin def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py index 872bcf177c..3de636361b 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py @@ -374,7 +374,7 @@ class FluxControlInpaintPipeline( def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py index 1438d4a902..a39b9c9ce2 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py @@ -341,7 +341,7 @@ class FluxControlNetPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleF def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py index 52e15de53b..582c7bbad8 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py @@ -335,7 +335,7 @@ class FluxControlNetImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, From def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py index d1e874d0b8..f7f34ef231 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py @@ -346,7 +346,7 @@ class FluxControlNetInpaintPipeline(DiffusionPipeline, FluxLoraLoaderMixin, From def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_fill.py b/src/diffusers/pipelines/flux/pipeline_flux_fill.py index ddfb284eaf..d50db407a8 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_fill.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_fill.py @@ -419,7 +419,7 @@ class FluxFillPipeline( def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py index 1c4cf3b1cd..08e2f12778 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py @@ -333,7 +333,7 @@ class FluxImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py index eeacd9b19b..0494146693 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py @@ -337,7 +337,7 @@ class FluxInpaintPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FluxIPAdapterM def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py index 3c78aeaf36..ce2941f3dd 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py @@ -358,7 +358,7 @@ class FluxKontextPipeline( def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py index 6dc621901c..56a5e934a4 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py @@ -391,7 +391,7 @@ class FluxKontextInpaintPipeline( def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py index b5ccfb31a3..e79db337b2 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py @@ -292,7 +292,7 @@ class FluxPriorReduxPipeline(DiffusionPipeline): def encode_prompt( self, prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], + prompt_2: Optional[Union[str, List[str]]] = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.FloatTensor] = None,