@@ -147,19 +147,20 @@ class Replicate {
147
147
const { wait, signal, ...data } = options ;
148
148
149
149
const identifier = ModelVersionIdentifier . parse ( ref ) ;
150
+ const isBlocking = typeof wait === "boolean" || typeof wait === "number" ;
150
151
151
152
let prediction ;
152
153
if ( identifier . version ) {
153
154
prediction = await this . predictions . create ( {
154
155
...data ,
155
156
version : identifier . version ,
156
- wait : wait ,
157
+ wait : isBlocking ? wait : false ,
157
158
} ) ;
158
159
} else if ( identifier . owner && identifier . name ) {
159
160
prediction = await this . predictions . create ( {
160
161
...data ,
161
162
model : `${ identifier . owner } /${ identifier . name } ` ,
162
- wait : wait ,
163
+ wait : isBlocking ? wait : false ,
163
164
} ) ;
164
165
} else {
165
166
throw new Error ( "Invalid model version identifier" ) ;
@@ -170,23 +171,26 @@ class Replicate {
170
171
progress ( prediction ) ;
171
172
}
172
173
173
- prediction = await this . wait (
174
- prediction ,
175
- wait || { } ,
176
- async ( updatedPrediction ) => {
177
- // Call progress callback with the updated prediction object
178
- if ( progress ) {
179
- progress ( updatedPrediction ) ;
174
+ const isDone = isBlocking && prediction . status !== "starting" ;
175
+ if ( ! isDone ) {
176
+ prediction = await this . wait (
177
+ prediction ,
178
+ isBlocking ? { } : wait ,
179
+ async ( updatedPrediction ) => {
180
+ // Call progress callback with the updated prediction object
181
+ if ( progress ) {
182
+ progress ( updatedPrediction ) ;
183
+ }
184
+
185
+ // We handle the cancel later in the function.
186
+ if ( signal && signal . aborted ) {
187
+ return true ; // stop polling
188
+ }
189
+
190
+ return false ; // continue polling
180
191
}
181
-
182
- // We handle the cancel later in the function.
183
- if ( signal && signal . aborted ) {
184
- return true ; // stop polling
185
- }
186
-
187
- return false ; // continue polling
188
- }
189
- ) ;
192
+ ) ;
193
+ }
190
194
191
195
if ( signal && signal . aborted ) {
192
196
prediction = await this . predictions . cancel ( prediction . id ) ;
0 commit comments