Update pipeline.py

pipeline.py

@@ -219,10 +219,15 @@ class AAS_XL(AttentionBase):
         self.mask = mask  # mask with shape (1, 1 ,h, w)
         self.ss_steps = ss_steps
         self.ss_scale = ss_scale
-        self.mask_16 = F.max_pool2d(mask, (1024 // 16, 1024 // 16)).round().squeeze().squeeze()
-        self.mask_32 = F.max_pool2d(mask, (1024 // 32, 1024 // 32)).round().squeeze().squeeze()
-        self.mask_64 = F.max_pool2d(mask, (1024 // 64, 1024 // 64)).round().squeeze().squeeze()
-        self.mask_128 = F.max_pool2d(mask, (1024 // 128, 1024 // 128)).round().squeeze().squeeze()
+        # self.mask_16 = F.max_pool2d(mask, (1024 // 16, 1024 // 16)).round().squeeze().squeeze()
+        # self.mask_32 = F.max_pool2d(mask, (1024 // 32, 1024 // 32)).round().squeeze().squeeze()
+        # self.mask_64 = F.max_pool2d(mask, (1024 // 64, 1024 // 64)).round().squeeze().squeeze()
+        # self.mask_128 = F.max_pool2d(mask, (1024 // 128, 1024 // 128)).round().squeeze().squeeze()
+        # target_size: 입력된 mask의 크기 (height, width)
+        self.mask_16 = F.max_pool2d(mask, (mask.shape[-2] // 16, mask.shape[-1] // 16)).round().squeeze().squeeze()
+        self.mask_32 = F.max_pool2d(mask, (mask.shape[-2] // 32, mask.shape[-1] // 32)).round().squeeze().squeeze()
+        self.mask_64 = F.max_pool2d(mask, (mask.shape[-2] // 64, mask.shape[-1] // 64)).round().squeeze().squeeze()
+        self.mask_128 = F.max_pool2d(mask, (mask.shape[-2] // 128, mask.shape[-1] // 128)).round().squeeze().squeeze()
 
     def attn_batch(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, is_mask_attn, mask, **kwargs):
         B = q.shape[0] // num_heads