fix gpu OOM

i2v_test.py

@@ -22,7 +22,7 @@ class Image2Video():
         model_list = []
         for gpu_id in range(gpu_num):
             model = instantiate_from_config(model_config)
-            model = model.cuda(gpu_id)
+            # model = model.cuda(gpu_id)
             assert os.path.exists(ckpt_path), "Error: checkpoint Not Found!"
             model = load_model_checkpoint(model, ckpt_path)
             model.eval()
@@ -38,6 +38,7 @@ class Image2Video():
         if steps > 60:
             steps = 60 
         model = self.model_list[gpu_id]
+        model = model.cuda()
         batch_size=1
         channels = model.model.diffusion_model.in_channels
         frames = model.temporal_length
@@ -65,6 +66,7 @@ class Image2Video():
 
         save_videos(batch_samples, self.result_dir, filenames=[prompt_str], fps=self.save_fps)
         print(f"Saved in {prompt_str}. Time used: {(time.time() - start):.2f} seconds")
+        model = model.cpu()
         return os.path.join(self.result_dir, f"{prompt_str}.mp4")
     
     def download_model(self):

t2v_test.py

@@ -20,7 +20,7 @@ class Text2Video():
         model_list = []
         for gpu_id in range(gpu_num):
             model = instantiate_from_config(model_config)
-            model = model.cuda(gpu_id)
+            # model = model.cuda(gpu_id)
             assert os.path.exists(ckpt_path), "Error: checkpoint Not Found!"
             model = load_model_checkpoint(model, ckpt_path)
             model.eval()
@@ -36,6 +36,7 @@ class Text2Video():
         if steps > 60:
             steps = 60 
         model = self.model_list[gpu_id]
+        model = model.cuda()
         batch_size=1
         channels = model.model.diffusion_model.in_channels
         frames = model.temporal_length
@@ -56,6 +57,7 @@ class Text2Video():
 
         save_videos(batch_samples, self.result_dir, filenames=[prompt_str], fps=self.save_fps)
         print(f"Saved in {prompt_str}. Time used: {(time.time() - start):.2f} seconds")
+        model=model.cpu()
         return os.path.join(self.result_dir, f"{prompt_str}.mp4")
     
     def download_model(self):