more refactor

Author: yiyixuxu

src/diffusers/pipelines/components_manager.py

@@ -256,13 +256,99 @@ def remove(self, name):
         if self._auto_offload_enabled:
             self.enable_auto_cpu_offload(self._auto_offload_device)
 
+    # YiYi TODO: looking into improving the search pattern
     def get(self, names: Union[str, List[str]]):
+        """
+        Get components by name with simple pattern matching.
+        
+        Args:
+            names: Component name(s) or pattern(s)
+                Patterns:
+                - "unet" : exact match
+                - "!unet" : everything except exact match "unet"
+                - "base_*" : everything starting with "base_"
+                - "!base_*" : everything NOT starting with "base_"
+                - "*unet*" : anything containing "unet"
+                - "!*unet*" : anything NOT containing "unet"
+                - "refiner|vae|unet" : anything containing any of these terms
+                - "!refiner|vae|unet" : anything NOT containing any of these terms
+        
+        Returns:
+            Single component if names is str and matches one component,
+            dict of components if names matches multiple components or is a list
+        """
         if isinstance(names, str):
-            if names not in self.components:
+            # Check if this is a "not" pattern
+            is_not_pattern = names.startswith('!')
+            if is_not_pattern:
+                names = names[1:]  # Remove the ! prefix
+            
+            # Handle OR patterns (containing |)
+            if '|' in names:
+                terms = names.split('|')
+                matches = {
+                    name: comp for name, comp in self.components.items() 
+                    if any((term in name) != is_not_pattern for term in terms)  # Flip condition if not pattern
+                }
+                if is_not_pattern:
+                    logger.info(f"Getting components NOT containing any of {terms}: {list(matches.keys())}")
+                else:
+                    logger.info(f"Getting components containing any of {terms}: {list(matches.keys())}")
+            
+            # Exact match
+            elif names in self.components:
+                if is_not_pattern:
+                    matches = {
+                        name: comp for name, comp in self.components.items() 
+                        if name != names
+                    }
+                    logger.info(f"Getting all components except '{names}': {list(matches.keys())}")
+                else:
+                    logger.info(f"Getting component: {names}")
+                    return self.components[names]
+            
+            # Prefix match (ends with *)
+            elif names.endswith('*'):
+                prefix = names[:-1]
+                matches = {
+                    name: comp for name, comp in self.components.items() 
+                    if name.startswith(prefix) != is_not_pattern  # Flip condition if not pattern
+                }
+                if is_not_pattern:
+                    logger.info(f"Getting components NOT starting with '{prefix}': {list(matches.keys())}")
+                else:
+                    logger.info(f"Getting components starting with '{prefix}': {list(matches.keys())}")
+            
+            # Contains match (starts with *)
+            elif names.startswith('*'):
+                search = names[1:-1] if names.endswith('*') else names[1:]
+                matches = {
+                    name: comp for name, comp in self.components.items() 
+                    if (search in name) != is_not_pattern  # Flip condition if not pattern
+                }
+                if is_not_pattern:
+                    logger.info(f"Getting components NOT containing '{search}': {list(matches.keys())}")
+                else:
+                    logger.info(f"Getting components containing '{search}': {list(matches.keys())}")
+            
+            else:
                 raise ValueError(f"Component '{names}' not found in ComponentsManager")
-            return self.components[names]
+            
+            if not matches:
+                raise ValueError(f"No components found matching pattern '{names}'")
+            return matches if len(matches) > 1 else next(iter(matches.values()))
+        
         elif isinstance(names, list):
-            return {n: self.components[n] for n in names}
+            results = {}
+            for name in names:
+                result = self.get(name)
+                if isinstance(result, dict):
+                    results.update(result)
+                else:
+                    results[name] = result
+            logger.info(f"Getting multiple components: {list(results.keys())}")
+            return results
+        
         else:
             raise ValueError(f"Invalid type for names: {type(names)}")
 
