working masking...

hydra_node/models.py

@@ -268,17 +268,21 @@ class StableDiffusionModel(nn.Module):
         if request.image is None:
             main_noise = []
             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:
+                    noise_x = sample_start_noise_special(request.seed, request, self.device)
+                else:
+                    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 = torch.from_numpy(mask).clone().to(self.device).permute(2, 0, 1)
                         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].unsqueeze(0)
+                        mask = torch.repeat_interleave(mask, request.latent_channels, dim=0)
                         mask = (mask > 0.5).float()
 
                         # interpolate start noise