remove download

sample/sample.py

@@ -12,14 +12,14 @@
     import utils
 
     from diffusion import create_diffusion
-    from download import find_model
+    from utils import find_model
 except:
     sys.path.append(os.path.split(sys.path[0])[0])
 
     import utils
 
     from diffusion import create_diffusion
-    from download import find_model
+    from utils import find_model
 
 import torch
 import argparse

sample/sample_ddp.py

@@ -16,7 +16,7 @@
 import torch
 sys.path.append(os.path.split(sys.path[0])[0])
 import torch.distributed as dist
-from download import find_model
+from utils import find_model
 from diffusion import create_diffusion
 from diffusers.models import AutoencoderKL
 from tqdm import tqdm

sample/sample_t2x.py

@@ -15,7 +15,6 @@
 
 import os, sys
 sys.path.append(os.path.split(sys.path[0])[0])
-from download import find_model
 from pipeline_videogen import VideoGenPipeline
 from models import get_models
 from utils import save_video_grid

utils.py

@@ -271,6 +271,20 @@ def save_video_grid(video, nrow=None):
 
     return video_grid
 
+def find_model(model_name):
+    """
+    Finds a pre-trained Latte model, downloading it if necessary. Alternatively, loads a model from a local path.
+    """
+    assert os.path.isfile(model_name), f'Could not find Latte checkpoint at {model_name}'
+    checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
+
+    if "ema" in checkpoint:  # supports checkpoints from train.py
+        print('Using Ema!')
+        checkpoint = checkpoint["ema"]
+    else:
+        print('Using model!')
+        checkpoint = checkpoint['model']
+    return checkpoint
 
 #################################################################################
 #                             MMCV  Utils                                    #