Update app.py

app.py

@@ -116,13 +116,15 @@ os.makedirs("./gradio_tmp", exist_ok=True)
 upscale_model = load_sd_upscale("model_real_esran/RealESRGAN_x4.pth", device)
 frame_interpolation_model = load_rife_model("model_rife")
 
-def infer(
+@spaces.GPU(duration=300)
+def generate(
     prompt: str,
     image_input: str,
     num_inference_steps: int,
     guidance_scale: float,
     seed: int = 42,
-    progress=gr.Progress(track_tqdm=True),
+    scale_status: bool = False,
+    rife_status: bool = False,
 ):
     if seed == -1:
         seed = random.randint(0, 2**8 - 1)
@@ -167,6 +169,12 @@ def infer(
     ).frames
     
     free_memory()
+    
+    if scale_status:
+        video_pt = upscale_batch_and_concatenate(upscale_model, video_pt, device)
+    if rife_status:
+        video_pt = rife_inference_with_latents(frame_interpolation_model, video_pt)
+
     return (video_pt, seed)
 
 
@@ -320,8 +328,8 @@ with gr.Blocks() as demo:
     </table>
         """)
 
-    @spaces.GPU(duration=180)
-    def generate(
+
+    def run(
         prompt,
         image_input,
         seed_value,
@@ -329,18 +337,15 @@ with gr.Blocks() as demo:
         rife_status,
         progress=gr.Progress(track_tqdm=True)
     ):
-        latents, seed = infer(
+        latents, seed = generate(
             prompt,
             image_input,
             num_inference_steps=50,
             guidance_scale=7.0,
             seed=seed_value,
-            progress=progress,
-        )
-        if scale_status:
-            latents = upscale_batch_and_concatenate(upscale_model, latents, device)
-        if rife_status:
-            latents = rife_inference_with_latents(frame_interpolation_model, latents)
+            scale_status=scale_status,
+            rife_status=rife_status,
+        )   
 
         batch_size = latents.shape[0]
         batch_video_frames = []
@@ -361,11 +366,11 @@ with gr.Blocks() as demo:
         return video_path, video_update, gif_update, seed_update
 
     generate_button.click(
-        generate,
+        fn=run,
         inputs=[prompt, image_input, seed_param, enable_scale, enable_rife],
         outputs=[video_output, download_video_button, download_gif_button, seed_text],
     )
 
 if __name__ == "__main__":
     demo.queue(max_size=15)
-    demo.launch()
+    demo.launch()