Upload pipeline.py

pipeline.py

@@ -20,30 +20,50 @@ 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
 
@@ -269,6 +289,11 @@ class StableDiffusionControlNetPipeline(
             do_convert_rgb=True,
             do_normalize=False,
         )
+        self.control_mask_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor,
+            do_normalize=False,
+            do_convert_grayscale=True,
+        )
         self.register_to_config(requires_safety_checker=requires_safety_checker)
 
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
@@ -1349,6 +1374,17 @@ class StableDiffusionControlNetPipeline(
         else:
             assert False
 
+        if effective_region_mask is not None:
+            effective_region_mask = self.control_mask_processor.preprocess(
+                effective_region_mask, height=height, width=width
+            ).to(dtype=torch.float32)
+
+            print("mask shape:")
+            print(effective_region_mask.shape)
+            print()
+
+            print(effective_region_mask)
+
         # 5. Prepare timesteps
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler, num_inference_steps, device, timesteps, sigmas