From d034840b9f10dc7ce24255301a8f970f2f26be56 Mon Sep 17 00:00:00 2001
From: Zeke Sikelianos <zeke@sikelianos.com>
Date: Mon, 25 Sep 2023 14:34:18 -0700
Subject: [PATCH] default to public model and make deployments opt-in

---
 .env.example                   |  7 ++++++-
 pages/api/predictions/index.js | 27 +++++++++++++++++++--------
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/.env.example b/.env.example
index f259d4a..0cde9f3 100644
--- a/.env.example
+++ b/.env.example
@@ -9,4 +9,9 @@ DATABASE_URL=
 # so the Replicate API can send it webhooks
 # 
 # e.g. https://8db01fea81ad.ngrok.io 
-NGROK_HOST=
\ No newline at end of file
+NGROK_HOST=
+
+
+# Optional: Set a value for this to run the private Replicate deployment of the model
+# instead of the public ControlNet Scribble model.
+USE_REPLICATE_DEPLOYMENT=
\ No newline at end of file
diff --git a/pages/api/predictions/index.js b/pages/api/predictions/index.js
index 229d36e..010105d 100644
--- a/pages/api/predictions/index.js
+++ b/pages/api/predictions/index.js
@@ -27,21 +27,32 @@ export default async function handler(req) {
 
   const input = await getObjectFromRequestBodyStream(req.body);
 
-  // https://replicate.com/rossjillian/controlnet
-
   const replicate = new Replicate({
     auth: process.env.REPLICATE_API_TOKEN,
   });
 
-  const prediction = await replicate.deployments.predictions.create(
-    "replicate",
-    "scribble-diffusion-jagilley-controlnet",
-    {
+  let prediction;
+
+  if (process.env.USE_REPLICATE_DEPLOYMENT) {
+    prediction = await replicate.deployments.predictions.create(
+      "replicate",
+      "scribble-diffusion-jagilley-controlnet",
+      {
+        input,
+        webhook: `${WEBHOOK_HOST}/api/replicate-webhook`,
+        webhook_events_filter: ["start", "completed"],
+      }
+    );
+  } else {
+    // https://replicate.com/jagilley/controlnet-scribble/versions
+    prediction = await replicate.predictions.create({
+      version:
+        "435061a1b5a4c1e26740464bf786efdfa9cb3a3ac488595a2de23e143fdb0117",
       input,
       webhook: `${WEBHOOK_HOST}/api/replicate-webhook`,
       webhook_events_filter: ["start", "completed"],
-    }
-  );
+    });
+  }
 
   if (prediction?.error) {
     return NextResponse.json({ detail: prediction.error }, { status: 500 });