Skip to content

Conversation

@sayakpaul
Copy link
Member

@sayakpaul sayakpaul commented Sep 12, 2025

What does this PR do?

As discussed internally, this PR introduces pipeline-specific Mixin classes that hold common methods shared across different task-specific pipelines. The actual pipelines (such as StableDiffusionPipeline, QwenImagePipeline, QwenImageImg2ImgPipeline, etc.) would then subclass from these mixins.

Examples of such common methods vary from pipeline to pipeline, but some examples include encode_prompt(), properties such as guidance_scale, etc.

Additionally, we have a couple of methods like retrieve_latents() and retrieve_timesteps() originated in stable_diffusion and we copied them over to other pipelines that use these methods. Same for calculate_shift that originated in flux. This PR also considers this situation and follows a reasonable approach to treat it (more below).

Note

This PR introduces Mixin classes for four popular pipelines: Flux, Qwen, SDXL, and SD.


Guiding principles

A sample structure for src/diffusers/pipelines/flux:

.
├── __init__.py
├── modeling_flux.py
├── pipeline_flux_control_img2img.py
├── pipeline_flux_control_inpaint.py
├── pipeline_flux_control.py
├── pipeline_flux_controlnet_image_to_image.py
├── pipeline_flux_controlnet_inpainting.py
├── pipeline_flux_controlnet.py
├── pipeline_flux_fill.py
├── pipeline_flux_img2img.py
├── pipeline_flux_inpaint.py
├── pipeline_flux_kontext_inpaint.py
├── pipeline_flux_kontext.py
├── pipeline_flux_prior_redux.py
├── pipeline_flux_utils.py <------ holds the Mixin
├── pipeline_flux.py
└── pipeline_output.py

pipeline_flux_utils.py here has the Mixin class, holding the common methods. It additionally has the calculate_shift() method, which is imported by the other pipelines under src/diffusers/pipelines/flux. Those pipelines won't have to maintain # Copied from ... versions of the calculate_shift() method anymore. They can just import like from .pipeline_flux_utils import calculate_shift. This is how we reduce LOCs, further on an intra-pipeline basis.

However, if a new pipeline (other than Flux), wants to leverage calculate_shift(), it should still follow # Copied from ... versions of the calculate_shift() instead of a direct import. This way, we can let individual pipeline-level methods evolve independently without introducing complex dependency patterns. For example, pipeline_flux_utils.py uses the # Copied from ... versions of methods like retrieve_latents().

TODO

Move the pipeline-specific LoRA loader classes to their respective pipeline modules.


Notes

  • In the Qwen mixin, kept encode_prompt() varying while unifying the _get_qwen_prompt_embeds() method in QwenImageMixin. Comments are in line.
  • There are a few methods that are shared by 2/3 pipelines such as prepare_image() is shared by QwenImageControlNetInpaintPipeline and QwenImageControlNetPipeline but not others. We can choose to move them to pipeline_qwen_utils.py but I am not strongly opinionated on this. The same applies to methods like prepare_mask_latents(). My preference here is to keep them as is for now.

Took similar principles to do FluxMixin as well. I looked for methods and properties shared by all the Flux family of pipelines first. I moved them to FluxMixin. And then I looked for a clear pattern of methods that are shared by a particular class of Flux pipelines. prepare_image() fit the bill, and decided to move it to FluxControlMixin(FluxMixin). This method is used by all Flux Control pipelines (including ControlNets).

Another alternative could be to separate w.r.t text-to-image, image-to-image, etc., as we do in modular. But I think the separation should be such that the Mixin class is able to accommodate methods that are shared by all its children. I am not too strongly opinionated about this. So, open to suggestions.


Questions

We have

class StableDiffusionMixin:

This PR also has a similar Mixin class but without the methods from the above Mixin. To avoid confusion, I have named it SDMixin.

WDYT about collating the two classes (StableDiffusionMixin and SDMixin (introduced in this PR)) into one?

@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Copy link
Collaborator

@DN6 DN6 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking good 👍🏽 Still on the fence on whether we use multiple levels of inheritance with QwenImageMixin. But other than that this looks good to me.

Perhaps we try with to do the same with Flux and SDXL to see if some common patterns show up for these Mixins.

@sayakpaul
Copy link
Member Author

Looking good 👍🏽 Still on the fence on whether we use multiple levels of inheritance with QwenImageMixin. But other than that this looks good to me.

With you. I am also on the fence. Since the mixins for QwenImage and QwenImageEdit primarily differ in encode_prompt(), I decided to design it like that. We can still rewrite the encode_prompt() method in a single mixin so that it can deal with the images, but it could complicate things. It will also end up changing the signature of the method, leaving room for breaking changes.

But we can definitely uniformize the _get_qwen_prompt_embeds() method as it's private anyway. WDYT?

Perhaps we try with to do the same with Flux and SDXL to see if some common patterns show up for these Mixins.

Yes, I will work on that.


return split_result

def _get_qwen_prompt_embeds(
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Decided to unify it to further reduce LOC. It's a private method anyway, so, we should be safe.



class QwenImagePipelineMixin(QwenImageMixin):
def encode_prompt(
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only differing encode_prompt() to not introduce any breaking change.

@sayakpaul sayakpaul changed the title [wip] unbloat QwenImage pipelines by introducing a pipeline-specific mixins to hold common methods [wip] unbloat QwenImage and Flux pipelines by introducing a pipeline-specific mixins to hold common methods Sep 23, 2025
@sayakpaul
Copy link
Member Author

@DN6 pushed updates for Flux. LMK what you think.

@sayakpaul sayakpaul requested a review from DN6 September 23, 2025 06:29
@sayakpaul sayakpaul changed the title [wip] unbloat QwenImage and Flux pipelines by introducing a pipeline-specific mixins to hold common methods [wip] [core] introduce pipeline-specific mixins to hold common methods Nov 27, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants