[modular] add auto_docstring & more doc related refactors #12958

yiyixuxu · 2026-01-10T02:36:39Z

This PR adds a utility script utils/modular_auto_docstring.py that automatically generates docstrings for modular pipeline block classes from their doc property.

Usage

Mark classes with # auto_docstring comment:

# auto_docstring
class QwenImageAutoVaeEncoderStep(AutoPipelineBlocks):
    block_classes = [QwenImageInpaintVaeEncoderStep, QwenImageImg2ImgVaeEncoderStep]
    block_names = ["inpaint", "img2img"]
    block_trigger_inputs = ["mask_image", "image"]

    @property
    def doc(self):
        return (
            "Vae encoder step that encodes image inputs into latent representations.\n"
            "This is an auto pipeline block.\n"
            " - `QwenImageInpaintVaeEncoderStep` (inpaint) is used when `mask_image` is provided.\n"
            " - `QwenImageImg2ImgVaeEncoderStep` (img2img) is used when `image` is provided."
        )

Run the script to insert docstrings:

python utils/modular_auto_docstring.py --fix_and_overwrite

HuggingFaceDocBuilderDev · 2026-01-10T02:44:15Z

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.

yiyixuxu · 2026-01-10T11:31:37Z

src/diffusers/modular_pipelines/modular_pipeline_utils.py

+    # ======================================================
+
+    @classmethod
+    def prompt(cls) -> "InputParam":


our pipeline parameter are pretty consistent across different pipelines, e.g. you always have prompt, height, width, num_inference_steps, etc. I made template for these common ones, so that it is easier to define

before you need

InputParam( name="prompt", type_hint=str, required=True, description="The prompt or prompts to guide image generation." )

now you do

InputParam.prompt() InputParam.height(default=1024) InputParam.num_inference_steps(default=28) InputParam.generator()

I'm a bit apprehensive about introducing dedicated class methods for common parameters in this way. I think the class can become quite large as common inputs expand.

I would prefer to keep current syntax (IMO this ensures InputParams are defined in a consistent way) and use post init on the dataclass to automatically add a description. e.g

# centralised descriptions would live somewhere like constants.py # can be used for modular + non-modular INPUT_PARAM_TEMPLATES = { "prompt": {"type_hint": str, "required": True, "description": "The prompt or prompts to guide image generation."}, "height": {"type_hint": int, "description": "The height in pixels of the generated image."}, "width": {"type_hint": int, "description": "The width in pixels of the generated image."}, "generator": {"type_hint": torch.Generator, "description": "Torch generator for deterministic generation."}, # ... } @dataclass class InputParam: name: str = None type_hint: Any = None required: bool = False default: Any = None description: str = None def __post_init__(self): if not self.name or self.name not in INPUT_PARAM_TEMPLATES: return template = INPUT_PARAM_TEMPLATES[self.name] if self.type_hint is None: self.type_hint = template.get("type_hint") if self.description is None: self.description = template.get("description")

If we feel that methods for these inputs are necessary, one way to address it without adding individual methods to the InputParam is to use a metaclass. It would result in the InputParam object being less crowded.

class InputParamMeta(type): def __getattr__(cls, name: str): if name in INPUT_PARAM_TEMPLATES: def factory(**overrides): return cls(name=name, **overrides) return factory raise AttributeError(f"No template named '{name}'") @dataclass class InputParam(metaclass=InputParamMeta):

I removed the class methods!
also moved the common parameter definitions out of InputParam/OutputParam as well. They're currently still in the same file, but we can move them elsewhere later.

To use a predefined param:

InputParam.template("image")

i went with this API, instead of the __post_init__ because you have to explicitly opt into using the template, even though it's slightly more verbose.

The main thing I wanted to avoid is with __post_init__ auto-filling, if you make a param that exists in our template but actually intend to customize its docstring later, e.g. you do something likeInputParam(name="image"), you will silently get the template defaults, and it is hard to notice.

With .template(), it's more clear when the user intends to use the pre-defined definition vs write a custom param. so for the same example, if someone writes InputParam(name="image") without filling in fields, the docstring will show "TODO: Add description".

we can also override fields in the template like this:

InputParam.template("image", note="resized") # appends a note to the description InputParam.template("image", required=False)

yiyixuxu · 2026-01-10T11:32:59Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py


