Disable pytorch attention in VAE for AMD.

comfyanonymous · Feb 14, 2025 · 1cd6cd6 · 1cd6cd6
1 parent d7b4bf2
commit 1cd6cd6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py
@@ -297,7 +297,7 @@ def vae_attention():
     if model_management.xformers_enabled_vae():
         logging.info("Using xformers attention in VAE")
         return xformers_attention
-    elif model_management.pytorch_attention_enabled():
+    elif model_management.pytorch_attention_enabled_vae():
         logging.info("Using pytorch attention in VAE")
         return pytorch_attention
     else:

diff --git a/comfy/model_management.py b/comfy/model_management.py
@@ -912,6 +912,11 @@ def pytorch_attention_enabled():
     global ENABLE_PYTORCH_ATTENTION
     return ENABLE_PYTORCH_ATTENTION
 
+def pytorch_attention_enabled_vae():
+    if is_amd():
+        return False  # enabling pytorch attention on AMD currently causes crash when doing high res
+    return pytorch_attention_enabled()
+
 def pytorch_attention_flash_attention():
     global ENABLE_PYTORCH_ATTENTION
     if ENABLE_PYTORCH_ATTENTION: