Support Iluvatar CoreX

README.md

@@ -297,6 +297,13 @@ For models compatible with Cambricon Extension for PyTorch (torch_mlu). Here's a
 2. Next, install the PyTorch(torch_mlu) following the instructions on the [Installation](https://www.cambricon.com/docs/sdk_1.15.0/cambricon_pytorch_1.17.0/user_guide_1.9/index.html)
 3. Launch ComfyUI by running `python main.py`
 
+#### Iluvatar Corex
+
+For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step guide tailored to your platform and installation method:
+
+1. Install the Iluvatar Corex Toolkit by adhering to the platform-specific instructions on the [Installation](https://support.iluvatar.com/#/DocumentCentre?id=1&nameCenter=2&productId=520117912052801536)
+2. Launch ComfyUI by running `python main.py`
+
 # Running
 
 ```python main.py```

comfy/model_management.py

@@ -128,6 +128,11 @@ def get_supported_float8_types():
 except:
     mlu_available = False
 
+try:
+    ixuca_available = hasattr(torch, "corex")
+except:
+    ixuca_available = False
+
 if args.cpu:
     cpu_state = CPUState.CPU
 
@@ -151,6 +156,12 @@ def is_mlu():
         return True
     return False
 
+def is_ixuca():
+    global ixuca_available
+    if ixuca_available:
+        return True
+    return False
+
 def get_torch_device():
     global directml_enabled
     global cpu_state
@@ -288,7 +299,7 @@ def is_amd():
         if torch_version_numeric[0] >= 2:
             if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
                 ENABLE_PYTORCH_ATTENTION = True
-    if is_intel_xpu() or is_ascend_npu() or is_mlu():
+    if is_intel_xpu() or is_ascend_npu() or is_mlu() or is_ixuca():
         if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
             ENABLE_PYTORCH_ATTENTION = True
 except:
@@ -1027,6 +1038,8 @@ def xformers_enabled():
         return False
     if is_mlu():
         return False
+    if is_ixuca():
+        return False
     if directml_enabled:
         return False
     return XFORMERS_IS_AVAILABLE
@@ -1062,6 +1075,8 @@ def pytorch_attention_flash_attention():
             return True
         if is_amd():
             return True #if you have pytorch attention enabled on AMD it probably supports at least mem efficient attention
+        if is_ixuca():
+            return True
     return False
 
 def force_upcast_attention_dtype():
@@ -1181,6 +1196,9 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
     if is_mlu():
         return True
 
+    if is_ixuca():
+        return True
+
     if torch.version.hip:
         return True
 
@@ -1240,6 +1258,9 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
 
     if is_ascend_npu():
         return True
+
+    if is_ixuca():
+        return True
 
     if is_amd():
         arch = torch.cuda.get_device_properties(device).gcnArchName

cuda_malloc.py

@@ -50,7 +50,33 @@ def enum_display_devices():
                 "GeForce GTX 1650", "GeForce GTX 1630", "Tesla M4", "Tesla M6", "Tesla M10", "Tesla M40", "Tesla M60"
                 }
 
+def _load_torch_submodule(filename):
+    """Helper to load and check a submodule from torch's installation"""
+    try:
+        torch_spec = importlib.util.find_spec("torch")
+        if torch_spec and torch_spec.submodule_search_locations:
+            for folder in torch_spec.submodule_search_locations:
+                file_path = os.path.join(folder, filename)
+                if os.path.isfile(file_path):
+                    spec = importlib.util.spec_from_file_location("torch_submodule_check", file_path)
+                    module = importlib.util.module_from_spec(spec)
+                    spec.loader.exec_module(module)
+                    return module
+        return None
+    except:
+        return None
+
+def is_ixuca():
+    ixuca_module = _load_torch_submodule("corex.py")
+    if ixuca_module and hasattr(ixuca_module, "corex"):
+        return bool(ixuca_module.corex)
+    else:
+        return False
+
 def cuda_malloc_supported():
+    if is_ixuca():
+        return False
+
     try:
         names = get_gpu_names()
     except:
@@ -64,20 +90,11 @@ def cuda_malloc_supported():
 
 
 if not args.cuda_malloc:
-    try:
-        version = ""
-        torch_spec = importlib.util.find_spec("torch")
-        for folder in torch_spec.submodule_search_locations:
-            ver_file = os.path.join(folder, "version.py")
-            if os.path.isfile(ver_file):
-                spec = importlib.util.spec_from_file_location("torch_version_import", ver_file)
-                module = importlib.util.module_from_spec(spec)
-                spec.loader.exec_module(module)
-                version = module.__version__
-        if int(version[0]) >= 2: #enable by default for torch version 2.0 and up
+    version_module = _load_torch_submodule("version.py")
+    if version_module and hasattr(version_module, "__version__"):
+        version = version_module.__version__
+        if version and int(version[0]) >= 2: #enable by default for torch version 2.0 and up
             args.cuda_malloc = cuda_malloc_supported()
-    except:
-        pass
 
 
 if args.cuda_malloc and not args.disable_cuda_malloc: