dynamic aspect support

pipeline.py

@@ -96,11 +96,12 @@ class SEGCFGSelfAttnProcessor:
     Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
     """
 
-    def __init__(self, blur_sigma=1.0, do_cfg=True, inf_blur_threshold=9999.0):
+    def __init__(self, latent_ratio=1.0, blur_sigma=1.0, do_cfg=True, inf_blur_threshold=9999.0):
         if not hasattr(F, "scaled_dot_product_attention"):
             raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
         self.blur_sigma = blur_sigma
         self.do_cfg = do_cfg
+        self.latent_ratio = latent_ratio
         if self.blur_sigma > inf_blur_threshold:
             self.inf_blur = True
         else:
@@ -157,13 +158,16 @@ class SEGCFGSelfAttnProcessor:
         head_dim = inner_dim // attn.heads
 
         query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
         key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
         value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
 
         # the output of sdp = (batch, num_heads, seq_len, head_dim)
         # TODO: add support for attn.scale when we move to Torch 2.1
-        height = width = math.isqrt(query.shape[2])
+        # height = width = math.isqrt(query.shape[2])
+        height = round(self.latent_ratio * math.sqrt(query.shape[2] / self.latent_ratio))
+        width = round(math.sqrt(query.shape[2] / self.latent_ratio))
+        # A = query.shape[2]
+        # height, width = int(self.latent_ratio * torch.sqrt(A/self.latent_ratio)), int(torch.sqrt(A/self.latent_ratio))
         if self.do_cfg:
             query_uncond, query_org, query_ptb = query.chunk(3)
             query_ptb = query_ptb.permute(0, 1, 3, 2).view(batch_size//3, attn.heads * head_dim, height, width)
@@ -1409,7 +1413,9 @@ class StableDiffusionXLPipeline(
                 # change attention layer in UNet if use SEG
                 if self.do_seg:
                     
-                    replace_processor = SEGCFGSelfAttnProcessor(blur_sigma=seg_blur_sigma, do_cfg=self.do_classifier_free_guidance)
+                    # h/w
+                    latent_ratio = height / width
+                    replace_processor = SEGCFGSelfAttnProcessor(latent_ratio=latent_ratio, blur_sigma=seg_blur_sigma, do_cfg=self.do_classifier_free_guidance)
 
                     if self.seg_applied_layers_index:
                         drop_layers = self.seg_applied_layers_index