+# auto_docstring
+class QwenImageAutoTextEncoderStep(AutoPipelineBlocks):
+    """


these are all auto generated docstring

yiyixuxu · 2026-01-10T11:33:19Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

 # ====================


+# auto_docstring


add this mark and then run

python utils/modular_auto_docstring.py --fix_and_overwrite

We could add this to make style and error out in the CI if it's not the case? 👀

This will help us enforce consistency.

yes, we should do that eventually!
(won't include in this PR though)

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

DN6 · 2026-01-14T07:57:37Z

src/diffusers/modular_pipelines/modular_pipeline_utils.py

+    # ======================================================
+
+    @classmethod
+    def prompt(cls) -> "InputParam":


I'm a bit apprehensive about introducing dedicated class methods for common parameters in this way. I think the class can become quite large as common inputs expand.

I would prefer to keep current syntax (IMO this ensures InputParams are defined in a consistent way) and use post init on the dataclass to automatically add a description. e.g

# centralised descriptions would live somewhere like constants.py # can be used for modular + non-modular INPUT_PARAM_TEMPLATES = { "prompt": {"type_hint": str, "required": True, "description": "The prompt or prompts to guide image generation."}, "height": {"type_hint": int, "description": "The height in pixels of the generated image."}, "width": {"type_hint": int, "description": "The width in pixels of the generated image."}, "generator": {"type_hint": torch.Generator, "description": "Torch generator for deterministic generation."}, # ... } @dataclass class InputParam: name: str = None type_hint: Any = None required: bool = False default: Any = None description: str = None def __post_init__(self): if not self.name or self.name not in INPUT_PARAM_TEMPLATES: return template = INPUT_PARAM_TEMPLATES[self.name] if self.type_hint is None: self.type_hint = template.get("type_hint") if self.description is None: self.description = template.get("description")

If we feel that methods for these inputs are necessary, one way to address it without adding individual methods to the InputParam is to use a metaclass. It would result in the InputParam object being less crowded.

class InputParamMeta(type): def __getattr__(cls, name: str): if name in INPUT_PARAM_TEMPLATES: def factory(**overrides): return cls(name=name, **overrides) return factory raise AttributeError(f"No template named '{name}'") @dataclass class InputParam(metaclass=InputParamMeta):

sayakpaul

I guess it will be better to review the PR after https://github.com/huggingface/diffusers/pull/12958/files#r2689367739 ? 👀

sayakpaul

Super cool feature!

sayakpaul · 2026-01-16T10:34:35Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

 # ====================


+# auto_docstring


We could add this to make style and error out in the CI if it's not the case? 👀

This will help us enforce consistency.

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

sayakpaul · 2026-01-16T10:37:12Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

+        return "Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block."
+        " - `QwenImageTextEncoderStep` (text_encoder) is used when `prompt` is provided."
+        " - if `prompt` is not provided, step will be skipped."


Could this be autogenerated as well provided a class name?

good idea!
I think it should be doable, I will look into it (not in this PR)

sayakpaul · 2026-01-16T10:38:25Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

+
+      Inputs:
+
+          mask_image (`Image`):


No strong opinion but I find PIL.Image.Image a tad bit more informative as a type hint.

sayakpaul · 2026-01-16T10:38:41Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

+
+          processed_image (`None`):
+
+          image_latents (`Tensor`):


Similarly torch.Tensor?

I will look into improving them in a follow-up

sayakpaul · 2026-01-16T10:38:51Z

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+
+          generator (`Generator`, *optional*):


torch.Generator perhaps?

src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py

src/diffusers/modular_pipelines/modular_pipeline_utils.py

src/diffusers/modular_pipelines/qwenimage/before_denoise.py

… into modular-doc-improv

yiyixuxu · 2026-01-19T09:22:35Z

@bot /style

github-actions · 2026-01-19T09:23:05Z

Style bot fixed some files and pushed the changes.

up

7b499de

yiyixuxu added 10 commits January 10, 2026 10:52

up up

b29873d

Merge branch 'main' into modular-doc-improv

fbfe5c8

update outputs

43ab148

style

34a743e

add modular_auto_docstring!

ff09bf1

more auto docstring

d20f413

style

2a81f2e

up up up

f0555af

more more

507953f

up

1c90ce3

yiyixuxu changed the title ~~[modular] more doc related refactors~~ [modular] add auto_docstring & more doc related refactors Jan 10, 2026

yiyixuxu commented Jan 10, 2026

View reviewed changes

yiyixuxu requested review from DN6 and sayakpaul January 10, 2026 11:34

DN6 reviewed Jan 14, 2026

View reviewed changes

sayakpaul reviewed Jan 16, 2026

View reviewed changes

yiyixuxu added 9 commits January 17, 2026 09:36

address feedbacks

aea0d04

add TODO in the description for empty docstring

25c968a

refactor based on dhruv's feedback: remove the class method

de03d7f

add template method

002c3e8

up

1f2dbc9

up up up

fb15752

apply auto docstring

8d45ff5

make style

f056af1

rmove space in make docstring

9452520

yiyixuxu commented Jan 19, 2026

View reviewed changes

src/diffusers/modular_pipelines/modular_pipeline_utils.py Outdated Show resolved Hide resolved

yiyixuxu commented Jan 19, 2026

View reviewed changes

src/diffusers/modular_pipelines/qwenimage/before_denoise.py Outdated Show resolved Hide resolved

yiyixuxu added 5 commits January 18, 2026 22:44

Apply suggestions from code review

7e9d2b9

revert change in z

b7127ce

Merge branch 'modular-doc-improv' of github.com:huggingface/diffusers…

d75fbc4

… into modular-doc-improv

fix

1f9576a

Merge branch 'main' into modular-doc-improv

aba551c

Apply style fixes

23d0642

[modular] add auto_docstring & more doc related refactors #12958

Are you sure you want to change the base?

[modular] add auto_docstring & more doc related refactors #12958

Uh oh!

Conversation

yiyixuxu commented Jan 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Usage

Uh oh!

HuggingFaceDocBuilderDev commented Jan 10, 2026

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiyixuxu commented Jan 19, 2026

Uh oh!

github-actions bot commented Jan 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yiyixuxu commented Jan 10, 2026 •

edited

Loading

github-actions bot commented Jan 19, 2026 •

edited

Loading