Merge pull request #64 from replicate/reduce-prediction-load

Rely on app database rather than replicate for prediction polling

.env.example

@@ -6,10 +6,12 @@ OPENAI_API_KEY=
 
 # setup a supabase account and create a new project https://supabase.com/docs/guides/getting-started/quickstarts/nextjs
 # this is only required if you want your outputs to persist, or if you're using controlnet (input image uploads need to be saved to a DB)
-SUPABASE_SERVICE_ROLE=
 SUPABASE_URL=
-SUPABASE_JWT=
 SUPABASE_KEY=
 
+# host for your supabase "images" bucket used to store and serve predicitions
+# e.g. xyz.supabase.co
+NEXT_PUBLIC_SUPABASE_IMAGES_HOST=
+
 # install https://ngrok.com/, run it, and set the host here
 NGROK_HOST=

components/controlnet-prediction.js

@@ -19,7 +19,7 @@ export default function ControlnetPrediction({ prediction }) {
   }
 
   async function getOutput(prediction) {
-    const bucket = `https://ennwjiitmiqwdrgxkevm.supabase.co/storage/v1/object/public/images/public`;
+    const bucket = `${process.env.NEXT_PUBLIC_SUPABASE_IMAGES_URL}/storage/v1/object/public/images/public`;
     const predictionUrl = `${bucket}/${prediction.id}.png`;
     const predictionAnnotationUrl = `${bucket}/annotated-${prediction.id}.png`;
 

components/popup.js

@@ -49,7 +49,7 @@ export default function Popup({ open, setOpen }) {
                       width={250}
                       height={250}
                       className="mx-auto"
-                      src="https://ennwjiitmiqwdrgxkevm.supabase.co/storage/v1/object/public/images/public/2370caef-bc8c-423e-88bb-0ca4a71e19c8.png"
+                      src={`${process.env.NEXT_PUBLIC_SUPABASE_IMAGES_URL}/storage/v1/object/public/images/public/2370caef-bc8c-423e-88bb-0ca4a71e19c8.png`}
                       alt=""
                       unoptimized={true}
                     />

components/prediction.js

@@ -13,7 +13,7 @@ export default function Prediction({ prediction }) {
     if (prediction.output) {
       return prediction.output;
     } else {
-      return `https://ennwjiitmiqwdrgxkevm.supabase.co/storage/v1/object/public/images/public/${prediction.id}.png`;
+      return `${process.env.NEXT_PUBLIC_SUPABASE_IMAGES_URL}/storage/v1/object/public/images/public/${prediction.id}.png`;
     }
   }
 

lib/db.js

@@ -8,4 +8,4 @@ if (process.env.SUPABASE_URL) {
   db = createClient(supabaseUrl, supabaseKey);
 }
 
-module.exports = db;
+export default db;

lib/upsertPrediction.js

@@ -58,10 +58,10 @@ export default async function upsertPrediction(prediction) {
 
   const { data, error } = await db
     .from("predictions")
-    .upsert([predictionObject], { onConflict: "id" });
+    .upsert(predictionObject, { onConflict: "id" });
 
   if (error) {
-    console.log("error upserting prediction ", error);
+    console.log("error upserting prediction ", predictionObject, error);
   }
 
   console.log(`prediction ${JSON.stringify(predictionObject.id)} upserted`);

pages/api/og.js

@@ -10,7 +10,7 @@ export default async function handler(req) {
   try {
     const { searchParams } = req.nextUrl;
     const bucketUrl =
-      "https://ennwjiitmiqwdrgxkevm.supabase.co/storage/v1/object/public/images/public/";
+      `${process.env.NEXT_PUBLIC_SUPABASE_IMAGES_URL}/storage/v1/object/public/images/public/`;
     let predictionIds = searchParams.get("ids").split(",");
     let outputUrls = predictionIds.map((id) => `${bucketUrl}${id}.png`);
     const prompt = searchParams.get("prompt");

pages/api/predictions/[id].js

@@ -1,14 +1,24 @@
-import Replicate from "replicate";
 import db from "../../../lib/db";
 
-const replicate = new Replicate({
-  auth: process.env.REPLICATE_API_TOKEN,
-});
-
 export default async function handler(req, res) {
   if (req.method === "GET") {
-    const prediction = await replicate.predictions.get(req.query.id);
-    res.end(JSON.stringify(prediction));
+    const {data, error} = await db.from("predictions").select().eq('id', req.query.id)
+
+    if (error) {
+      console.log("error retrieving prediction ", error, req.body);
+      return res.status(500).json({ error: error.message });
+    }
+
+    // If the prediction is not found, we are still waiting on the webhook
+    // so we just return success with an unknown status and id.
+    // TODO: We could store the initial prediction in the db on creation
+    // time but then we'd potentially be storing failed and partial
+    // predictions.
+    if (data.length === 0) {
+      return res.status(200).json({ id: req.query.id, status: "unknown"  });
+    }
+
+    res.end(JSON.stringify(data[0]));
   } else if (req.method === "PUT") {
     const { data, error } = await db
       .from("predictions")

pages/api/predictions/index.js

@@ -26,7 +26,7 @@ export default async function handler(req, res) {
 
   if (!WEBHOOK_HOST) {
     throw new Error(
-      "WEBHOOK HOST is not set. If you're on local, make sure you set it to an ngrok url. If this doesn't exist, replicate predictions won't save to DB."
+      "WEBHOOK_HOST is not set. If you're on local, make sure you set NGROK_HOST. If this doesn't exist, replicate predictions won't save to DB."
     );
   }
 

pages/api/submissions/[id].js

@@ -12,14 +12,14 @@ export default async function handler(req, res) {
     .eq("submission_id", req.query.id)
     .order("created_at", { ascending: true });
 
-  data.forEach((prediction) => {
-    delete prediction.output;
-  });
-
   if (error) {
     console.log("error getting predictions for submission id ", error);
     return res.status(500).json({ error: error.message });
   }
 
+  data.forEach((prediction) => {
+    delete prediction.output;
+  });
+
   res.end(JSON.stringify(data));
 }

pages/controlnet.js

@@ -22,7 +22,7 @@ const supabase = createClient(
   "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6ImVubndqaWl0bWlxd2RyZ3hrZXZtIiwicm9sZSI6ImFub24iLCJpYXQiOjE2ODM5Mjc3OTgsImV4cCI6MTk5OTUwMzc5OH0.zCHzwchIjcmKNmccb9D4OLVwrWrpLHMmf4a8W7UedFs"
 );
 
-const supabaseUrl = "https://ennwjiitmiqwdrgxkevm.supabase.co";
+const supabaseUrl = process.env.NEXT_PUBLIC_SUPABASE_IMAGES_HOST;
 
 const sleep = (ms) => new Promise((r) => setTimeout(r, ms));
 const fileTypes = ["JPG", "PNG"];

pages/index.js

@@ -544,7 +544,9 @@ export async function getServerSideProps({ req }) {
     const response = await fetch(`${baseUrl}/api/submissions/${submissionId}`, {
       method: "GET",
     });
-    submissionPredictions = await response.json();
+    if (response.ok) {
+      submissionPredictions = await response.json();
+    }
   }
 
   return { props: { baseUrl, submissionPredictions } };

pages/memories.js

@@ -7,7 +7,7 @@ export default function History() {
   const [anonId, setAnonId] = useState(null);
 
   function getPredictionOutput(prediction) {
-    return `https://ennwjiitmiqwdrgxkevm.supabase.co/storage/v1/object/public/images/public/${prediction.id}.png`;
+    return `${process.env.NEXT_PUBLIC_SUPABASE_IMAGES_URL}/storage/v1/object/public/images/public/${prediction.id}.png`;
   }
 
   const clearHistory = () => {