@@ -431,18 +517,34 @@ def __repr__(self):
 
         return output
 
-    def add_from_pretrained(self, pretrained_model_name_or_path, **kwargs):
+    def add_from_pretrained(self, pretrained_model_name_or_path, prefix: Optional[str] = None, **kwargs):
+        """
+        Load components from a pretrained model and add them to the manager.
+        
+        Args:
+            pretrained_model_name_or_path (str): The path or identifier of the pretrained model
+            prefix (str, optional): Prefix to add to all component names loaded from this model.
+                                  If provided, components will be named as "{prefix}_{component_name}"
+            **kwargs: Additional arguments to pass to DiffusionPipeline.from_pretrained()
+        """
         from ..pipelines.pipeline_utils import DiffusionPipeline
 
         pipe = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path, **kwargs)
         for name, component in pipe.components.items():
-            if name not in self.components and component is not None:
-                self.add(name, component)
-            elif name in self.components:
+                
+            if component is None:
+                continue
+            
+            # Add prefix if specified
+            component_name = f"{prefix}_{name}" if prefix else name
+            
+            if component_name not in self.components:
+                self.add(component_name, component)
+            else:
                 logger.warning(
-                    f"Component '{name}' already exists in ComponentsManager and will not be added. To add it, either:\n"
-                    f"1. remove the existing component with remove('{name}')\n"
-                    f"2. Use a different name: add('{name}_2', component)"
+                    f"Component '{component_name}' already exists in ComponentsManager and will not be added. To add it, either:\n"
+                    f"1. remove the existing component with remove('{component_name}')\n"
+                    f"2. Use a different prefix: add_from_pretrained(..., prefix='{prefix}_2')"
                 )
 
 def summarize_dict_by_value_and_parts(d: Dict[str, Any]) -> Dict[str, Any]:

src/diffusers/pipelines/modular_pipeline.py

@@ -975,10 +975,16 @@ def intermediates_inputs(self) -> List[str]:
         for block in self.blocks.values():
             # Add inputs that aren't in outputs yet
             inputs.extend(input_name for input_name in block.intermediates_inputs if input_name.name not in outputs)
-            # Add this block's outputs
-            block_intermediates_outputs = [out.name for out in block.intermediates_outputs]
-            outputs.update(block_intermediates_outputs)
 
+            # Only add outputs if the block cannot be skipped
+            should_add_outputs = True
+            if hasattr(block, "block_trigger_inputs") and None not in block.block_trigger_inputs:
+                should_add_outputs = False
+            
+            if should_add_outputs:
+                # Add this block's outputs
+                block_intermediates_outputs = [out.name for out in block.intermediates_outputs]
+                outputs.update(block_intermediates_outputs)
         return inputs
 
     @property
@@ -1035,47 +1041,59 @@ def trigger_inputs(self):
         return self._get_trigger_inputs()
 
     def _traverse_trigger_blocks(self, trigger_inputs):
+        # Convert trigger_inputs to a set for easier manipulation
+        active_triggers = set(trigger_inputs)
 
-        def fn_recursive_traverse(block, block_name, trigger_inputs):
+        def fn_recursive_traverse(block, block_name, active_triggers):
             result_blocks = OrderedDict()
+            
             # sequential or PipelineBlock
             if not hasattr(block, 'block_trigger_inputs'):
                 if hasattr(block, 'blocks'):
                     # sequential
                     for block_name, block in block.blocks.items():
-                        blocks_to_update = fn_recursive_traverse(block, block_name, trigger_inputs)
+                        blocks_to_update = fn_recursive_traverse(block, block_name, active_triggers)
                         result_blocks.update(blocks_to_update)
                 else:
                     # PipelineBlock
                     result_blocks[block_name] = block
+                    # Add this block's output names to active triggers if defined
+                    if hasattr(block, 'outputs'):
+                        active_triggers.update(out.name for out in block.outputs)
                 return result_blocks
 
             # auto
             else:
