[WIP] Modular Diffusers support custom code/pipeline blocks #11539

DN6 · 2025-05-12T14:16:41Z

What does this PR do?

Add support for loading custom pipeline blocks with Modular Diffusers. PR is still in very rough shape, but is functional.

Snippet to test

from diffusers.pipelines.modular_pipeline import PipelineBlock

block = PipelineBlock.from_pretrained(
        "diffusers-internal-dev/modular-depth-block", trust_remote_code=True
)

Note I think the formatting changes might have been because of a difference in ruff versions.

TODOs:

Possibly move logic to fetch custom modules into AutoPipelineBlock rather than have it in ModularPipelineMixin
Clean up custom code fetching in get_class_from_dynamic_module. Probably don't need the is_modular argument.
Remove get_class_in_modular_module
trust_remote_code=None is currently broken. Need to fix.
Add logic to save custom pipeline blocks
Add support for local custom code.
Add tests

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

sayakpaul

Thanks, Dhruv!

I couldn't run the code as there seems to be some problems with undefined variables. I can give the saving logic a look after that is fixed.

sayakpaul · 2025-05-13T03:37:30Z

src/diffusers/utils/dynamic_modules_utils.py

@@ -154,15 +157,132 @@ def check_imports(filename):
    return get_relative_imports(filename)


-def get_class_in_module(class_name, module_path):
+def resolve_trust_remote_code(trust_remote_code, model_name, has_local_code, has_remote_code):


(nit): Would prefer this to be a private method.

src/diffusers/utils/dynamic_modules_utils.py

sayakpaul · 2025-05-13T03:41:18Z

src/diffusers/utils/dynamic_modules_utils.py

+                        f"The repository for {model_name} contains custom code which must be executed to correctly "
+                        f"load the model. You can inspect the repository content at https://hf.co/{model_name}.\n"
+                        f"You can avoid this prompt in future by passing the argument `trust_remote_code=True`.\n\n"
+                        f"Do you wish to run the custom code? [y/N] "
+                    )
+                    if answer.lower() in ["yes", "y", "1"]:
+                        trust_remote_code = True
+                    elif answer.lower() in ["no", "n", "0", ""]:
+                        trust_remote_code = False
+                signal.alarm(0)


If it's just about passing trust_remote_code=Yes, would it be too much to just enforce that and avoid signal altogether?

Could be a case where as a user you don't know it's custom code. This would open a terminal prompt that asks you to confirm whether or not you want to proceed loading the code.

No strong opinions but my preference would be to just tell the user how they can resolve that issue with a sensible message instead of opening a terminal to get their choice.

sayakpaul · 2025-05-13T03:44:16Z

src/diffusers/utils/dynamic_modules_utils.py

@@ -37,6 +39,7 @@

 # See https://huggingface.co/datasets/diffusers/community-pipelines-mirror
 COMMUNITY_PIPELINES_MIRROR_ID = "diffusers/community-pipelines-mirror"
+_HF_REMOTE_CODE_LOCK = threading.Lock()


Why do we have to acquire a lock?

It's used in the step that modifies sys.modules. Assuming it is to avoid having background threads concurrently read/write from there during that operation. It's taken from transformers custom code loading.

src/diffusers/utils/dynamic_modules_utils.py

HuggingFaceDocBuilderDev · 2025-05-15T07:16:05Z

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.

DN6 · 2025-05-15T09:16:47Z

@yiyixuxu PR is ready for an initial review. Lmk if it's fine for this logic to live in ModularPipelineMixin or if it's more appropriate in AutoPipelineBlocks (IMO feels better for it to be here).

Apologies for the formatting issues 🙈 I can try to remove them if it's too much.

I have not added saving custom modules yet. Think @sayakpaul Wanted to take a look at it.

sayakpaul · 2025-05-15T10:16:07Z

Yes, I will take a look into saving.

sayakpaul · 2025-05-15T10:16:30Z

src/diffusers/pipelines/modular_pipeline.py

 if is_accelerate_available():
-    import accelerate
+    pass


(nit): remove.

sayakpaul · 2025-05-15T10:18:24Z

src/diffusers/pipelines/modular_pipeline.py

+            trust_remote_code, pretrained_model_name_or_path, has_remote_code
+        )
+        if not (has_remote_code and trust_remote_code):
+            raise ValueError("")


Suggested change

raise ValueError("")

raise ValueError("TODO")

sayakpaul · 2025-05-15T10:24:03Z

src/diffusers/utils/dynamic_modules_utils.py

    """
    Import a module on the cache directory for modules and extract a class from it.
    """
-    module_path = module_path.replace(os.path.sep, ".")
-    module = importlib.import_module(module_path)
+    name = os.path.normpath(module_path)


No strong opinions but currently there seems to be a mix of Path and os usage. Maybe it's necessary that way but if it's not I would prefer using one of the two.

