Disable xformers memory-efficient attention in VAE (unsupported on Blackwell/sm_120 GPUs)

vae_wrapper.py

@@ -15,7 +15,7 @@ def default(value, default_value):
 
 def load_stable_model(model_path):
     vae_model = StableDiffusionPipeline.from_pretrained(model_path)
-    vae_model.set_use_memory_efficient_attention_xformers(True)
+    vae_model.set_use_memory_efficient_attention_xformers(False)
     return vae_model.vae
 
 
@@ -74,7 +74,7 @@ class VaeWrapper(nn.Module):
         if latent_type == "stable":
             vae_model = load_stable_model("stabilityai/stable-diffusion-x4-upscaler")
             vae_model.enable_slicing()
-            vae_model.set_use_memory_efficient_attention_xformers(True)
+            vae_model.set_use_memory_efficient_attention_xformers(False)
             self.down_factor = 4
         elif latent_type == "video":
             vae_model = AutoencoderKLTemporalDecoder.from_pretrained(
@@ -83,7 +83,7 @@ class VaeWrapper(nn.Module):
                 torch_dtype=torch.float16 if variant == "fp16" else torch.float32,
                 variant="fp16" if variant == "fp16" else None,
             )
-            vae_model.set_use_memory_efficient_attention_xformers(True)
+            vae_model.set_use_memory_efficient_attention_xformers(False)
             self.down_factor = 8
         elif latent_type == "refiner":
             vae_model = AutoencoderKL.from_pretrained(
@@ -92,7 +92,7 @@ class VaeWrapper(nn.Module):
                 revision=None,
             )
             vae_model.enable_slicing()
-            vae_model.set_use_memory_efficient_attention_xformers(True)
+            vae_model.set_use_memory_efficient_attention_xformers(False)
             self.down_factor = 8
 
         vae_model.eval()