Use comfy xformers check

This commit is contained in:
Kijai
2024-04-24 17:34:24 +03:00
parent e9f4f69be5
commit 98f08a4229
+2 -9
View File
@@ -7,6 +7,7 @@ import torch.nn.functional as F
from typing import Optional, Any
from .util import checkpoint
import model_management
class Conv2d(torch.nn.Conv2d):
def reset_parameters(self):
@@ -16,14 +17,6 @@ class Linear(torch.nn.Linear):
def reset_parameters(self):
return None
try:
import xformers
import xformers.ops
XFORMERS_IS_AVAILBLE = True
except:
print("WARNING: xformers is not available, inference might be slow.")
XFORMERS_IS_AVAILBLE = False
# CrossAttn precision handling
import os
@@ -261,7 +254,7 @@ class BasicTransformerBlock(nn.Module):
disable_self_attn=False,
):
super().__init__()
attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
attn_mode = "softmax-xformers" if model_management.XFORMERS_IS_AVAILABLE else "softmax"
assert attn_mode in self.ATTENTION_MODES
attn_cls = self.ATTENTION_MODES[attn_mode]
self.disable_self_attn = disable_self_attn