Skip to content

[Modular] Updates for Custom Pipeline Blocks #11940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 18, 2025
Merged

[Modular] Updates for Custom Pipeline Blocks #11940

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a633289
update
DN6 Jul 16, 2025
5c43924
update
DN6 Jul 17, 2025
532b395
update
DN6 Jul 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/diffusers/commands/custom_blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from . import BaseDiffusersCLICommand


EXPECTED_PARENT_CLASSES = ["ModularPipelineBlocks"]
EXPECTED_PARENT_CLASSES = ["PipelineBlock"]
Copy link
Collaborator

Choose a reason for hiding this comment

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

why it's only PipelineBlock?

Copy link
Collaborator

Choose a reason for hiding this comment

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

users should be able to define custom ModularPipelineBlocks too where they assmeble their own presets so that it is easier to use

in diff-diff remote code example, https://huggingface.co/YiYiXu/modular-diffdiff-0704/blob/main/config.json

Copy link
Collaborator

Choose a reason for hiding this comment

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

maybe support all block types?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Ah. Yeah that make sense will update.

CONFIG = "config.json"


Expand Down
37 changes: 34 additions & 3 deletions src/diffusers/modular_pipelines/modular_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,14 +351,17 @@ def from_pretrained(
"token",
]
hub_kwargs = {name: kwargs.pop(name) for name in hub_kwargs_names if name in kwargs}
hub_kwargs.update({"trust_remote_code": trust_remote_code})

config = cls.load_config(pretrained_model_name_or_path)
has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"]
trust_remote_code = resolve_trust_remote_code(
trust_remote_code, pretrained_model_name_or_path, has_remote_code
)
if not (has_remote_code and trust_remote_code):
raise ValueError("TODO")
raise ValueError(
"Selected model repository does not happear to have any custom code or does not have a valid `config.json` file."
)

class_ref = config["auto_map"][cls.__name__]
module_file, class_name = class_ref.split(".")
Expand All @@ -367,7 +370,6 @@ def from_pretrained(
pretrained_model_name_or_path,
module_file=module_file,
class_name=class_name,
is_modular=True,
**hub_kwargs,
**kwargs,
)
Expand All @@ -384,10 +386,39 @@ def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):

full_mod = type(self).__module__
module = full_mod.rsplit(".", 1)[-1].replace("__dynamic__", "")
parent_module = self.save_pretrained.__func__.__qualname__.split(".", 1)[0]
parent_module = self.__class__.__bases__[0].__name__
auto_map = {f"{parent_module}": f"{module}.{cls_name}"}

self.register_to_config(auto_map=auto_map)

_component_specs = {spec.name: deepcopy(spec) for spec in self.expected_components}
_config_specs = {spec.name: deepcopy(spec) for spec in self.expected_configs}

register_components_dict = {}
for name, component_spec in _component_specs.items():
if component_spec.type_hint is not None:
lib_name, cls_name = _fetch_class_library_tuple(component_spec.type_hint)
else:
lib_name = cls_name = None
load_spec_dict = {k: getattr(component_spec, k) for k in component_spec.loading_fields()}

