Update app.py

app.py

@@ -389,11 +389,11 @@ def register(state, drawpad, model):
     return state
 
 
-@spaces.GPU(duration=(opt.prep_time + opt.run_time + 5))
+# @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()
+    # listener = Listener(opt.address, authkey=opt.authkey)
+    # conn = listener.accept()
 
     # Reset model.
     model.device = torch.device('cuda')
@@ -407,16 +407,16 @@ def run(state, drawpad):
     tic = time.time()
     while True:
         # Receive real-time mask inputs from the main process.
-        data = conn.recv()
-        if data is not None:
-            print('Received data!!!')
-            for i in range(opt.max_palettes):
-                model.update_single_layer(
-                    idx=i,
-                    mask=data['masks'][i],
-                    mask_strength=data['mask_strengths'][i],
-                    mask_std=data['mask_stds'][i],
-                )
+        # data = conn.recv()
+        # if data is not None:
+        #     print('Received data!!!')
+        #     for i in range(opt.max_palettes):
+        #         model.update_single_layer(
+        #             idx=i,
+        #             mask=data['masks'][i],
+        #             mask_strength=data['mask_strengths'][i],
+        #             mask_std=data['mask_stds'][i],
+        #         )
 
         yield [state, model()]
         toc = time.time()
@@ -441,7 +441,7 @@ def draw(state, drawpad):
     #     return
 
     # ZeroGPU hack.
-    conn = Client(opt.address, authkey=opt.authkey)
+    # 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)
@@ -466,20 +466,20 @@ 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],
-    #     )
-    data = dict(
-        masks=masks,
-        mask_strengths=mask_strengths,
-        mask_stds=mask_stds,
-    )
-    conn.send(data)
-    conn.close()
+    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],
+        )
+    # data = dict(
+    #     masks=masks,
+    #     mask_strengths=mask_strengths,
+    #     mask_stds=mask_stds,
+    # )
+    # conn.send(data)
+    # conn.close()
 
 ### Load examples