Update app.py

app.py

@@ -26,6 +26,7 @@ from torchvision import transforms
 
 import spaces
 from huggingface_hub import snapshot_download
+import time
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -116,9 +117,9 @@ def init_pipe():
             controlnet_transformer,
     )
 
-    pipe.vae.enable_slicing()
-    pipe.vae.enable_tiling()
-    pipe.enable_model_cpu_offload()
+    # pipe.vae.enable_slicing()
+    # pipe.vae.enable_tiling()
+    # pipe.enable_model_cpu_offload()
 
     return pipe
     
@@ -129,10 +130,11 @@ def inference(source_images,
         pipe, vae, guidance_scale, 
         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)
+
+    pipe.to(DEVICE)
+    # 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(DEVICE)
@@ -161,6 +163,7 @@ def inference(source_images,
             image_latents = None
             latents = source_latents
 
+    a=time.perf_counter()
     video = pipe(
         prompt = text_prompt,
         negative_prompt = negative_prompt,
@@ -169,11 +172,13 @@ def inference(source_images,
         height = h,
         width = w,
         num_frames = f,
-        num_inference_steps = 20,
+        num_inference_steps = 30,
         interval = 6,
         guidance_scale = guidance_scale,
         generator = torch.Generator(device=DEVICE).manual_seed(random_seed)
     ).frames[0]
+    b=time.perf_counter()
+    print(f"Denoise 30 steps in {b-a}s")
 
     return video