set 1 min timeout

visualizer_drag_gradio.py

@@ -5,6 +5,7 @@ import os.path as osp
 from argparse import ArgumentParser
 from functools import partial
 from pathlib import Path
+import time
 
 import gradio as gr
 import numpy as np
@@ -34,6 +35,7 @@ cache_dir = args.cache_dir
 
 device = 'cuda'
 IS_SPACE = "DragGan/DragGan" in os.environ.get('SPACE_ID', '')
+TIMEOUT = 60 if IS_SPACE else 0
 
 
 def reverse_point_pairs(points):
@@ -553,7 +555,13 @@ with gr.Blocks() as app:
             print(f'    Source: {p_in_pixels}')
             print(f'    Target: {t_in_pixels}')
             step_idx = 0
+            last_time = time.time()
             while True:
+                # add a TIMEOUT break
+                print(f'Running time: {time.time() - last_time}')
+                if time.time() - last_time > TIMEOUT:
+                    print('Timeout break!')
+                    break
                 if global_state["temporal_params"]["stop"]:
                     break
 
@@ -907,5 +915,5 @@ with gr.Blocks() as app:
 
 print("SHAReD: Start app", parser.parse_args())
 gr.close_all()
-app.queue(concurrency_count=2, max_size=200, api_open=False)
+app.queue(concurrency_count=3, max_size=200, api_open=False)
 app.launch(share=args.share, show_api=False)