Update app.py

app.py

@@ -28,7 +28,7 @@ import subprocess
 import spaces
 from huggingface_hub import snapshot_download
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def download_model():
@@ -50,7 +50,6 @@ def get_prompt(file:str):
         a=f.readlines()
     return a #a[0]:positive prompt, a[1] negative prompt
 
-@spaces.GPU(duration=120)
 def init_pipe():
     def unwarp_model(state_dict):
         new_state_dict = {}
@@ -73,7 +72,7 @@ def init_pipe():
     )
 
     text_encoder = T5EncoderModel.from_pretrained(f"./cogvideox-5b-{key}/", subfolder="text_encoder", torch_dtype=torch.float16)
-    vae = AutoencoderKLCogVideoX.from_pretrained(f"./cogvideox-5b-{key}/", subfolder="vae", torch_dtype=torch.float16).to(device)
+    vae = AutoencoderKLCogVideoX.from_pretrained(f"./cogvideox-5b-{key}/", subfolder="vae", torch_dtype=torch.float16)
     tokenizer = T5Tokenizer.from_pretrained(f"./cogvideox-5b-{key}/tokenizer", torch_dtype=torch.float16)
 
 
@@ -104,8 +103,8 @@ def init_pipe():
     transformer.load_state_dict(transformer_state_dict, strict=True)
     controlnet_transformer.load_state_dict(controlnet_transformer_state_dict, strict=True)
 
-    transformer = transformer.half().to(device)
-    controlnet_transformer = controlnet_transformer.half().to(device)
+    transformer = transformer.half()
+    controlnet_transformer = controlnet_transformer.half()
 
     vae = vae.eval()
     text_encoder = text_encoder.eval()
@@ -134,6 +133,10 @@ def inference(source_images,
         h, w, random_seed)->List[PIL.Image.Image]:
     torch.manual_seed(random_seed)
 
+    pipe.vae.to(DEVICE)
+    pipe.transformer.to(DEVICE)
+    pipe.controlnet_transformer.to(DEVICE)
+
     source_pixel_values = source_images/127.5 - 1.0
     source_pixel_values = source_pixel_values.to(torch.float16).to("cuda:0")
     if target_images is not None: