diff --git a/lib/deployments.js b/lib/deployments.js
index 6cab261..716c8e1 100644
--- a/lib/deployments.js
+++ b/lib/deployments.js
@@ -9,11 +9,11 @@ const { transformFileInputs } = require("./util");
  * @param {object} options.input - Required. An object with the model inputs
  * @param {string} [options.webhook] - An HTTPS URL for receiving a webhook when the prediction has new output
  * @param {string[]} [options.webhook_events_filter] - You can change which events trigger webhook requests by specifying webhook events (`start`|`output`|`logs`|`completed`)
- * @param {boolean} [options.block] - Whether to wait until the prediction is completed before returning. Defaults to false
+ * @param {boolean|integer} [options.wait] - Whether to wait until the prediction is completed before returning. If an integer is provided, it will wait for that many seconds. Defaults to false
  * @returns {Promise<object>} Resolves with the created prediction data
  */
 async function createPrediction(deployment_owner, deployment_name, options) {
-  const { input, block, ...data } = options;
+  const { input, wait, ...data } = options;
 
   if (data.webhook) {
     try {
@@ -25,8 +25,13 @@ async function createPrediction(deployment_owner, deployment_name, options) {
   }
 
   const headers = {};
-  if (block) {
-    headers["Prefer"] = "wait";
+  if (wait) {
+    if (typeof wait === "number") {
+      const n = Math.max(1, Math.ceil(Number(wait)) || 1);
+      headers["Prefer"] = `wait=${n}`;
+    } else {
+      headers["Prefer"] = "wait";
+    }
   }
 
   const response = await this.request(
diff --git a/package-lock.json b/package-lock.json
index 300e538..0579ebd 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "replicate",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "replicate",
-      "version": "1.0.0",
+      "version": "1.0.1",
       "license": "Apache-2.0",
       "devDependencies": {
         "@biomejs/biome": "^1.4.1",
diff --git a/package.json b/package.json
index fcd47e8..19e68fc 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "replicate",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "JavaScript client for Replicate",
   "repository": "github:replicate/replicate-javascript",
   "homepage": "https://github.com/replicate/replicate-javascript#readme",