Upload pipeline.py

pipeline.py

@@ -20,50 +20,30 @@ import numpy as np
 import PIL.Image
 import torch
 import torch.nn.functional as F
-from transformers import (
-    CLIPImageProcessor,
-    CLIPTextModel,
-    CLIPTokenizer,
-    CLIPVisionModelWithProjection,
-)
+from transformers import (CLIPImageProcessor, CLIPTextModel, CLIPTokenizer,
+                          CLIPVisionModelWithProjection)
 
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
-from diffusers.loaders import (
-    FromSingleFileMixin,
-    IPAdapterMixin,
-    StableDiffusionLoraLoaderMixin,
-    TextualInversionLoaderMixin,
-)
-from diffusers.models import (
-    AutoencoderKL,
-    ControlNetModel,
-    ImageProjection,
-    UNet2DConditionModel,
-)
+from diffusers.loaders import (FromSingleFileMixin, IPAdapterMixin,
+                               StableDiffusionLoraLoaderMixin,
+                               TextualInversionLoaderMixin)
+from diffusers.models import (AutoencoderKL, ControlNetModel, ImageProjection,
+                              UNet2DConditionModel)
 from diffusers.models.lora import adjust_lora_scale_text_encoder
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
-from diffusers.pipelines.stable_diffusion.pipeline_output import (
-    StableDiffusionPipelineOutput,
-)
-from diffusers.pipelines.stable_diffusion.safety_checker import (
-    StableDiffusionSafetyChecker,
-)
+from diffusers.pipelines.pipeline_utils import (DiffusionPipeline,
+                                                StableDiffusionMixin)
+from diffusers.pipelines.stable_diffusion.pipeline_output import \
+    StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import \
+    StableDiffusionSafetyChecker
 from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import (
-    USE_PEFT_BACKEND,
-    deprecate,
-    logging,
-    replace_example_docstring,
-    scale_lora_layers,
-    unscale_lora_layers,
-)
-from diffusers.utils.torch_utils import (
-    is_compiled_module,
-    is_torch_version,
-    randn_tensor,
-)
+from diffusers.utils import (USE_PEFT_BACKEND, deprecate, logging,
+                             replace_example_docstring, scale_lora_layers,
+                             unscale_lora_layers)
+from diffusers.utils.torch_utils import (is_compiled_module, is_torch_version,
+                                         randn_tensor)
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -691,6 +671,7 @@ class StableDiffusionControlNetPipeline(
         control_guidance_start=0.0,
         control_guidance_end=1.0,
         callback_on_step_end_tensor_inputs=None,
+        effective_region_mask=None,
     ):
         if callback_steps is not None and (
             not isinstance(callback_steps, int) or callback_steps <= 0
@@ -853,6 +834,9 @@ class StableDiffusionControlNetPipeline(
                     f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
                 )
 
+        if effective_region_mask is not None:
+            self.check_mask(effective_region_mask)
+
     def check_image(self, image, prompt, prompt_embeds):
         image_is_pil = isinstance(image, PIL.Image.Image)
         image_is_tensor = isinstance(image, torch.Tensor)
@@ -894,6 +878,16 @@ class StableDiffusionControlNetPipeline(
                 f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"
             )
 
+    def check_mask(self, mask):
+        image_is_pil = isinstance(mask, PIL.Image.Image)
+        image_is_tensor = isinstance(mask, torch.Tensor)
+        image_is_np = isinstance(mask, np.ndarray)
+
+        if not image_is_pil and not image_is_tensor and not image_is_np:
+            raise TypeError(
+                f"mask must be passed and be one of PIL image, numpy array, or torch tensor, but is {type(mask)}"
+            )
+
     def prepare_image(
         self,
         image,
@@ -995,6 +989,20 @@ class StableDiffusionControlNetPipeline(
         assert emb.shape == (w.shape[0], embedding_dim)
         return emb
 
+    def apply_effective_region_mask(
+        self, effective_region_mask: torch.Tensor, out: torch.Tensor
+    ) -> torch.Tensor:
+        if effective_region_mask is None:
+            return out
+
+        B, C, H, W = out.shape
+        mask = F.interpolate(
+            effective_region_mask.to(out.device),
+            size=(H, W),
+            mode="bilinear",
+        )
+        return out * mask
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -1173,6 +1181,8 @@ class StableDiffusionControlNetPipeline(
         callback = kwargs.pop("callback", None)
         callback_steps = kwargs.pop("callback_steps", None)
 
+        effective_region_mask = kwargs.pop("effective_region_mask", None)
+
         if callback is not None:
             deprecate(
                 "callback",
@@ -1233,6 +1243,7 @@ class StableDiffusionControlNetPipeline(
             control_guidance_start,
             control_guidance_end,
             callback_on_step_end_tensor_inputs,
+            effective_region_mask,
         )
 
         self._guidance_scale = guidance_scale
@@ -1439,6 +1450,7 @@ class StableDiffusionControlNetPipeline(
                         controlnet_cond_scale = controlnet_cond_scale[0]
                     cond_scale = controlnet_cond_scale * controlnet_keep[i]
 
+                # Controlnet is returning the residuals to be added to SD here
                 down_block_res_samples, mid_block_res_sample = self.controlnet(
                     control_model_input,
                     t,
@@ -1449,6 +1461,21 @@ class StableDiffusionControlNetPipeline(
                     return_dict=False,
                 )
 
+                # Apply mask here
+                # Note that downblocks are ordered from largest->smallest
+                if effective_region_mask is not None:
+                    masked_down_block_res_samples = ()
+                    for down_block_res_sample in down_block_res_samples:
+                        down_block_res_sample = self.apply_effective_region_mask(
+                            effective_region_mask, down_block_res_sample
+                        )
+                        masked_down_block_res_samples = (
+                            masked_down_block_res_samples + (down_block_res_sample,)
+                        )
+                    mid_block_res_sample = self.apply_effective_region_mask(
+                        effective_region_mask, mid_block_res_sample
+                    )
+
                 if guess_mode and self.do_classifier_free_guidance:
                     # Inferred ControlNet only for the conditional batch.
                     # To apply the output of ControlNet to both the unconditional and conditional batches,