fix

Author: yiyixuxu

src/diffusers/pipelines/modular_pipeline.py

@@ -82,9 +82,9 @@ def to_dict(self) -> Dict[str, Any]:
     def __repr__(self):
         def format_value(v):
             if hasattr(v, "shape") and hasattr(v, "dtype"):
-                return f"Tensor(\n      dtype={v.dtype}, shape={v.shape}\n      {v})"
+                return f"Tensor(dtype={v.dtype}, shape={v.shape})"
             elif isinstance(v, list) and len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
-                return f"[Tensor(\n      dtype={v[0].dtype}, shape={v[0].shape}\n      {v[0]}), ...]"
+                return f"[Tensor(dtype={v[0].dtype}, shape={v[0].shape}), ...]"
             else:
                 return repr(v)
 
@@ -238,6 +238,10 @@ def __init__(self):
         if not (len(self.block_classes) == len(self.block_names) == len(self.block_trigger_inputs)):
             raise ValueError(f"In {self.__class__.__name__}, the number of block_classes, block_names, and block_trigger_inputs must be the same.")
         default_blocks = [t for t in self.block_trigger_inputs if t is None]
+        # can only have 1 or 0 default block, and has to put in the last 
+        # the order of blocksmatters here because the first block with matching trigger will be dispatched
+        # e.g. blocks = [inpaint, img2img] and block_trigger_inputs = ["mask", "image"]
+        # if both mask and image are provided, it is inpaint; if only image is provided, it is img2img
         if len(default_blocks) > 1 or (
                 len(default_blocks) == 1 and self.block_trigger_inputs[-1] is not None
             ):
@@ -248,6 +252,7 @@ def __init__(self):
 
         # Map trigger inputs to block objects
         self.trigger_to_block_map = dict(zip(self.block_trigger_inputs, self.blocks.values()))
+        self.trigger_to_block_name_map = dict(zip(self.block_trigger_inputs, self.blocks.keys()))
         self.block_to_trigger_map = dict(zip(self.blocks.keys(), self.block_trigger_inputs))
 
     @property
@@ -324,6 +329,9 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState:
             if input_name is not None and state.get_input(input_name) is not None:
                 block = self.trigger_to_block_map[input_name]
                 break
+            elif input_name is not None and state.get_intermediate(input_name) is not None:
+                block = self.trigger_to_block_map[input_name]
+                break
 
         if block is None:
             logger.warning(f"skipping auto block: {self.__class__.__name__}")

src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_modular.py

@@ -1112,6 +1112,7 @@ def intermediates_inputs(self) -> List[str]:
             "mask", # inpainting
             "masked_image_latents", # inpainting
             "noise", # inpainting
+            "image_latents", # inpainting
         ]
 
     @property
@@ -2028,25 +2029,24 @@ class StableDiffusionXLAutoVaeEncoderStep(AutoPipelineBlocks):
     block_trigger_inputs = ["mask_image", "image"]
 
 
-class StableDiffusionXLAutoSetTimestepsStep(AutoPipelineBlocks):
-    block_classes = [StableDiffusionXLImg2ImgSetTimestepsStep, StableDiffusionXLSetTimestepsStep]
-    block_names = ["img2img", "text2img"]
-    block_trigger_inputs = ["image", None]
+class StableDiffusionXLBeforeDenoiseStep(SequentialPipelineBlocks):
+    block_classes = [StableDiffusionXLInputStep, StableDiffusionXLSetTimestepsStep, StableDiffusionXLPrepareLatentsStep, StableDiffusionXLPrepareAdditionalConditioningStep]
+    block_names = ["input", "set_timesteps", "prepare_latents", "prepare_add_cond"]
 
+class StableDiffusionXLImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks):
+    block_classes = [StableDiffusionXLInputStep, StableDiffusionXLImg2ImgSetTimestepsStep, StableDiffusionXLImg2ImgPrepareLatentsStep, StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep]
+    block_names = ["input", "set_timesteps", "prepare_latents", "prepare_add_cond"]
 
