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