Add timeout to prediction polling to prevent queue lock

Previously, the Gradio app would poll the Cog prediction status indefinitely, waiting for a "succeeded" or "failed" response. If the Cog server became unresponsive or a prediction hung (e.g. due to a bad input or internal model error), the polling loop would never exit.

This behavior caused the Gradio queue to back up and eventually stall completely under high load, since all workers could get stuck waiting for unresponsive predictions.

This change introduces a timeout mechanism (60 seconds) in the polling loop. If the prediction doesn't complete within that time, an error is raised and the request is gracefully terminated. This helps avoid queue deadlocks and improves reliability under concurrent usage.

app.py

@@ -60,46 +60,58 @@ def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
     headers = {'Content-Type': 'application/json'}
 
     payload = {"input": {}}
-    
-    
     base_url = "http://0.0.0.0:7860"
+
     for i, key in enumerate(names):
         value = args[i]
-        if value and (os.path.exists(str(value))):
+        if value and os.path.exists(str(value)):
             value = f"{base_url}/gradio_api/file=" + value
         if value is not None and value != "":
             payload["input"][key] = value
 
-    time.sleep(1.0) 
+    time.sleep(1.0)
     response = requests.post("http://0.0.0.0:5000/predictions", headers=headers, json=payload)
-    
-    
+
     if response.status_code == 201:
-        time.sleep(1.0)
         follow_up_url = response.json()["urls"]["get"]
-        response = requests.get(follow_up_url, headers=headers)
-        while response.json()["status"] != "succeeded":
-            if response.json()["status"] == "failed":
-                raise gr.Error("The submission failed!")
+
+        # ⏰ Timeout logic
+        max_wait_seconds = 60
+        poll_interval = 1
+        start_time = time.time()
+
+        while True:
             response = requests.get(follow_up_url, headers=headers)
-            
+            try:
+                response_json = response.json()
+            except ValueError:
+                raise gr.Error("Cog server response is not valid JSON.")
+
+            status = response_json.get("status")
+            if status == "succeeded":
+                break
+            if status == "failed":
+                raise gr.Error("The submission failed.")
+            if time.time() - start_time > max_wait_seconds:
+                raise gr.Error("Prediction timed out after 60 seconds.")
+            time.sleep(poll_interval)
+
     if response.status_code == 200:
-        
         json_response = response.json()
-        #If the output component is JSON return the entire output response 
-        if(outputs[0].get_config()["name"] == "json"):
+        if outputs[0].get_config()["name"] == "json":
             return json_response["output"]
+
         predict_outputs = parse_outputs(json_response["output"])
         processed_outputs = process_outputs(predict_outputs)
         print(f"processed_outputs: {processed_outputs}")
         return tuple(processed_outputs) if len(processed_outputs) > 1 else processed_outputs[0]
     else:
-        time.sleep(1)
-        if(response.status_code == 409):
-            raise gr.Error(f"Sorry, the Cog image is still processing. Try again in a bit.")
+        if response.status_code == 409:
+            raise gr.Error("Sorry, the Cog image is still processing. Try again in a bit.")
         raise gr.Error(f"The submission failed! Error: {response.status_code}")
 
 
+
 css = '''
 #col-container{max-width: 800px;margin: 0 auto;}
 '''