Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 os.system("pip install --upgrade torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html")
 os.system("git clone https://github.com/openai/CLIP")
 os.system("pip install -e ./CLIP")
-os.system("pip install einops ninja scipy numpy Pillow tqdm")
+os.system("pip install einops ninja scipy numpy Pillow tqdm imageio-ffmpeg imageio")
 import sys
 sys.path.append('./CLIP')
 import io
@@ -21,6 +21,8 @@ from tqdm.notebook import tqdm
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 from einops import rearrange
 import gradio as gr
+import imageio
+
 print(torch.cuda.get_device_name(0))
 device = torch.device('cuda:0')
 def fetch(url_or_path):
@@ -104,6 +106,7 @@ w_stds = G.mapping(zs, None).std(0)
     
     
 def inference(text):
+  all_frames = []
   target = clip_model.embed_text(text)
   steps = 20
   seed = 2
@@ -146,9 +149,14 @@ def inference(text):
     q_ema = q_ema * 0.9 + q * 0.1
     image = G.synthesis(q_ema * w_stds + G.mapping.w_avg, noise_mode='const')
     pil_image = TF.to_pil_image(image[0].add(1).div(2).clamp(0,1))
+    all_frames.append(pil_image)
     #os.makedirs(f'samples/{timestring}', exist_ok=True)
     #pil_image.save(f'samples/{timestring}/{i:04}.jpg')
-  return pil_image
+  writer = imageio.get_writer('test.mp4', fps=20)
+  for im in all_frames:
+      writer.append_data(np.array(im))
+  writer.close()
+  return pil_image, "test.mp4"
   
   
 title = "StyleGAN3+CLIP"
@@ -158,7 +166,7 @@ examples = [['elon musk']]
 gr.Interface(
     inference, 
     "text", 
-    gr.outputs.Image(type="pil", label="Output"),
+    [gr.outputs.Image(type="pil", label="Output"),"playable_video"],
     title=title,
     description=description,
     article=article,