Allow user to configure poll interval

Author: andreasjansson

replicate/client.py

@@ -19,6 +19,7 @@ def __init__(self, api_token=None) -> None:
         self.base_url = os.environ.get(
             "REPLICATE_API_BASE_URL", "https://api.replicate.com"
         )
+        self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
 
         # TODO: make thread safe
         self.session = requests.Session()

replicate/prediction.py

@@ -24,7 +24,7 @@ class Prediction(BaseModel):
     def wait(self):
         """Wait for prediction to finish."""
         while self.status not in ["succeeded", "failed", "canceled"]:
-            time.sleep(0.5)
+            time.sleep(self._client.poll_interval)
             self.reload()
 
     def output_iterator(self) -> Iterator[Any]:
@@ -36,7 +36,7 @@ def output_iterator(self) -> Iterator[Any]:
             for output in new_output:
                 yield output
             previous_output = output
-            time.sleep(0.5)
+            time.sleep(self._client.poll_interval)
             self.reload()
 
         if self.status == "failed":