Update app.py

app.py

@@ -31,6 +31,7 @@ import glob
 import pathlib
 from functools import partial
 from pprint import pprint
+from multiprocessing.connection import Client, Listener
 
 import numpy as np
 from PIL import Image
@@ -150,6 +151,11 @@ opt.excluded_keys = ['inpainting_mode', 'is_running', 'active_palettes', 'curren
 opt.prep_time = 20
 
 
+### Shared memory hack for ZeroGPU
+opt.address = ('localhost', 6000)
+ = b'secret password'
+
+
 ### Event handlers
 
 def add_palette(state):
@@ -385,6 +391,11 @@ def register(state, drawpad, model):
 
 @spaces.GPU(duration=(opt.prep_time + opt.run_time + 5))
 def run(state, drawpad):
+    # ZeroGPU hack.
+    listener = Listener(opt.address, authkey=opt.authkey)
+    conn = listener.accept()
+
+    # Reset model.
     model.device = torch.device('cuda')
     model.reset_seed(model.generator, opt.seed)
     model.reset_latent()
@@ -395,6 +406,17 @@ def run(state, drawpad):
 
     tic = time.time()
     while True:
+        # Receive real-time mask inputs from the main process.
+        msg = conn.recv()
+        print(msg + '  Received!!!')
+        # for i in range(opt.max_palettes):
+        #     model.update_single_layer(
+        #         idx=i,
+        #         mask=masks[i],
+        #         mask_strength=mask_strengths[i],
+        #         mask_std=mask_stds[i],
+        #     )
+
         yield [state, model()]
         toc = time.time()
         tdelta = toc - tic
@@ -412,12 +434,14 @@ def show_element():
     return gr.update(visible=True)
 
 
-@spaces.GPU
 def draw(state, drawpad):
     if not state.is_running:
         print('[WARNING] Streaming is currently off, update ignored.')
         return
 
+    # ZeroGPU hack.
+    conn = Client(opt.address, authkey=opt.authkey)
+
     user_input = np.asarray(drawpad['layers'][0]) # (H, W, 4)
     foreground_mask = torch.tensor(user_input[..., -1])[None, None] # (1, 1, H, W)
     user_input = torch.tensor(user_input[..., :-1]) # (H, W, 3)
@@ -441,13 +465,15 @@ def draw(state, drawpad):
     #     mask_strengths = [1] + [state.mask_strengths[v] for v in has_masks]
     #     mask_stds = [0] + [state.mask_stds[v] for v in has_masks]
 
-    for i in range(len(has_masks)):
-        model.update_single_layer(
-            idx=i,
-            mask=masks[i],
-            mask_strength=mask_strengths[i],
-            mask_std=mask_stds[i],
-        )
+    # for i in range(len(has_masks)):
+    #     model.update_single_layer(
+    #         idx=i,
+    #         mask=masks[i],
+    #         mask_strength=mask_strengths[i],
+    #         mask_std=mask_stds[i],
+    #     )
+    conn.send('Hello!!!!')
+    conn.close()
 
 ### Load examples