Update app.py

app.py

@@ -94,13 +94,25 @@ def sample_text2video(model, prompt, n_samples, batch_size,
     assert(all_videos.shape[0] >= n_samples)
     return all_videos
 
+def save_results(videos, 
+                 save_name="results", save_fps=8, save_mp4=True, 
+                 save_npz=False, save_mp4_sheet=False, save_jpg=False
+                 ):
+    
+    save_subdir = os.path.join("videos")
+    os.makedirs(save_subdir, exist_ok=True)
+    for i in range(videos.shape[0]):
+        npz_to_video_grid(videos[i:i+1,...], 
+                          os.path.join(save_subdir, f"{save_name}_{i:03d}.mp4"), 
+                          fps=save_fps)
+    return 'Successfully saved videos in ' + save_subdir
 
 def get_video(prompt):
     samples = sample_text2video(model, prompt, n_samples = 1, batch_size = 1,
                           sampler=ddim_sampler,
                           )
-    
-    return samples[0]
+
+    return save_results(samples)
 
 prompt_inp = gr.Textbox(label = "Prompt")
 iface = gr.Interface(fn=get_video, inputs=[prompt_inp], outputs="text")