fix tqdm

Author: Dazhi Zhong

Diff for: cfg_sample.py

@@ -14,6 +14,7 @@
 from torchvision import transforms
 from torchvision.transforms import functional as TF
 from torchvision.utils import save_image
+from tqdm.auto import trange
 
 from CLIP import clip
 from diffusion import get_model, get_models, sampling, utils, download_model
@@ -27,9 +28,6 @@ def isnotebook():
 IS_NOTEBOOK = isnotebook()
 if IS_NOTEBOOK:
     from IPython import display
-    from tqdm.notebook import trange
-else:
-    from tqdm import trange
 
 
 MODULE_DIR = Path(__file__).resolve().parent

Diff for: clip_sample.py

@@ -14,6 +14,7 @@
 from torchvision import transforms
 from torchvision.transforms import functional as TF
 from torchvision.utils import save_image
+from tqdm.auto import trange
 
 from CLIP import clip
 from diffusion import get_model, get_models, sampling, utils, download_model
@@ -28,9 +29,6 @@ def isnotebook():
 IS_NOTEBOOK = isnotebook()
 if IS_NOTEBOOK:
     from IPython import display
-    from tqdm.notebook import trange
-else:
-    from tqdm import trange