diff --git a/modules/initialize.py b/modules/initialize.py index 6c9c2388f..0365bbb30 100644 --- a/modules/initialize.py +++ b/modules/initialize.py @@ -51,6 +51,7 @@ def check_versions(): def initialize(): from modules import initialize_util initialize_util.fix_torch_version() + initialize_util.fix_pytorch_lightning() initialize_util.fix_asyncio_event_loop_policy() initialize_util.validate_tls_options() initialize_util.configure_sigint_handler() diff --git a/modules/initialize_util.py b/modules/initialize_util.py index b6767138d..8abe78814 100644 --- a/modules/initialize_util.py +++ b/modules/initialize_util.py @@ -24,6 +24,13 @@ def fix_torch_version(): torch.__long_version__ = torch.__version__ torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0) +def fix_pytorch_lightning(): + # Checks if pytorch_lightning.utilities.distributed already exists in the sys.modules cache + if 'pytorch_lightning.utilities.distributed' not in sys.modules: + import pytorch_lightning + # Lets the user know that the library was not found and then will set it to pytorch_lightning.utilities.rank_zero + print(f"Pytorch_lightning.distributed not found, attempting pytorch_lightning.rank_zero") + sys.modules["pytorch_lightning.utilities.distributed"] = pytorch_lightning.utilities.rank_zero def fix_asyncio_event_loop_policy(): """