From 8b091f02def32bce977802463adefd44c5f9cde1 Mon Sep 17 00:00:00 2001
From: Max Tretikov <max@tretikov.com>
Date: Fri, 14 Jun 2024 14:09:46 -0600
Subject: [PATCH] Add xformer.ops imports

---
 comfy/ldm/modules/attention.py              | 1 +
 comfy/ldm/modules/diffusionmodules/model.py | 4 ++--
 comfy/model_management.py                   | 3 ++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
index ad3b6c471..10c552a3e 100644
--- a/comfy/ldm/modules/attention.py
+++ b/comfy/ldm/modules/attention.py
@@ -12,6 +12,7 @@ from ... import model_management
 
 if model_management.xformers_enabled():
     import xformers  # pylint: disable=import-error
+    import xformers.ops # pylint: disable=import-error
 
 from ...cli_args import args
 from ... import ops
diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py
index bab6b0ca9..8e11c1eab 100644
--- a/comfy/ldm/modules/diffusionmodules/model.py
+++ b/comfy/ldm/modules/diffusionmodules/model.py
@@ -11,8 +11,8 @@ from .... import ops
 ops = ops.disable_weight_init
 
 if model_management.xformers_enabled_vae():
-    import xformers
-    import xformers.ops
+    import xformers  # pylint: disable=import-error
+    import xformers.ops # pylint: disable=import-error
 
 def get_timestep_embedding(timesteps, embedding_dim):
     """
diff --git a/comfy/model_management.py b/comfy/model_management.py
index a4840f688..e9287bb7e 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -158,7 +158,8 @@ if args.disable_xformers:
     XFORMERS_IS_AVAILABLE = False
 else:
     try:
-        import xformers # pylint: disable=import-error 
+        import xformers # pylint: disable=import-error
+        import xformers.ops # pylint: disable=import-error
 
         XFORMERS_IS_AVAILABLE = True
         try: