add support for manually setting a model to staging status

Author: zeke

src/index.ts

@@ -21,10 +21,10 @@ const getModelStatus = async (model: InferenceModel) => {
 // Get Replicate model warm/cold statuses in parallel
 const statuses = await Promise.all(inferenceModels.map(getModelStatus));
 
-// Add status to each model object
+// Set status (unless it's already manually defined on the model object)
 const replicateModels = inferenceModels.map((model, index) => ({
     ...model,
-    status: statuses[index],
+    status: model.status ?? statuses[index],
 }));
 
 console.log("\n\nReplicate model mappings:");

src/models.ts

@@ -2,6 +2,9 @@ export interface InferenceModel {
     hfModel: string;
     providerModel: string;
     task: string;
+
+    // You can set this to force the value, e.g. to keep a model as 'staging' even if it's
+    // warm/live on Replicate. If not set, the status will be inferred from the provider model
     status?: 'live' | 'staging';
 }