Update demo/t2v.py

demo/t2v.py

@@ -36,14 +36,14 @@ class Text2Video():
             self.download_model(model_folder, model_filename)
         if not os.path.isfile(os.path.join(model_folder, 'InternVideo2-stage2_1b-224p-f4.pt')):
             self.download_internvideo2(model_folder)
-        self.agent = torch.load(os.path.join(model_folder, model_filename))
+        self.agent = torch.load(os.path.join(model_folder, model_filename),map_location='cpu')
         model_name = 'internvideo2'
 
         # Get ViCLIP
         viclip_global_instance = ViCLIPGlobalInstance(model_name)
         if not viclip_global_instance._instantiated:
             print("Instantiating InternVideo2")
-            viclip_global_instance.instantiate()
+            viclip_global_instance.instantiate(device='cpu')
         self.clip = viclip_global_instance.viclip
         self.tokenizer = viclip_global_instance.viclip_tokenizer
 
@@ -51,8 +51,11 @@ class Text2Video():
         if not os.path.exists(self.result_dir):
             os.mkdir(self.result_dir)
 
+    @spaces.GPU
     def get_prompt(self, prompt, duration):
         torch.cuda.empty_cache()
+        self.agent.to('cuda')
+        self.clip.to('cuda')
         print('start:', prompt, time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
         start = time.time()
 
@@ -88,6 +91,9 @@ class Text2Video():
 
         save_videos(prior_recon.unsqueeze(0), self.result_dir, filenames=[prompt_str], fps=15)
         print(f"Saved in {prompt_str}.mp4. Time used: {(time.time() - start):.2f} seconds")
+        # Offload GPU
+        self.agent.to('cpu')
+        self.clip.to('cpu')
         return os.path.join(self.result_dir, f"{prompt_str}.mp4")
     
     def download_model(self, model_folder, model_filename):