Update app.py

app.py

@@ -12,6 +12,7 @@ import torchtext
 from stat import ST_CTIME
 from datetime import datetime, timedelta
 import shutil
+torch.hub.download_url_to_file('https://i.imgur.com/tXrot31.jpg', 'gpu.jpg')
 # Images
 torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2021/08/04/14/16/tower-6521842_1280.jpg', 'tower.jpg')
 torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2017/08/31/05/36/buildings-2699520_1280.jpg', 'city.jpg')
@@ -357,7 +358,6 @@ def crop(img, h, w):
 def main(input_path, model_path, output_dir, need_animation=False, resize_h=None, resize_w=None, serial=False):
     if not os.path.exists(output_dir):
         os.mkdir(output_dir)
-
     for entry in os.listdir(output_dir):
         path = os.path.join(output_dir, entry)
         stats = os.stat(path)
@@ -485,12 +485,25 @@ def gradio_inference(image):
          resize_h=400,         # resize original input to this size. None means do not resize.
          resize_w=400,         # resize original input to this size. None means do not resize.
          serial=True)          # if need animation, serial must be True.
-
+inferences_running = 0
+def throttled_inference(image):
+    global inferences_running
+    current = inferences_running
+    if current >= 5:
+        print(f"Rejected inference when we already had {current} running")
+        return "gpu.jpg",Image.open("gpu.jpg")
+    print(f"Inference starting when we already had {current} running")
+    inferences_running += 1
+    try:
+        return gradio_inference(image)
+    finally:
+        print("Inference finished")
+        inferences_running -= 1
 title = "Paint Transformer"
 description = "Gradio demo for Paint Transformer: Feed Forward Neural Painting with Stroke Prediction. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.03798'>Paint Transformer: Feed Forward Neural Painting with Stroke Prediction</a> | <a href='https://github.com/Huage001/PaintTransformer'>Github Repo</a></p>"
 gr.Interface(
-    gradio_inference, 
+    throttled_inference, 
     gr.inputs.Image(type="file", label="Input"), 
     [gr.outputs.Image(type="file", label="Output GIF"),
     gr.outputs.Image(type="pil", label="Output Image")],
@@ -500,4 +513,5 @@ gr.Interface(
     examples=[
     ['city.jpg'],
     ['tower.jpg']
-    ],enable_queue=True).launch(debug=True)
+    ]
+    ).launch(debug=True)