another way of stopping

1 files changed, 25 insertions, 21 deletions

diff --git a/cli/app/search/live.py b/cli/app/search/live.py
index 25d6ec7..627822b 100644
--- a/cli/app/search/live.py
+++ b/cli/app/search/live.py
@@ -198,6 +198,7 @@ class InterpolatorVariable:
 
 class Interpolator:
   def __init__(self):
+    self.stopped = False
     self.opts = {}
     self.sin_params = {}
     self.lerp_params = {}
@@ -376,6 +377,8 @@ class Interpolator:
       self.set_category(payload)
     if cmd == 'setEncoding':
       self.set_encoding(json.loads(payload))
+    if cmd == 'stop':
+      self.stopped = True
     pass
 
 # --------------------------
@@ -459,6 +462,9 @@ class Listener:
         print("Loading network...")
       elif i == 1:
         print("Processing!")
+      elif self.stopped:
+        print("Stopping...")
+        return
       gen_time = time.time()
       gen_images = interpolator.on_step(i, self.sess)
       if i == 0:
@@ -469,32 +475,30 @@ class Listener:
         print("Exiting...")
         break
 
-      if (i % 100) == 0 or i == 1:
-        print("Step {}. Generation time: {:.2f}s".format(i, time.time() - gen_time))
-
       to_pil_time = time.time()
       out_img = vs.data2pil(gen_images[0])
       to_pil_time_total += time.time() - to_pil_time
 
-      if out_img is not None:
-        save_time = time.time()
-        # out_img.save(os.path.join(path_out, "frame_{:05d}.png".format(i)), format='png', compression_level=3)
-        out_img.save(pipe.stdin, format='png', compression_level=3)
-        save_time_total += time.time() - save_time
+      if out_img is None:
+        return
+      save_time = time.time()
+      # out_img.save(os.path.join(path_out, "frame_{:05d}.png".format(i)), format='png', compression_level=3)
+      out_img.save(pipe.stdin, format='png', compression_level=3)
+      save_time_total += time.time() - save_time
 
-        resize_time = time.time()
-        img_to_send = out_img.resize((256, 256), Image.BICUBIC)
-        resize_time_total += time.time() - resize_time
+      resize_time = time.time()
+      img_to_send = out_img.resize((256, 256), Image.BICUBIC)
+      resize_time_total += time.time() - resize_time
 
-        meta = {
-          'i': i,
-          'sequence_i': i,
-          'skip_i': 0,
-          'sequence_len': 99999,
-        }
-        send_time = time.time()
-        self.rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, img_to_send, 'jpg')
-        send_time_total += time.time() - send_time
+      meta = {
+        'i': i,
+        'sequence_i': i,
+        'skip_i': 0,
+        'sequence_len': 99999,
+      }
+      send_time = time.time()
+      self.rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, img_to_send, 'jpg')
+      send_time_total += time.time() - send_time
       if (i % 100) == 0 or i == 1:
-        print("gen: {:2f}, pil: {:2f}, save: {:2f}, resize: {:2f}, send: {:2f}".format(gen_time_total / i, to_pil_time_total / i, save_time_total / i, resize_time_total / i, send_time_total / i))
+        print("step: {}, gen: {:2f}, pil: {:2f}, save: {:2f}, resize: {:2f}, send: {:2f}".format(i, gen_time_total / i, to_pil_time_total / i, save_time_total / i, resize_time_total / i, send_time_total / i))