# Since ModularPipelineBlocks can never have loaded components we set
# first two fields in the config dict to None
register_components_dict[name] = (
None,
None,
Copy link
Collaborator Author

@DN6 DN6 Jul 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

@yiyixuxu Here for serializing PipelineBlocks, we deliberately set the library, class tuple to None. But type_hint also provides this information. Do we need to have both of these in the config then? Will modular configs ever have the class library tuple be non null values?

Copy link
Collaborator

Choose a reason for hiding this comment

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

I don't think we need to set them to None, they should just be the actual library/class

but, yes that's a good question, let's try to find a solution here:

we do not need these two fields technically, I left it there because that's I want to have an intuitive way to be able to tell if a components are loaded or not, in DiffusionPipeline these tuples will be null values if not loaded so I did same with modular.

maybe we can keep these 2 fields (or something else that can indicate whehter they are loaded or not) in config dict ( which gets print out when you do print(modular_pipeline) ), but remove them when saving to the modular_model_index.json?

let me know what you think!

Copy link
Collaborator

Choose a reason for hiding this comment

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

ohh actually, do you include the components info in the config.json for pipeline blocks too?
I was thinking maybe we don't need this, because pipeline blocks are just definations, so it does not need to include the loading related info - these are already saved in the pipeline, https://huggingface.co/YiYiXu/modular-diffdiff-0704

Copy link
Collaborator

Choose a reason for hiding this comment

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

updates ModularPipelineBlock serialization so that type hint, repo, subfolder, variant information gets saved when calling block.save_pretrained. Useful for when you have custom pipeline blocks that have components that use from_pretrained

so not sure if I understand here, I think the block's config.json is only used to create block, the auto_map points to the code defination is enough becausee you can define the default loading spec inside the custom code; user may want to change the loading spec at run time but they need to do that with the ModuarPipeline's config, not the blocks, no?

Copy link
Collaborator Author

@DN6 DN6 Jul 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Hmm so what about custom pipeline blocks that are meant to work with a particular model type? e.g Trying to create a custom block here to work with the Florence2-large model.
https://huggingface.co/diffusers-internal-dev/florence2-image-annotator/blob/main/config.json

How should one structure this? The goal is to insert a step that automatically creates a mask for an image
for inpainting based on a text prompt.

import torch
from diffusers.modular_pipelines import ModularPipelineBlocks, SequentialPipelineBlocks
from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS
from diffusers.utils import load_image

# fetch the Florence2 image annotator block that will create our mask
image_annotator_block = ModularPipelineBlocks.from_pretrained("diffusers-internal-dev/florence2-image-annotator", trust_remote_code=True)

my_blocks = INPAINT_BLOCKS.copy()

# insert the annotation block before the image encoding step
my_blocks.insert("image_annotator", image_annotator_block, 1)

# Create our initial set of inpainting blocks
blocks = SequentialPipelineBlocks.from_blocks_dict(my_blocks)
pipe = blocks.init_pipeline()

If the component spec information like repo/subfolder for the model is not present in the config, the init_pipeline call won't work?

Copy link
Collaborator

@yiyixuxu yiyixuxu Jul 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

so I see you already defined here https://huggingface.co/diffusers-internal-dev/florence2-image-annotator/blob/main/block.py#L17

so once you convert the blocks into pipeline, it will show up in pipeline's config like this

>>> image_annotator_block.init_pipeline()
ModularPipeline {
  "_blocks_class_name": "Florence2ImageAnnotatorBlock",
  "_class_name": "ModularPipeline",
  "_diffusers_version": "0.35.0.dev0",
  "image_annotator": [
    null,
    null,
    {
      "repo": "microsoft/Florence-2-large",
      "revision": null,
      "subfolder": "",
      "type_hint": [
        "transformers",
        "AutoModelForCausalLM"
      ],
      "variant": null
    }
  ],
  "image_annotator_processor": [
    null,
    null,
    {
      "repo": "microsoft/Florence-2-large",
      "revision": null,
      "subfolder": "",
      "type_hint": [
        "transformers",
        "AutoProcessor"
      ],
      "variant": null
    }
  ]
}

and if you run save_pretrained() on the pipeline, you will get the modular_model_index.json contains the loading specs for the cusom model

pipe = image_annotator_block.init_pipeline()
pipe.save_pretraind(...)

e.g. this is what I got by just adding pipe.save_pretrained("YiYiXU/test_image_annotator", push_to_hub=True) to the code example you shared https://huggingface.co/YiYiXu/test_image_annotator

ideally, we also push the custom code in the same repo, so everything can be together with ModularPipeline.from_pretrained() (the custom blocks + loading configs -> you got a fully initialized custom pipeline)

and since everything is in the same folder, the author of the custom pipeline also have the option to fill out the repo/subfolder info in modular_model_index.json if they like (instead of defining it in the code)

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Ah I see it is expected to load from the custom code itself. Okay, so I made changes to ModularPipelineBlocks to allow running init_pipeline lmk if those look okay.

Re: Having all components + custom code in the same repo. Assuming we load two custom blocks from two different repos, each has a block.py file defining the custom code. If we serialize the Pipeline with code into one repo, how do we figure out where to put all the components? Also how do we handle components with custom code in modular_model_index.json?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Re: (None, None, <config details>) format in modular_model_index.json and this

maybe we can keep these 2 fields (or something else that can indicate whehter they are loaded or not) in config dict ( which gets print out when you do print(modular_pipeline) ), but remove them when saving to the modular_model_index.json?

I think this could work. I can put together a quick PR to just see how it feels.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Re: Having all components + custom code in the same repo. Assuming we load two custom blocks from two different repos, each has a block.py file defining the custom code. If we serialize the Pipeline with code into one repo, how do we figure out where to put all the components? Also how do we handle components with custom code in modular_model_index.json?

For the blocks author, they can define the default component loading specs (repo/subfolder etc) when they write the pipeline blocks code, in the expected_components field. It's optional to include the repo info, but if they did that, like you did here with the florence-image-annotator block https://huggingface.co/diffusers-internal-dev/florence2-image-annotator/blob/main/block.py#L17 - these info info will stick with the block wherever it goes forever. lol

so in your question, if you load two different custom blocks from two different repos and mix them together into one pipeline, the loading specs for these custom components from both blocks should show up as default in your final pipeline. If you save the final pipeline into a repo, it will serialize the pipeline config dict into a modular_model_index.json in that repo, which will contain the default loading spec for the custom components, but the person who put together the final pipeline can edit the modular_model_index.json if they want to use something else. The end user could use a different loading specs at run time as well

how do we handle components with custom code in modular_model_index.json?

not sure I understand the question correctly, but ModularPipeline.from_pretrained() will look for custom code first, so if you save the code in the same repo, it should work
the code is here:
https://github.com/huggingface/diffusers/blame/main/src/diffusers/modular_pipelines/modular_pipeline.py#L2053

if your pipeline uses two different blocks from two different repos, I imagine your repo also need a block.py, and in that block.py you import these custom code, so I think it would also work lol but have not tested with such cases so let me know

Copy link
Collaborator

@yiyixuxu yiyixuxu Jul 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

and if you only want to host the block in your repo, (not the entire pipeline and do not want to save a modular_model_index.json) - defining the custom componenet in the code is enough,

I think your current code alreaedy work out of box very well: if the user just want to import that custom block, mix into their existing inpanting pipeline, they can pass the inpaiting repo in init_pipeline() - the final pipeline will contain the loading info for the custom component (since we defined a default) and all the loading info for other components (fetchd from the modular_model_index.json in that inpainting repo we passed)

e.g.
your code, only difference is I passed a "YiYiXu/modular-demo-auto" to init_pipeline() - "YiYiXu/modular-demo-auto" does not contain the custom component but we have a default defined so pipeline still have that info

import torch
from diffusers.modular_pipelines import ModularPipelineBlocks, SequentialPipelineBlocks
from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS
from diffusers.utils import load_image

# fetch the Florence2 image annotator block that will create our mask
from diffusers.modular_pipelines import PipelineBlock
image_annotator_block = ModularPipelineBlocks.from_pretrained("diffusers-internal-dev/florence2-image-annotator", trust_remote_code=True)

my_blocks = INPAINT_BLOCKS.copy()

# insert the annotation block before the image encoding step
my_blocks.insert("image_annotator", image_annotator_block, 1)

# Create our initial set of inpainting blocks
blocks = SequentialPipelineBlocks.from_blocks_dict(my_blocks)
pipe = blocks.init_pipeline("YiYiXu/modular-demo-auto")
print(pipe)

the pipeline is ready to use and knows how to load all the components

StableDiffusionXLModularPipeline {
  "_blocks_class_name": "SequentialPipelineBlocks",
  "_class_name": "StableDiffusionXLModularPipeline",
  "_diffusers_version": "0.35.0.dev0",
  "force_zeros_for_empty_prompt": true,
  "image_annotator": [
    null,
    null,
    {
      "repo": "microsoft/Florence-2-large",
      "revision": null,
      "subfolder": "",
      "type_hint": [
        "transformers",
        "AutoModelForCausalLM"
      ],
      "variant": null
    }
  ],
  "image_annotator_processor": [
    null,
    null,
    {
      "repo": "microsoft/Florence-2-large",
      "revision": null,
      "subfolder": "",
      "type_hint": [
        "transformers",
        "AutoProcessor"
      ],
      "variant": null
    }
  ],
  "requires_aesthetics_score": false,
  "scheduler": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "scheduler",
      "type_hint": [
        "diffusers",
        "EulerDiscreteScheduler"
      ],
      "variant": null
    }
  ],
  "text_encoder": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "text_encoder",
      "type_hint": [
        "transformers",
        "CLIPTextModel"
      ],
      "variant": null
    }
  ],
  "text_encoder_2": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "text_encoder_2",
      "type_hint": [
        "transformers",
        "CLIPTextModelWithProjection"
      ],
      "variant": null
    }
  ],
  "tokenizer": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "tokenizer",
      "type_hint": [
        "transformers",
        "CLIPTokenizer"
      ],
      "variant": null
    }
  ],
  "tokenizer_2": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "tokenizer_2",
      "type_hint": [
        "transformers",
        "CLIPTokenizer"
      ],
      "variant": null
    }
  ],
  "unet": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "unet",
      "type_hint": [
        "diffusers",
        "UNet2DConditionModel"
      ],
      "variant": null
    }
  ],
  "vae": [
    null,
    null,
    {
      "repo": "SG161222/RealVisXL_V4.0",
      "revision": null,
      "subfolder": "vae",
      "type_hint": [
        "diffusers",
        "AutoencoderKL"
      ],
      "variant": null
    }
  ]
}

{
"type_hint": (lib_name, cls_name),
**load_spec_dict,
},
)
self.register_to_config(**register_components_dict)

default_configs = {}
for name, config_spec in _config_specs.items():
default_configs[name] = config_spec.default
self.register_to_config(**default_configs)

self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
config = dict(self.config)
self._internal_dict = FrozenDict(config)
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/modular_pipelines/modular_pipeline_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ class ComponentSpec:
config: Optional[FrozenDict] = None
# YiYi Notes: should we change it to pretrained_model_name_or_path for consistency? a bit long for a field name
repo: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": True})
subfolder: Optional[str] = field(default=None, metadata={"loading": True})
subfolder: Optional[str] = field(default="", metadata={"loading": True})
variant: Optional[str] = field(default=None, metadata={"loading": True})
revision: Optional[str] = field(default=None, metadata={"loading": True})
default_creation_method: Literal["from_config", "from_pretrained"] = "from_pretrained"
Expand Down