infilling maybe works?

hydra_node/models.py

@@ -68,6 +68,14 @@ def sample_start_noise(seed, C, H, W, f, device="cuda"):
     noise = torch.randn([C, (H) // f, (W) // f], device=device).unsqueeze(0)
     return noise
 
+def sample_start_noise_special(seed, request, device="cuda"):
+    if request.seed is not None:
+        torch.manual_seed(request.seed)
+        np.random.seed(request.seed)
+    
+    noise = torch.randn([request.latent_channels, request.height // request.downsampling_factor, request.width // request.downsampling_factor], device=device).unsqueeze(0)
+    return noise
+
 @torch.no_grad()
 #@torch.autocast("cuda", enabled=True, dtype=torch.float16)
 def encode_image(image, model):
@@ -215,7 +223,7 @@ class StableDiffusionModel(nn.Module):
         if request.module is not None:
             if request.module == "vanilla":
                 pass
-            
+
             else:
                 module = self.premodules[request.module]
                 CrossAttention.set_hypernetwork(module)
@@ -259,9 +267,25 @@ class StableDiffusionModel(nn.Module):
 
         if request.image is None:
             main_noise = []
-            for seed in range(request.seed, request.seed+request.n_samples):
-                main_noise.append(sample_start_noise(seed, request.latent_channels, request.height, request.width, request.downsampling_factor, self.device))
-            
+            for seed_offset in range(request.n_samples):
+                noise_x = sample_start_noise_special(request.seed+seed_offset, request, self.device)
+                if request.masks is not None:
+                    for maskobj in request.masks:
+                        mask_seed = maskobj["seed"]
+                        mask = maskobj["mask"]
+                        mask = np.asarray(mask)
+                        mask = torch.from_numpy(mask).clone().to(self.device)
+                        mask = mask.float() / 255.0
+                        # convert RGB or grayscale image into 4-channel
+                        mask = mask[0]
+                        mask = torch.repeat_interleave(mask, request.latent_channels, dim=0).unsqueeze(0)
+                        mask = (mask > 0.5).float()
+
+                        # interpolate start noise
+                        noise_x = (noise_x * (1-mask)) + (sample_start_noise_special(mask_seed+seed_offset, request, self.device) * mask)
+
+                main_noise.append(noise_x)
+
             main_noise = torch.cat(main_noise, dim=0)
             start_code = main_noise
         

hydra_node/sanitize.py

@@ -136,7 +136,6 @@ def sanitize_stable_diffusion(request, config):
 
     if request.masks is not None:
         masks = request.masks
-        images = []
         for x in range(len(masks)):
             image = masks[x]["mask"]
             try:
@@ -160,6 +159,7 @@ def sanitize_stable_diffusion(request, config):
                 image = Image.open(BytesIO(image))
                 #image = image.convert('RGB')
                 image = image.resize((request.width//request.downsampling_factor, request.height//request.downsampling_factor), resample=Image.Resampling.LANCZOS)
+
             except Exception as e:
                 traceback.print_exc()
                 return False, "Error while opening and cleaning image"