@@ -147,19 +147,20 @@ class Replicate {
147147 const { wait, signal, ...data } = options ;
148148
149149 const identifier = ModelVersionIdentifier . parse ( ref ) ;
150+ const isBlocking = typeof wait === "boolean" || typeof wait === "number" ;
150151
151152 let prediction ;
152153 if ( identifier . version ) {
153154 prediction = await this . predictions . create ( {
154155 ...data ,
155156 version : identifier . version ,
156- wait : wait ,
157+ wait : isBlocking ? wait : false ,
157158 } ) ;
158159 } else if ( identifier . owner && identifier . name ) {
159160 prediction = await this . predictions . create ( {
160161 ...data ,
161162 model : `${ identifier . owner } /${ identifier . name } ` ,
162- wait : wait ,
163+ wait : isBlocking ? wait : false ,
163164 } ) ;
164165 } else {
165166 throw new Error ( "Invalid model version identifier" ) ;
@@ -170,23 +171,26 @@ class Replicate {
170171 progress ( prediction ) ;
171172 }
172173
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
180191 }
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+ }
190194
191195 if ( signal && signal . aborted ) {
192196 prediction = await this . predictions . cancel ( prediction . id ) ;
0 commit comments