Merge pull request #4 from replicate/add-support-for-manually-setting-a-model-to-staging-status

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';
 }
 
@@ -10,6 +13,7 @@ export const inferenceModels: InferenceModel[] = [
         hfModel: "deepseek-ai/DeepSeek-R1",
         providerModel: "deepseek-ai/deepseek-r1",
         task: "conversational",
+        status: "staging",
     },
     {
         hfModel: "black-forest-labs/FLUX.1-dev",