-                # Find first block_trigger_input that matches any value in our trigger_value tuple
+                # Find first block_trigger_input that matches any value in our active_triggers
                 this_block = None
+                matching_trigger = None
                 for trigger_input in block.block_trigger_inputs:
-                    if trigger_input is not None and trigger_input in trigger_inputs:
+                    if trigger_input is not None and trigger_input in active_triggers:
                         this_block = block.trigger_to_block_map[trigger_input]
+                        matching_trigger = trigger_input
                         break
 
                 # If no matches found, try to get the default (None) block
                 if this_block is None and None in block.block_trigger_inputs:
                     this_block = block.trigger_to_block_map[None]
+                    matching_trigger = None
 
                 if this_block is not None:
                     # sequential/auto
                     if hasattr(this_block, 'blocks'):
-                        result_blocks.update(fn_recursive_traverse(this_block, block_name, trigger_inputs))
+                        result_blocks.update(fn_recursive_traverse(this_block, block_name, active_triggers))
                     else:
                         # PipelineBlock
                         result_blocks[block_name] = this_block
+                        # Add this block's output names to active triggers if defined
+                        if hasattr(this_block, 'outputs'):
+                            active_triggers.update(out.name for out in this_block.outputs)
 
             return result_blocks
 
         all_blocks = OrderedDict()
         for block_name, block in self.blocks.items():
-            blocks_to_update = fn_recursive_traverse(block, block_name, trigger_inputs)
+            blocks_to_update = fn_recursive_traverse(block, block_name, active_triggers)
             all_blocks.update(blocks_to_update)
         return all_blocks
 

src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_modular.py

@@ -2994,14 +2994,28 @@ def description(self):
                " - `StableDiffusionXLInpaintDecodeStep` (inpaint) is used when `padding_mask_crop` is provided.\n" + \
                " - `StableDiffusionXLDecodeStep` (non-inpaint) is used when `padding_mask_crop` is not provided."
 
+class StableDiffusionXLAutoIPAdapterStep(AutoPipelineBlocks):
+    block_classes = [StableDiffusionXLIPAdapterStep]
+    block_names = ["ip_adapter"]
+    block_trigger_inputs = ["ip_adapter_image"]
 
-class StableDiffusionAutoPipeline(SequentialPipelineBlocks):
-    block_classes = [StableDiffusionXLTextEncoderStep, StableDiffusionXLAutoVaeEncoderStep, StableDiffusionXLAutoBeforeDenoiseStep, StableDiffusionXLAutoDenoiseStep, StableDiffusionXLAutoDecodeStep]
-    block_names = ["text_encoder", "image_encoder", "before_denoise", "denoise", "decode"]
+    @property
+    def description(self):
+        return "Run IP Adapter step if `ip_adapter_image` is provided."
+
+class StableDiffusionXLAutoPipeline(SequentialPipelineBlocks):
+    block_classes = [StableDiffusionXLTextEncoderStep, StableDiffusionXLAutoIPAdapterStep, StableDiffusionXLAutoVaeEncoderStep, StableDiffusionXLAutoBeforeDenoiseStep, StableDiffusionXLAutoDenoiseStep, StableDiffusionXLAutoDecodeStep]
+    block_names = ["text_encoder", "ip_adapter", "image_encoder", "before_denoise", "denoise", "decode"]
 
     @property
     def description(self):
-        return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL.\n"
+        return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL.\n" + \
+               "- for image-to-image generation, you need to provide either `image` or `image_latents`\n" + \
+               "- for inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop` \n" + \
+               "- to run the controlnet workflow, you need to provide `control_image`\n" + \
+               "- to run the controlnet_union workflow, you need to provide `control_image` and `control_mode`\n" + \
+               "- to run the ip_adapter workflow, you need to provide `ip_adapter_image`\n" + \
+               "- for text-to-image generation, all you need to provide is `prompt`"
 
 # block mapping 
 TEXT2IMAGE_BLOCKS = OrderedDict([