[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.

update

2904639

sayakpaul reviewed May 13, 2025

View reviewed changes

yiyixuxu mentioned this pull request May 15, 2025

[modular diffusers] more refactor #11235

Open

18 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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

DN6 commented May 12, 2025 •

edited

Loading

sayakpaul left a comment

sayakpaul May 13, 2025

sayakpaul May 13, 2025

DN6 May 15, 2025

sayakpaul May 15, 2025

sayakpaul May 13, 2025

DN6 May 15, 2025

HuggingFaceDocBuilderDev commented May 15, 2025

DN6 commented May 15, 2025

sayakpaul commented May 15, 2025 •

edited

Loading

sayakpaul May 15, 2025

sayakpaul May 15, 2025

sayakpaul May 15, 2025

sayakpaul May 15, 2025

yiyixuxu 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

sayakpaul commented May 16, 2025

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

Are you sure you want to change the base?

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

Conversation

DN6 commented May 12, 2025 • edited Loading

What does this PR do?

Before submitting

Who can review?

sayakpaul left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented May 15, 2025

DN6 commented May 15, 2025

sayakpaul commented May 15, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yiyixuxu 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

sayakpaul commented May 16, 2025

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