sayakpaul · 2025-05-15T10:24:13Z

src/diffusers/utils/dynamic_modules_utils.py

@@ -454,4 +531,4 @@ def get_class_from_dynamic_module(
        revision=revision,
        local_files_only=local_files_only,
    )
-    return get_class_in_module(class_name, final_module.replace(".py", ""))
+    return get_class_in_module(class_name, final_module)


yiyixuxu · 2025-05-15T18:48:16Z

super super cool!
end-to-end, works with AutoPipelineBlocks and SequentialPipelineBlocks too

from diffusers.pipelines.modular_pipeline import ModularPipelineMixin
block = ModularPipelineMixin.from_pretrained(
        "YiYiXu/modular-depth-block", trust_remote_code=True
)

print(block)
block.setup_loader()
print(block.loader)
block.loader.load()



url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png"
depth_processed = block.run(url = url, output="image")
print(f"depth_processed.shape: {depth_processed.shape}")

yiyixuxu · 2025-05-15T20:03:35Z

Lmk if it's fine for this logic to live in ModularPipelineMixin or if it's more appropriate in AutoPipelineBlocks (IMO feels better for it to be here).

does that mean only AutoPipelineBlocks will be able to use remote code then? I think all types of blocks should be able to load from remote, no?

Maybe we should change ModularPipelineMixin to ModularPipelineBlock or ModularBlock or ModularPipeline? so people can just load with ModularPipeline.from_pretrained() for anything, similar to DiffusionPipeline.

Also in Modular system, Block and Pipeline pretty much mean the same thing, there is no clear difference between them. Should we pick up a term and stick to it? which one is better?

on a somewhat related note , should I change AutoPipelineBlocks to something else? The "Auto" in AutoPipelineBlocks has pretty different meaning from "auto" in AutoPipeline and AutoModel - it is similar in the sense that the specific class(es) is automatically selected, but it is the sub-blocks added to it that are selected

let me know what you think!

yiyixuxu · 2025-05-15T20:10:45Z

also, can you think a bit more how would this work with ModularLoader? I think there is some overlap here that can be refactored and simplified later, maybe we do not need a separate ModularLoader for the model components

for now, I think if I just change the config_name for ModularLoader into config.json it should already work with these remote pipeline repo to pick up the models specified there

yiyixuxu · 2025-05-15T20:28:17Z

let's merge this now? we can do follow-up changes in a new PR

I'm waiting for this one to get in first so I can merge in all the new changes into the main modular PR and it is a good place for a review (otherwise there might be a lot of merge conflicts for you to resolve)

sayakpaul · 2025-05-16T03:14:24Z

Sorry for chiming in

on a somewhat related note , should I change AutoPipelineBlocks to something else? The "Auto" in AutoPipelineBlocks has pretty different meaning from "auto" in AutoPipeline and AutoModel - it is similar in the sense that the specific class(es) is automatically selected, but it is the sub-blocks added to it that are selected

I think the essence is still similar to "automatic selection and configuration". Users are getting the blocks / classes they are expecting to get "automatically". So, I think it's fine.

DN6 · 2025-05-20T09:42:24Z

@yiyixuxu Merging this. I will think about your comments re: ModularLoader and Blocks vs Pipeline and open a follow up PR with some ideas.

yiyixuxu · 2025-05-20T19:30:58Z

@DN6 @sayakpaul
I put together a complete example here https://huggingface.co/YiYiXu/modular-diffdiff

sayakpaul · 2025-05-21T03:28:59Z

@yiyixuxu thanks! How were the files (config.json, modular_model_index.json) under https://huggingface.co/YiYiXu/modular-diffdiff/tree/main saved?

yiyixuxu · 2025-05-21T22:17:29Z

@sayakpaul I manually made the config.json since it is just one line

modular_model_index.json is saved with ModularLoader, I'm going to write more documentation on that ASAP there are some info here #11462

note that in modular model components save and loading are handled by the ModuarLoader class, at least for current design so the 'trust_remote_code` isn't necessarily because by default it only load/save code, I think something that directly push a pipeline block into hub would be more appropriate, like what's done in smolageents

https://github.com/huggingface/smolagents/blob/89d4e17a104e822a53e3072787f8a06544fcd7a0/src/smolagents/tools.py#L329

cc @DN6 here too, but don't need to do anything now I think, this works and it's sufficient for now

