From 86276832e0204aebb6353b22dcce1c832a2ae865 Mon Sep 17 00:00:00 2001 From: Dalton Date: Tue, 19 Mar 2024 14:45:07 -0400 Subject: [PATCH] Update sd_hijack_ddpm_v1.py --- extensions-builtin/LDSR/sd_hijack_ddpm_v1.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py b/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py index 21e6c61c8..594dbc776 100644 --- a/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py +++ b/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py @@ -4,6 +4,7 @@ # Some models such as LDSR require VQ to work correctly # The classes are suffixed with "V1" and added back to the "ldm.models.diffusion.ddpm" module +import sys import torch import torch.nn as nn import numpy as np @@ -14,6 +15,8 @@ from contextlib import contextmanager from functools import partial from tqdm import tqdm from torchvision.utils import make_grid +import pytorch_lightning.utilities.rank_zero +sys.modules['pytorch_lightning.utilities.distributed'] = sys.modules['pytorch_lightning.utilities.rank_zero'] from pytorch_lightning.utilities.rank_zero import rank_zero_only from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config