mazpie commited on
Commit
a8f4e91
1 Parent(s): 67ebba2

Update demo/t2v.py

Browse files
Files changed (1) hide show
  1. demo/t2v.py +3 -0
demo/t2v.py CHANGED
@@ -57,6 +57,8 @@ class Text2Video():
57
  torch.cuda.empty_cache()
58
  self.agent.to('cuda')
59
  self.clip.to('cuda')
 
 
60
  print('start:', prompt, time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
61
  start = time.time()
62
 
@@ -95,6 +97,7 @@ class Text2Video():
95
  # Offload GPU
96
  self.agent.to('cpu')
97
  self.clip.to('cpu')
 
98
  return os.path.join(self.result_dir, f"{prompt_str}.mp4")
99
 
100
  def download_model(self, model_folder, model_filename):
 
57
  torch.cuda.empty_cache()
58
  self.agent.to('cuda')
59
  self.clip.to('cuda')
60
+ self.tokenizer.to('cuda')
61
+
62
  print('start:', prompt, time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
63
  start = time.time()
64
 
 
97
  # Offload GPU
98
  self.agent.to('cpu')
99
  self.clip.to('cpu')
100
+ self.tokenizer.to('cpu')
101
  return os.path.join(self.result_dir, f"{prompt_str}.mp4")
102
 
103
  def download_model(self, model_folder, model_filename):