Update app.py

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 from torchvision import transforms
 import torchtext
 
+torch.hub.download_url_to_file('https://i.imgur.com/WEHmKef.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')
@@ -512,7 +513,6 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
                save_all=True, append_images=all_frames[1:], optimize=False, duration=40, loop=0)
         return os.path.join(frame_dir, "animation.gif"), tensor_to_pil
 
-
 def gradio_inference(image):
     return main(input_path=image.name,
          model_path='model.pth',
@@ -522,12 +522,26 @@ def gradio_inference(image):
          resize_w=300,         # resize original input to this size. None means do not resize.
          serial=True)          # if need animation, serial must be True.
 
+def throttled_inference(image):
+    global inferences_running
+    current = inferences_running
+    if current >= 4:
+        print(f"Rejected inference when we already had {current} running")
+        return load_image("./gpu.jpg")
+    print(f"Inference starting when we already had {current} running")
+    inferences_running += 1
+    try:
+        return 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")],