From eeeb184071e2573f20e303cf0cbef622060bc2d0 Mon Sep 17 00:00:00 2001 From: Ali Hassani Date: Thu, 15 Jun 2023 16:34:44 -0700 Subject: [PATCH] Fix import error in newer version of lighting. Newer lighting versions move `rank_zero_only` from the distributed subpackage back to utilities. This commit adds a try-except to catch that import error, and retry with the alternative. --- ldm/models/diffusion/ddpm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ldm/models/diffusion/ddpm.py b/ldm/models/diffusion/ddpm.py index 3350c03..ce8c51b 100644 --- a/ldm/models/diffusion/ddpm.py +++ b/ldm/models/diffusion/ddpm.py @@ -17,7 +17,10 @@ from functools import partial import itertools from tqdm import tqdm from torchvision.utils import make_grid -from pytorch_lightning.utilities.distributed import rank_zero_only +try: + from pytorch_lightning.utilities import rank_zero_only +except: + from pytorch_lightning.utilities.distributed import rank_zero_only from omegaconf import ListConfig from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config