diff --git a/SUPIR/utils/tilevae.py b/SUPIR/utils/tilevae.py index de92ef8..c7e2c81 100644 --- a/SUPIR/utils/tilevae.py +++ b/SUPIR/utils/tilevae.py @@ -66,20 +66,20 @@ import torch import torch.version import torch.nn.functional as F from einops import rearrange -#from diffusers.utils.import_utils import is_xformers_available - -#import SUPIR.utils.devices as devices import comfy.model_management device = comfy.model_management.get_torch_device() -try: - import xformers - import xformers.ops - XFORMERS_IS_AVAILABLE = True -except: +if comfy.model_management.XFORMERS_IS_AVAILABLE: + try: + import xformers + import xformers.ops + XFORMERS_IS_AVAILABLE = True + except: + XFORMERS_IS_AVAILABLE = False + print("no module 'xformers'. Processing without...") +else: XFORMERS_IS_AVAILABLE = False - print("no module 'xformers'. Processing without...") sd_flag = True diff --git a/nodes.py b/nodes.py index 2b9a78c..083bb6f 100644 --- a/nodes.py +++ b/nodes.py @@ -21,15 +21,6 @@ from transformers import ( ) script_directory = os.path.dirname(os.path.abspath(__file__)) -try: - import xformers - import xformers.ops - - XFORMERS_IS_AVAILABLE = True -except: - XFORMERS_IS_AVAILABLE = False - - def dummy_build_vision_tower(*args, **kwargs): # Monkey patch the CLIP class before you create an instance. return None @@ -236,7 +227,7 @@ class SUPIR_Upscale: config.model.params.sampler_config.target = f".sgm.modules.diffusionmodules.sampling.{sampler}" print("Using non-tiled sampling") - if XFORMERS_IS_AVAILABLE: + if mm.XFORMERS_IS_AVAILABLE: config.model.params.control_stage_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.network_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.first_stage_config.params.ddconfig.attn_type = "vanilla-xformers" diff --git a/nodes_v2.py b/nodes_v2.py index 94f6af1..23b3595 100644 --- a/nodes_v2.py +++ b/nodes_v2.py @@ -21,15 +21,6 @@ from transformers import ( ) script_directory = os.path.dirname(os.path.abspath(__file__)) -try: - import xformers - import xformers.ops - - XFORMERS_IS_AVAILABLE = True -except: - XFORMERS_IS_AVAILABLE = False - - def dummy_build_vision_tower(*args, **kwargs): # Monkey patch the CLIP class before you create an instance. return None @@ -644,7 +635,8 @@ class SUPIR_model_loader: config = OmegaConf.load(config_path) - if XFORMERS_IS_AVAILABLE: + if mm.XFORMERS_IS_AVAILABLE: + print("Using XFORMERS") config.model.params.control_stage_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.network_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.first_stage_config.params.ddconfig.attn_type = "vanilla-xformers" @@ -784,8 +776,8 @@ class SUPIR_model_loader_v2: mm.soft_empty_cache() config = OmegaConf.load(config_path) - - if XFORMERS_IS_AVAILABLE: + if mm.XFORMERS_IS_AVAILABLE: + print("Using XFORMERS") config.model.params.control_stage_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.network_config.params.spatial_transformer_attn_type = "softmax-xformers" config.model.params.first_stage_config.params.ddconfig.attn_type = "vanilla-xformers"