test x formers

app.py

@@ -22,6 +22,8 @@ from diffusers import StableDiffusionUpscalePipeline
 from diffusers import LDMSuperResolutionPipeline
 import cv2
 import onnxruntime
+import xformers
+# from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
 
 def removeFurniture(input_img1,
             input_img2,
@@ -87,6 +89,9 @@ def upscale(image, prompt):
     # pipe = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
     pipe = pipe.to(device)
     pipe.enable_attention_slicing()
+    pipe.enable_xformers_memory_efficient_attention(attention_op=xformers.ops.MemoryEfficientAttentionFlashAttentionOp)
+    # Workaround for not accepting attention shape using VAE for Flash Attention
+    pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)
 
     ret = pipe(prompt=prompt, 
                    image=image,