-class StableDiffusionXLAutoPrepareLatentsStep(AutoPipelineBlocks):
-    block_classes = [StableDiffusionXLInpaintPrepareLatentsStep, StableDiffusionXLImg2ImgPrepareLatentsStep, StableDiffusionXLPrepareLatentsStep]
-    block_names = ["inpaint","img2img", "text2img"]
-    block_trigger_inputs = ["mask_image", "image", None]
+class StableDiffusionXLInpaintBeforeDenoiseStep(SequentialPipelineBlocks):
+    block_classes = [StableDiffusionXLInputStep, StableDiffusionXLImg2ImgSetTimestepsStep, StableDiffusionXLInpaintPrepareLatentsStep, StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep]
+    block_names = ["input", "set_timesteps", "prepare_latents", "prepare_add_cond"]
 
 
-class StableDiffusionXLAutoPrepareAdditionalConditioningStep(AutoPipelineBlocks):
-    block_classes = [
-        StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep,
-        StableDiffusionXLPrepareAdditionalConditioningStep,
-    ]
-    block_names = ["img2img", "text2img"]
-    block_trigger_inputs = ["image", None]
+class StableDiffusionXLAutoBeforeDenoiseStep(AutoPipelineBlocks):
+    block_classes = [StableDiffusionXLInpaintBeforeDenoiseStep, StableDiffusionXLImg2ImgBeforeDenoiseStep, StableDiffusionXLBeforeDenoiseStep]
+    block_names = ["inpaint", "img2img", "text2img"]
+    block_trigger_inputs = ["mask", "image_latents", None]
+
 
 
 class StableDiffusionXLAutoDenoiseStep(AutoPipelineBlocks):
@@ -2064,10 +2064,10 @@ class StableDiffusionXLAutoDecodeStep(AutoPipelineBlocks):
 TEXT2IMAGE_BLOCKS = OrderedDict([
     ("text_encoder", StableDiffusionXLTextEncoderStep),
     ("input", StableDiffusionXLInputStep),
-    ("set_timesteps", StableDiffusionXLAutoSetTimestepsStep),
-    ("prepare_latents", StableDiffusionXLAutoPrepareLatentsStep),
-    ("prepare_add_cond", StableDiffusionXLAutoPrepareAdditionalConditioningStep),
-    ("denoise", StableDiffusionXLAutoDenoiseStep),
+    ("set_timesteps", StableDiffusionXLSetTimestepsStep),
+    ("prepare_latents", StableDiffusionXLPrepareLatentsStep),
+    ("prepare_add_cond", StableDiffusionXLPrepareAdditionalConditioningStep),
+    ("denoise", StableDiffusionXLDenoiseStep),
     ("decode", StableDiffusionXLDecodeStep)
 ])
 
@@ -2099,11 +2099,8 @@ class StableDiffusionXLAutoDecodeStep(AutoPipelineBlocks):
 
 AUTO_BLOCKS = OrderedDict([
     ("text_encoder", StableDiffusionXLTextEncoderStep),
-    ("input", StableDiffusionXLInputStep),
     ("image_encoder", StableDiffusionXLAutoVaeEncoderStep),
-    ("set_timesteps", StableDiffusionXLAutoSetTimestepsStep),
-    ("prepare_latents", StableDiffusionXLAutoPrepareLatentsStep),
-    ("prepare_add_cond", StableDiffusionXLAutoPrepareAdditionalConditioningStep),
+    ("before_denoise", StableDiffusionXLAutoBeforeDenoiseStep),
     ("denoise", StableDiffusionXLAutoDenoiseStep),
     ("decode", StableDiffusionXLAutoDecodeStep)
 ])
@@ -2138,11 +2135,18 @@ def vae_scale_factor(self):
             vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         return vae_scale_factor
 
+    @property
+    def num_channels_unet(self):
+        num_channels_unet = 4
+        if hasattr(self, "unet") and self.unet is not None:
+            num_channels_unet = self.unet.config.in_channels
+        return num_channels_unet
+
     @property
     def num_channels_latents(self):
         num_channels_latents = 4
-        if hasattr(self, "unet") and self.unet is not None:
-            num_channels_latents = self.unet.config.in_channels
+        if hasattr(self, "vae") and self.vae is not None:
+            num_channels_latents = self.vae.config.latent_channels
         return num_channels_latents
 
     # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids