style

Author: yiyixuxu

src/diffusers/modular_pipelines/components_manager.py

@@ -142,6 +142,7 @@ def custom_offload_with_hook(
     user_hook.attach()
     return user_hook
 
+
 # this is the class that user can customize to implement their own offload strategy
 class AutoOffloadStrategy:
     """
@@ -277,37 +278,36 @@ def find_common_prefix(keys: List[str]) -> str:
 class ComponentsManager:
     """
     A central registry and management system for model components across multiple pipelines.
-    
-    [`ComponentsManager`] provides a unified way to register, track, and reuse model components
-    (like UNet, VAE, text encoders, etc.) across different modular pipelines. It includes
-    features for duplicate detection, memory management, and component organization.
-    
+
+    [`ComponentsManager`] provides a unified way to register, track, and reuse model components (like UNet, VAE, text
+    encoders, etc.) across different modular pipelines. It includes features for duplicate detection, memory
+    management, and component organization.
+
     <Tip warning={true}>
 
         This is an experimental feature and is likely to change in the future.
 
     </Tip>
-        
+
     Example:
         ```python
         from diffusers import ComponentsManager
-        
+
         # Create a components manager
         cm = ComponentsManager()
-        
+
         # Add components
         cm.add("unet", unet_model, collection="sdxl")
         cm.add("vae", vae_model, collection="sdxl")
-        
+
         # Enable auto offloading
         cm.enable_auto_cpu_offload(device="cuda")
-        
+
         # Retrieve components
         unet = cm.get_one(name="unet", collection="sdxl")
         ```
     """
 
-
     _available_info_fields = [
         "model_id",
         "added_time",
@@ -382,7 +382,7 @@ def add(self, name: str, component: Any, collection: Optional[str] = None):
             collection (Optional[str]): The collection to add the component to
 
         Returns:
-            str: The unique component ID, which is generated as "{name}_{id(component)}" where 
+            str: The unique component ID, which is generated as "{name}_{id(component)}" where
                  id(component) is Python's built-in unique identifier for the object
         """
         component_id = f"{name}_{id(component)}"
@@ -669,9 +669,9 @@ def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = "cuda"
 
         Args:
             device (Union[str, int, torch.device]): The execution device where models are moved for forward passes
-            memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of 
-                                        memory to keep free on the device to avoid running out of memory during 
-                                        model execution (e.g., for intermediate activations, gradients, etc.)
+            memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of
+                                        memory to keep free on the device to avoid running out of memory during model
+                                        execution (e.g., for intermediate activations, gradients, etc.)
         """
         if not is_accelerate_available():
             raise ImportError("Make sure to install accelerate to use auto_cpu_offload")
@@ -727,11 +727,13 @@ def get_model_info(
 
         Args:
             component_id (str): Name of the component to get info for
-            fields (Optional[Union[str, List[str]]]): Field(s) to return. Can be a string for single field or list of fields.
-                   If None, uses the available_info_fields setting.
+            fields (Optional[Union[str, List[str]]]):
+                   Field(s) to return. Can be a string for single field or list of fields. If None, uses the
+                   available_info_fields setting.
 
         Returns:
-            Dictionary containing requested component metadata. If fields is specified, returns only those fields. Otherwise, returns all fields.
+            Dictionary containing requested component metadata. If fields is specified, returns only those fields.
+            Otherwise, returns all fields.
         """
         if component_id not in self.components:
             raise ValueError(f"Component '{component_id}' not found in ComponentsManager")
@@ -938,10 +940,10 @@ def get_one(
         load_id: Optional[str] = None,
     ) -> Any:
         """
-        Get a single component by either: 
-        - searching name (pattern matching), collection, or load_id. 
+        Get a single component by either:
+        - searching name (pattern matching), collection, or load_id.
         - passing in a component_id
-        Raises an error if multiple components match or none are found. 
+        Raises an error if multiple components match or none are found.
 
         Args:
             component_id (Optional[str]): Optional component ID to get
@@ -998,13 +1000,13 @@ def get_components_by_ids(self, ids: List[str], return_dict_with_names: Optional
         Get components by a list of IDs.
 
         Args:
-            ids (List[str]): 
+            ids (List[str]):
                 List of component IDs
             return_dict_with_names (Optional[bool]):
                 Whether to return a dictionary with component names as keys:
 
         Returns:
-            Dict[str, Any]: Dictionary of components. 
+            Dict[str, Any]: Dictionary of components.
                 - If return_dict_with_names=True, keys are component names.
                 - If return_dict_with_names=False, keys are component IDs.
 
@@ -1042,4 +1044,3 @@ def get_components_by_names(self, names: List[str], collection: Optional[str] =
         """
         ids = self.get_ids(names, collection)
         return self.get_components_by_ids(ids)
-

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -85,10 +85,9 @@ def set_input(self, key: str, value: Any, kwargs_type: str = None):
         """
         Add an input to the immutable pipeline state, i.e, pipeline_state.inputs.
 
-        The kwargs_type parameter allows you to associate inputs with specific input types.
-        For example, if you call set_input(prompt_embeds=..., kwargs_type="guider_kwargs"),
-        this input will be automatically fetched when a pipeline block has "guider_kwargs" 
-        in its expected_inputs list.
+        The kwargs_type parameter allows you to associate inputs with specific input types. For example, if you call
+        set_input(prompt_embeds=..., kwargs_type="guider_kwargs"), this input will be automatically fetched when a
+        pipeline block has "guider_kwargs" in its expected_inputs list.
 
         Args:
             key (str): The key for the input
@@ -106,10 +105,9 @@ def set_intermediate(self, key: str, value: Any, kwargs_type: str = None):
         """
         Add an intermediate value to the mutable pipeline state, i.e, pipeline_state.intermediates.
 
-        The kwargs_type parameter allows you to associate intermediate values with specific input types.
-        For example, if you call set_intermediate(latents=..., kwargs_type="latents_kwargs"),
-        this intermediate value will be automatically fetched when a pipeline block has "latents_kwargs" 
-        in its expected_intermediate_inputs list.
+        The kwargs_type parameter allows you to associate intermediate values with specific input types. For example,
+        if you call set_intermediate(latents=..., kwargs_type="latents_kwargs"), this intermediate value will be
+        automatically fetched when a pipeline block has "latents_kwargs" in its expected_intermediate_inputs list.
 
         Args:
             key (str): The key for the intermediate value
@@ -414,13 +412,13 @@ def init_pipeline(
             collection=collection,
         )
         return modular_pipeline
-    
+
     @staticmethod
     def combine_inputs(*named_input_lists: List[Tuple[str, List[InputParam]]]) -> List[InputParam]:
         """
-        Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current
-        default value is None and new default value is not None. Warns if multiple non-None default values exist for the
-        same input.
+        Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if
+        current default value is None and new default value is not None. Warns if multiple non-None default values
+        exist for the same input.
 
         Args:
             named_input_lists: List of tuples containing (block_name, input_param_list) pairs
@@ -482,8 +480,6 @@ def combine_outputs(*named_output_lists: List[Tuple[str, List[OutputParam]]]) ->
         return list(combined_dict.values())
 
 
-
-
 class PipelineBlock(ModularPipelineBlocks):
     """
     A Pipeline Block is the basic building block of a Modular Pipeline.
@@ -499,15 +495,33 @@ class PipelineBlock(ModularPipelineBlocks):
 
     Args:
         description (str, optional): A description of the block, defaults to None. Define as a property in subclasses.
-        expected_components (List[ComponentSpec], optional): A list of components that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        expected_configs (List[ConfigSpec], optional): A list of configs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        inputs (List[InputParam], optional): A list of inputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        intermediate_inputs (List[InputParam], optional): A list of intermediate inputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        intermediate_outputs (List[OutputParam], optional): A list of intermediate outputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        outputs (List[OutputParam], optional): A list of outputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        required_inputs (List[str], optional): A list of required inputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        required_intermediate_inputs (List[str], optional): A list of required intermediate inputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
-        required_intermediate_outputs (List[str], optional): A list of required intermediate outputs that are expected to be used in the block, defaults to []. To override, define as a property in subclasses.
+        expected_components (List[ComponentSpec], optional):
+            A list of components that are expected to be used in the block, defaults to []. To override, define as a
+            property in subclasses.
+        expected_configs (List[ConfigSpec], optional):
+            A list of configs that are expected to be used in the block, defaults to []. To override, define as a
+            property in subclasses.
+        inputs (List[InputParam], optional):
+            A list of inputs that are expected to be used in the block, defaults to []. To override, define as a
+            property in subclasses.
+        intermediate_inputs (List[InputParam], optional):
+            A list of intermediate inputs that are expected to be used in the block, defaults to []. To override,
+            define as a property in subclasses.
+        intermediate_outputs (List[OutputParam], optional):
+            A list of intermediate outputs that are expected to be used in the block, defaults to []. To override,
+            define as a property in subclasses.
+        outputs (List[OutputParam], optional):
+            A list of outputs that are expected to be used in the block, defaults to []. To override, define as a
+            property in subclasses.
+        required_inputs (List[str], optional):
+            A list of required inputs that are expected to be used in the block, defaults to []. To override, define as
+            a property in subclasses.
+        required_intermediate_inputs (List[str], optional):
+            A list of required intermediate inputs that are expected to be used in the block, defaults to []. To
+            override, define as a property in subclasses.
+        required_intermediate_outputs (List[str], optional):
+            A list of required intermediate outputs that are expected to be used in the block, defaults to []. To
+            override, define as a property in subclasses.
     """
 
     model_name = None
@@ -997,7 +1011,8 @@ def doc(self):
 
 class SequentialPipelineBlocks(ModularPipelineBlocks):
     """
-    A Pipeline Blocks that combines multiple pipeline block classes into one. When called, it will call each block in sequence.
+    A Pipeline Blocks that combines multiple pipeline block classes into one. When called, it will call each block in
+    sequence.
 
     This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the
     library implements for all the pipeline blocks (such as loading or saving etc.)
@@ -1373,8 +1388,8 @@ def doc(self):
 
 class LoopSequentialPipelineBlocks(ModularPipelineBlocks):
     """
-    A Pipeline blocks that combines multiple pipeline block classes into a For Loop. When called, it will call each block in
-    sequence.
+    A Pipeline blocks that combines multiple pipeline block classes into a For Loop. When called, it will call each
+    block in sequence.
 
     This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the
     library implements for all the pipeline blocks (such as loading or saving etc.)

src/diffusers/modular_pipelines/node_utils.py

@@ -348,8 +348,8 @@ def get_group_name(name, group_params_keys=DEFAULT_PARAMS_GROUPS_KEYS):
 
 class ModularNode(ConfigMixin):
     """
-    A ModularNode is a base class to build UI nodes using diffusers. Currently only supports Mellon.
-    It is a wrapper around a ModularPipelineBlocks object.
+    A ModularNode is a base class to build UI nodes using diffusers. Currently only supports Mellon. It is a wrapper
+    around a ModularPipelineBlocks object.
 
     <Tip warning={true}>