Added support for RunwayML inpainting model

1 files changed, 32 insertions, 2 deletions

diff --git a/modules/processing.py b/modules/processing.py
index bcb0c32c..a6c308f9 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -546,7 +546,16 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         if not self.enable_hr:

             x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

-            samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning)

+            

+            # The "masked-image" in this case will just be all zeros since the entire image is masked.

+            image_conditioning = torch.zeros(x.shape[0], 3, self.height, self.width, device=x.device)

+            image_conditioning = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image_conditioning)) 

+

+            # Add the fake full 1s mask to the first dimension.

+            image_conditioning = torch.nn.functional.pad(image_conditioning, (0, 0, 0, 0, 1, 0), value=1.0)

+            image_conditioning = image_conditioning.to(x.dtype)

+

+            samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=image_conditioning)

             return samples

 

         x = create_random_tensors([opt_C, self.firstphase_height // opt_f, self.firstphase_width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

@@ -714,10 +723,31 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             elif self.inpainting_fill == 3:

                 self.init_latent = self.init_latent * self.mask

 

+        if self.image_mask is not None:

+            conditioning_mask = np.array(self.image_mask.convert("L"))

+            conditioning_mask = conditioning_mask.astype(np.float32) / 255.0

+            conditioning_mask = torch.from_numpy(conditioning_mask[None, None])

+

+            # Inpainting model uses a discretized mask as input, so we round to either 1.0 or 0.0

+            conditioning_mask = torch.round(conditioning_mask)

+        else:

+            conditioning_mask = torch.ones(1, 1, *image.shape[-2:])

+

+        # Create another latent image, this time with a masked version of the original input.

+        conditioning_mask = conditioning_mask.to(image.device)

+        conditioning_image = image * (1.0 - conditioning_mask)

+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(conditioning_image))

+

+        # Create the concatenated conditioning tensor to be fed to `c_concat`

+        conditioning_mask = torch.nn.functional.interpolate(conditioning_mask, size=self.init_latent.shape[-2:])

+        conditioning_mask = conditioning_mask.expand(conditioning_image.shape[0], -1, -1, -1)

+        self.image_conditioning = torch.cat([conditioning_mask, conditioning_image], dim=1)

+        self.image_conditioning = self.image_conditioning.to(shared.device).type(self.sd_model.dtype)

+

     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength):

         x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

 

-        samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning)

+        samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning, image_conditioning=self.image_conditioning)

 

         if self.mask is not None:

             samples = samples * self.nmask + self.init_latent * self.mask