add support for ZeroGPU

app.py

@@ -30,9 +30,9 @@ if __name__ == "__main__":
         gr.Interface(
             fn=send_to_model,
             inputs=[
-                gr.Video(value=None, label="Source Image"),
-                gr.Textbox(value="", label="Prompt"),
-                gr.Textbox(value="", label="Negative Prompt"),
+                gr.Video(value=None, label="source_video"),
+                gr.Textbox(value="", label="prompt"),
+                gr.Textbox(value="", label="neg_prompt"),
                 gr.Slider(
                     value   = 15,
                     minimum = 10,
@@ -41,8 +41,15 @@ if __name__ == "__main__":
                     label   = "guidance_scale",
                     info    = "The scale of the guidance field.",
                 ),
-                gr.Textbox(value=16, label="Video Length", info="The length of the video, must be less than 16 frames in the online demo to avoid timeout. However, you can run the model locally to process longer videos."),
-                gr.Dropdown(value=0, choices=[0, 1], label="Choose Option", info="Select 0 or 1."),
+                gr.Slider(
+                    value   = 16,
+                    minimum = 8,
+                    maximum = 32,
+                    step    = 2,
+                    label   = "video_length",
+                    info="The length of the video, must be less than 16 frames in the online demo to avoid timeout. However, you can run the model locally to process longer videos.",
+                ),
+                gr.Dropdown(value=0, choices=[0, 1], label="old_qk", info="Select 0 or 1."),
             ],
             outputs=[gr.Video(label="output", autoplay=True)],
             allow_flagging="never",

inference_utils.py

@@ -29,7 +29,8 @@ def init_pipeline(device):
 
     pipe = FlattenPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler, inverse_scheduler=inverse)
     pipe.enable_vae_slicing()
-    pipe.enable_xformers_memory_efficient_attention()
+    # xforers is not available in ZeroGPU?
+    # pipe.enable_xformers_memory_efficient_attention()
     pipe.to(device)
     return pipe
 
@@ -44,16 +45,16 @@ pipe = init_pipeline(device)
 
 
 def inference(
-    seed           = 66,
-    prompt         = None,
-    neg_prompt     = "",
-    guidance_scale = 10.0,
-    video_length   = 16,
-    video_path     = None,
-    output_dir     = None,
-    frame_rate     = 1,
-    fps            = 15,
-    old_qk         = 0,
+    seed          : int = 66,
+    prompt        : str = None,
+    neg_prompt    : str = "",
+    guidance_scale: float = 10.0,
+    video_length  : int = 16,
+    video_path    : str = None,
+    output_dir    : str = None,
+    frame_rate    : int = 1,
+    fps           : int = 15,
+    old_qk        : int = 0,
 ):
     generator = torch.Generator(device=device)
     generator.manual_seed(seed)
@@ -73,21 +74,28 @@ def inference(
 
     for k in trajectories.keys():
         trajectories[k] = trajectories[k].to(device)
-    sample = (pipe(
-        prompt,
-        video_length        = video_length,
-        frames              = real_frames,
-        num_inference_steps = sample_steps,
-        generator           = generator,
-        guidance_scale      = guidance_scale,
-        negative_prompt     = neg_prompt,
-        width               = width,
-        height              = height,
-        trajs               = trajectories,
-        output_dir          = "tmp/",
-        inject_step         = inject_step,
-        old_qk              = old_qk,
-    ).videos[0].permute(1, 2, 3, 0).cpu().numpy() * 255).astype(np.uint8)
+    sample = (
+        pipe(
+            prompt,
+            video_length        = video_length,
+            frames              = real_frames,
+            num_inference_steps = sample_steps,
+            generator           = generator,
+            guidance_scale      = guidance_scale,
+            negative_prompt     = neg_prompt,
+            width               = width,
+            height              = height,
+            trajs               = trajectories,
+            output_dir          = "tmp/",
+            inject_step         = inject_step,
+            old_qk              = old_qk,
+        )
+        .videos[0]
+        .permute(1, 2, 3, 0)
+        .cpu()
+        .numpy()
+        * 255
+    ).astype(np.uint8)
     temp_video_name = f"/tmp/{prompt}_{neg_prompt}_{str(guidance_scale)}_{time.time()}.mp4".replace(" ", "-")
     video_writer = imageio.get_writer(temp_video_name, fps=fps)
     for frame in sample: