Update demo/t2v.py

demo/t2v.py

@@ -52,12 +52,11 @@ class Text2Video():
         if not os.path.exists(self.result_dir):
             os.mkdir(self.result_dir)
 
-        self.agent.to('cuda')
-        self.clip.to('cuda')
-
     @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()
@@ -94,6 +93,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):