Initial video commit

app.py

@@ -142,7 +142,8 @@ def optimize_network(
     neuron,
     class_token,
     maximize,
-    display_rate = 20
+    display_rate,
+    video_writer
 ):
     global itt
     itt = 0
@@ -210,7 +211,7 @@ def optimize_network(
                         if display_augs:
                             aug_grid = torchvision.utils.make_grid(cutouts, nrow=math.ceil(math.sqrt(cutn)))
                             display.display(TF.to_pil_image(aug_grid.clamp(0, 1)))
-                    if save_progress_video and itt > 15:
+                    if save_progress_video:
                         video_writer.append_data(np.asarray(image))
 
             if anneal_lr:
@@ -240,13 +241,13 @@ def inference(
     num_iterations = int(num_iterations)
     neuron = int(neuron)
     display_rate = int(display_rate)
-    
+
     opt_type = 'MADGRAD'
     seed = 20
     save_progress_video = True
     timestring = time.strftime('%Y%m%d%H%M%S')
     if save_progress_video:
-        video_writer = imageio.get_writer(f'dip_{timestring}.mp4', mode='I', fps=30, codec='libx264', quality=7, pixelformat='yuv420p')
+        video_writer = imageio.get_writer(f'video.mp4', mode='I', fps=30, codec='libx264', quality=7, pixelformat='yuv420p')
     
     # Begin optimization / generation
     gc.collect()
@@ -261,17 +262,18 @@ def inference(
             neuron,
             class_token,
             maximize,
-            display_rate
+            display_rate,
+            video_writer = video_writer
         )
 
     out.save(f'dip_{timestring}.png', quality=100)
     if save_progress_video:
         video_writer.close()
-    return out 
+    return out, 'video.mp4'
 
 iface = gr.Interface(fn=inference, 
                      inputs=["number", "number", "number", "number", "number", 
                                 gr.inputs.Checkbox(default=False, label="class_token"),
                                 gr.inputs.Checkbox(default=True, label="maximise"),
                                 "number"], 
-                    outputs="image").launch()
+                    outputs=["image","video"]).launch()