fred-dev/comfy_ui_ali

1import importlib.util2import shutil3import os4import ctypes5import logging6 7 8def fix_pytorch_libomp():9    """10    Fix PyTorch libomp DLL issue on Windows by copying the correct DLL file if needed.11    """12    torch_spec = importlib.util.find_spec("torch")13    for folder in torch_spec.submodule_search_locations:14        lib_folder = os.path.join(folder, "lib")15        test_file = os.path.join(lib_folder, "fbgemm.dll")16        dest = os.path.join(lib_folder, "libomp140.x86_64.dll")17        if os.path.exists(dest):18            break19 20        with open(test_file, "rb") as f:21            contents = f.read()22            if b"libomp140.x86_64.dll" not in contents:23                break24        try:25            ctypes.cdll.LoadLibrary(test_file)26        except FileNotFoundError:27            logging.warning("Detected pytorch version with libomp issue, patching.")28            shutil.copyfile(os.path.join(lib_folder, "libiomp5md.dll"), dest)29