* add componentspec and configspec * up * up * move methods to blocks * Modular Diffusers Guiders (#11311) * cfg; slg; pag; sdxl without controlnet * support sdxl controlnet * support controlnet union * update * update * cfg zero* * use unwrap_module for torch compiled modules * remove guider kwargs * remove commented code * remove old guider * fix slg bug * remove debug print * autoguidance * smoothed energy guidance * add note about seg * tangential cfg * cfg plus plus * support cfgpp in ddim * apply review suggestions * refactor * rename enable/disable * remove cfg++ for now * rename do_classifier_free_guidance->prepare_unconditional_embeds * remove unused * [modular diffusers] introducing ModularLoader (#11462) * cfg; slg; pag; sdxl without controlnet --------- Co-authored-by: Aryan <[email protected]> * make loader optional * remove lora step and ip-adapter step -> no longer needed * rename pipeline -> components, data -> block_state * seperate controlnet step into input + denoise * refactor controlnet union * reefactor pipeline/block states so that it can dynamically accept kwargs * remove controlnet union denoise step, refactor & reuse controlnet denoisee step to accept aditional contrlnet kwargs * allow input_fields as input & update message * update input formating, consider kwarggs_type inputs with no name, e/g *_controlnet_kwargs * refactor the denoiseestep using LoopSequential! also add a new file for denoise step * change warning to debug * fix get_execusion blocks with loopsequential * fix auto denoise so all tests pass * update imports on guiders * remove modular reelated change from pipelines folder * made a modular_pipelines folder! * update __init__ * add notes * add block state will also make sure modifed intermediates_inputs will be updated * move block mappings to its own file * make inputs truly immutable, remove the output logic in sequential pipeline, and update so that intermediates_outputs are only new variables * decode block, if skip decoding do not need to update latent * fix imports * fix import * fix more * remove the output step * make generator intermediates (it is mutable) * after_denoise -> decoders * add a to-do for guider cconfig mixin * refactor component spec: replace create/create_from_pretrained/create_from_config to just create and load method * refactor modular loader: 1. load only load (pretrained components only if not specific names) 2. update acceept create spec 3. move the updte _componeent_spec logic outside register_components to each method that create/update the component: __init__/update/load * update components manager * up * [WIP] Modular Diffusers support custom code/pipeline blocks (#11539) * update * update * remove the duplicated components_manager file I forgot to deletee * fix import in block mapping * add a to-do for modular loader * prepare_latents_img2img pipeline method -> function, maybe do the same for others? * update input for loop blocks, do not need to include intermediate * solve merge conflict: manually add back the remote code change to modular_pipeline * add node_utils * modular node! * add * refator based on dhruv's feedbacks * update doc format for kwargs_type * up * updatee modular_pipeline.from_pretrained, modular_repo ->pretrained_model_name_or_path * save_pretrained for serializing config. (#11603) * save_pretrained for serializing config. * remove pushtohub * diffusers-cli rough --------- Co-authored-by: YiYi Xu <[email protected]> --------- Co-authored-by: Aryan <[email protected]> Co-authored-by: Dhruv Nair <[email protected]> Co-authored-by: Sayak Paul <[email protected]>

update

2904639

sayakpaul reviewed May 13, 2025

View reviewed changes

yiyixuxu mentioned this pull request May 15, 2025

[modular diffusers] more refactor #11235

Merged

29 tasks

update

e3bc60b

DN6 marked this pull request as ready for review May 15, 2025 09:12

DN6 requested a review from yiyixuxu May 15, 2025 09:16

sayakpaul reviewed May 15, 2025

View reviewed changes

DN6 merged commit 808dff0 into modular-refactor May 20, 2025
2 checks passed

sayakpaul deleted the modular-refactor-custom-blocks branch May 21, 2025 03:27

sayakpaul mentioned this pull request May 22, 2025

save_pretrained for serializing config. #11603

Merged

DN6 added the roadmap Add to current release roadmap label Jun 5, 2025

github-project-automation bot added this to Diffusers Roadmap 0.34 Jun 5, 2025

github-project-automation bot moved this to In Progress in Diffusers Roadmap 0.34 Jun 5, 2025

DN6 removed this from Diffusers Roadmap 0.34 Jun 5, 2025

[WIP] Modular Diffusers support custom code/pipeline blocks #11539

[WIP] Modular Diffusers support custom code/pipeline blocks #11539

Uh oh!

Conversation

DN6 commented May 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented May 15, 2025

Uh oh!

DN6 commented May 15, 2025

Uh oh!

sayakpaul commented May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yiyixuxu commented May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yiyixuxu commented May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yiyixuxu commented May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yiyixuxu commented May 15, 2025

Uh oh!

sayakpaul commented May 16, 2025

Uh oh!

DN6 commented May 20, 2025

Uh oh!

Uh oh!

yiyixuxu commented May 20, 2025

Uh oh!

sayakpaul commented May 21, 2025

Uh oh!

yiyixuxu commented May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

DN6 commented May 12, 2025 •

edited

Loading

sayakpaul commented May 15, 2025 •

edited

Loading

yiyixuxu commented May 15, 2025 •

edited

Loading

yiyixuxu commented May 15, 2025 •

edited

Loading

yiyixuxu commented May 15, 2025 •

edited

Loading

yiyixuxu commented May 21, 2025 •

edited

Loading