add time to wait before taking a screenshot

predict.py

@@ -3,6 +3,7 @@
 from selenium.webdriver.common.keys import Keys
 
 from cog import BasePredictor, Input, Path
+import time
 
 
 class Predictor(BasePredictor):
@@ -27,9 +28,17 @@ def predict(
             description="Height of the screenshot",
             default=1080
         ),
+        wait_until: int = Input(
+            description="Time in seconds to wait before a screenshot is taken",
+            default=0
+        )
     ) -> Path:
         self.browser.set_window_size(w, h)
         self.browser.get(url)
+        if wait_until > 0:
+            for i in range(wait_until):
+                print(f"Elapsed time: {i+1}/{wait_until} seconds")
+                time.sleep(1)
 
         # debug details
         print(f"Page title: {self.browser.title}")