Merge pull request #60 from replicate/add-latent-consistency

Add latent consistency model
replicate · Oct 19, 2023 · ed258c3 · ed258c3 · vercel · Oct 19, 2023
2 parents cbbd0fc + c7b7e4e
commit ed258c3
Show file tree

Hide file tree

Showing 4 changed files with 108 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -26,6 +26,12 @@ npm run dev
 
 Open [http://localhost:3000](http://localhost:3000) with your browser.
 
+## Want to add a model?
+
+1. Check out `lib/models.js` and add your model to the MODELS array.
+2. Optionally, generate some example predictions by adding some submissions to `lib/seeds.js`.
+3. Push your PR!
+
 ## Deploying on Vercel
 
 Alternatively, you can [deploy Zoo on Vercel](./doc/deploy_vercel/README.md).
diff --git a/components/save-image.js b/components/save-image.js
@@ -3,9 +3,19 @@ import { Dialog, Transition } from "@headlessui/react";
 import { XMarkIcon } from "@heroicons/react/20/solid";
 import FileSaver from "file-saver";
 
-export function SaveImage({ open, setOpen, prediction, url, annotatedUrl = false }) {
+export function SaveImage({
+  open,
+  setOpen,
+  prediction,
+  url,
+  annotatedUrl = false,
+}) {
   const download = async (url, id) => {
-    FileSaver.saveAs(url, `${id}.png`);
+    if (typeof url !== "list") {
+      FileSaver.saveAs(url[0], `${id}.png`);
+    } else {
+      FileSaver.saveAs(url, `${id}.png`);
+    }
   };
 
   return (
@@ -38,7 +48,11 @@ export function SaveImage({ open, setOpen, prediction, url, annotatedUrl = false
             leaveFrom="opacity-100 scale-100"
             leaveTo="opacity-0 scale-95"
           >
-            <Dialog.Panel className={`relative transform overflow-hidden rounded-lg bg-white px-4 pt-5 pb-4 text-left shadow-xl transition-all sm:my-8 sm:w-full ${annotatedUrl ? 'sm:max-w-3xl' : 'sm:max-w-sm'} mx-auto sm:p-6`}>
+            <Dialog.Panel
+              className={`relative transform overflow-hidden rounded-lg bg-white px-4 pt-5 pb-4 text-left shadow-xl transition-all sm:my-8 sm:w-full ${
+                annotatedUrl ? "sm:max-w-3xl" : "sm:max-w-sm"
+              } mx-auto sm:p-6`}
+            >
               <div className="absolute top-0 right-0 hidden pt-4 pr-4 sm:block">
                 <button
                   type="button"
@@ -58,7 +72,6 @@ export function SaveImage({ open, setOpen, prediction, url, annotatedUrl = false
                     Zoo Prediction
                   </Dialog.Title>
                   <div className="mt-4">
-
                     {annotatedUrl ? (
                       <div className="grid grid-cols-3 mt-4 gap-x-4">
                         <div>
@@ -103,7 +116,6 @@ export function SaveImage({ open, setOpen, prediction, url, annotatedUrl = false
                         loading="lazy"
                       />
                     )}
-
                   </div>
                   <div>
                     <p className="mt-4 block truncate text-sm font-medium text-gray-900">
@@ -113,11 +125,16 @@ export function SaveImage({ open, setOpen, prediction, url, annotatedUrl = false
                       {prediction.input.prompt}
                     </p>
 
-                    {Object.keys(prediction.input).map(key => {
+                    {Object.keys(prediction.input).map((key) => {
                       const value = prediction.input[key];
-                      return <p key={key} className="block text-sm font-medium text-gray-500">
-                        {key}: {value}
-                      </p>
+                      return (
+                        <p
+                          key={key}
+                          className="block text-sm font-medium text-gray-500"
+                        >
+                          {key}: {value}
+                        </p>
+                      );
                     })}
 
                     {prediction.metrics && (

diff --git a/lib/models.js b/lib/models.js
@@ -1,4 +1,38 @@
-// the ordering of these matters, it determines the order in the UI
+/**
+ * These are the models that are runnable on Zoo.
+ * To add your own, add a model object to this array.
+ * The ordering of these model objects matter — it determines the order in the UI!
+ *
+ * Each model object looks like this:
+ *
+ * {
+    id: 101 (integer ID, must be unique otherwise you'll get weirdness),
+    owner: "stability-ai", // model owner on Replicate
+    name: "SDXL", // model name on Replicate
+    default_params: { // default parameters used for running the model on Replicate
+      width: 1024,
+      height: 1024,
+      scheduler: "K_EULER",
+    },
+    version: "98d6bab2dd21e4ffc4cc626420ab4f24b99ec60728c5d835ff9c3439396aca45", // Replicate version
+    checked: true, // whether or no the model is checked by default in the UI
+    source: "replicate", // source of the model, either "replicate" or "openai"
+    url: "https://replicate.com/stability-ai/sdxl?utm_source=project&utm_campaign=zoo", // url to the model
+    description:
+      "A text-to-image generative AI model that creates beautiful 1024x1024 images",
+    links: [
+      {
+        name: "replicate",
+        url: "https://replicate.com/stability-ai/sdxl?utm_source=project&utm_campaign=zoo",
+      },
+      {
+        name: "github",
+        url: "https://github.com/Stability-AI/generative-models",
+      },
+    ],
+  },
+ *
+ */
 const MODELS = [
   {
     id: 31,
@@ -26,6 +60,31 @@ const MODELS = [
       },
     ],
   },
+  {
+    id: 33,
+    owner: "luosiallen",
+    name: "latent-consistency-model",
+    default_params: {
+      width: 512,
+      height: 512,
+      num_inference_steps: 4,
+    },
+    checked: true,
+    version: "553803fd018b3cf875a8bc774c99da9b33f36647badfd88a6eec90d61c5f62fc",
+    source: "replicate",
+    url: "https://replicate.com/luosiallen/latent-consistency-model?utm_source=project&utm_campaign=zoo",
+    description: "Synthesizing High-Resolution Images with Few-Step Inference",
+    links: [
+      {
+        name: "replicate",
+        url: "https://replicate.com/luosiallen/latent-consistency-model?utm_source=project&utm_campaign=zoo",
+      },
+      {
+        name: "github",
+        url: "https://github.com/luosiallen/latent-consistency-model",
+      },
+    ],
+  },
   {
     id: 1,
     owner: "stability-ai",

diff --git a/lib/seeds.js b/lib/seeds.js
@@ -1,3 +1,17 @@
+/**
+ * These seeds are ids pointing to a collection of predictions run on Zoo.
+ *
+ * To add a new seed, you need to generate predictions on Zoo and then
+ * save the id url parameter of the submission.
+ *
+ * Here's how:
+ * 1. Run Zoo locally with `npm run dev`
+ * 2. Run ngrok with `ngrok http 3000`. Set your NGROK_HOST environment variable to the url ngrok gives you.
+ * This is required because it serves as the endpoint for the webhook that saves predictions to the database.
+ * 3. Create a submission (enter a prompt and press go)
+ * 4. Save the id URL param and add it to the list here!
+ */
+
 const seeds = [
   "a-detailed-painting-of-fish-american-barbizon-school-by-diego-rivera-pxsz48",
   "a-digital-painting-of-ocean-waves-rayonism-by-alejandro-obregon-51pr569",
@@ -19,6 +33,8 @@ const seeds = [
   "a9a7cb97-04d5-4ec3-a3ea-21d178092acb",
   "08fe73b4-2ab4-46b1-8a91-a4392b54997e",
   "a1920791-41e2-45f4-8cca-bcfc01440fd4",
+  "831bbaa3-3071-43ad-aaa4-4765ffb95cd1",
+  "53d5acf3-7092-46a9-8219-f788b986c196",
 ];
 
 module.exports = seeds;