From cee8caa3e8bc7c81e9cd8d03373ea1f72e0bd581 Mon Sep 17 00:00:00 2001
From: Austin Chen <akrolsmir@gmail.com>
Date: Wed, 12 Oct 2022 09:53:04 -0700
Subject: [PATCH] Generate images from StableDiffusion (#1035)

* Generate images from StableDiffusion

* Update yarn.lock

* Log an error, remove extra comment

* Code cleanup

* Note about the API
---
 web/.gitignore                        |   3 +-
 web/components/editor.tsx             |  16 ++-
 web/components/editor/image-modal.tsx | 157 ++++++++++++++++++++++++++
 web/package.json                      |   1 +
 web/pages/api/v0/dream.ts             |  90 +++++++++++++++
 yarn.lock                             |  72 ++++++++++++
 6 files changed, 334 insertions(+), 5 deletions(-)
 create mode 100644 web/components/editor/image-modal.tsx
 create mode 100644 web/pages/api/v0/dream.ts

diff --git a/web/.gitignore b/web/.gitignore
index 32d4a19e..06749214 100644
--- a/web/.gitignore
+++ b/web/.gitignore
@@ -2,4 +2,5 @@
 .next
 node_modules
 out
-tsconfig.tsbuildinfo
\ No newline at end of file
+tsconfig.tsbuildinfo
+.env*
diff --git a/web/components/editor.tsx b/web/components/editor.tsx
index 4b42c6be..8de39b70 100644
--- a/web/components/editor.tsx
+++ b/web/components/editor.tsx
@@ -18,7 +18,6 @@ import { useCallback, useEffect, useState } from 'react'
 import { Linkify } from './linkify'
 import { uploadImage } from 'web/lib/firebase/storage'
 import { useMutation } from 'react-query'
-import { FileUploadButton } from './file-upload-button'
 import { linkClass } from './site-link'
 import { DisplayMention } from './editor/mention'
 import { DisplayContractMention } from './editor/contract-mention'
@@ -43,6 +42,7 @@ import ItalicIcon from 'web/lib/icons/italic-icon'
 import LinkIcon from 'web/lib/icons/link-icon'
 import { getUrl } from 'common/util/parse'
 import { TiptapSpoiler } from 'common/util/tiptap-spoiler'
+import { ImageModal } from './editor/image-modal'
 import {
   storageStore,
   usePersistentState,
@@ -255,6 +255,7 @@ export function TextEditor(props: {
   children?: React.ReactNode // additional toolbar buttons
 }) {
   const { editor, upload, children } = props
+  const [imageOpen, setImageOpen] = useState(false)
   const [iframeOpen, setIframeOpen] = useState(false)
   const [marketOpen, setMarketOpen] = useState(false)
 
@@ -268,12 +269,19 @@ export function TextEditor(props: {
           {/* Toolbar, with buttons for images and embeds */}
           <div className="flex h-9 items-center gap-5 pl-4 pr-1">
             <Tooltip text="Add image" noTap noFade>
-              <FileUploadButton
-                onFiles={upload.mutate}
+              <button
+                type="button"
+                onClick={() => setImageOpen(true)}
                 className="-m-2.5 flex h-10 w-10 items-center justify-center rounded-full text-gray-400 hover:text-gray-500"
               >
+                <ImageModal
+                  editor={editor}
+                  upload={upload}
+                  open={imageOpen}
+                  setOpen={setImageOpen}
+                />
                 <PhotographIcon className="h-5 w-5" aria-hidden="true" />
-              </FileUploadButton>
+              </button>
             </Tooltip>
             <Tooltip text="Add embed" noTap noFade>
               <button
diff --git a/web/components/editor/image-modal.tsx b/web/components/editor/image-modal.tsx
new file mode 100644
index 00000000..6728513f
--- /dev/null
+++ b/web/components/editor/image-modal.tsx
@@ -0,0 +1,157 @@
+import { UploadIcon } from '@heroicons/react/outline'
+import { Editor } from '@tiptap/react'
+import { useState } from 'react'
+import { AlertBox } from '../alert-box'
+import { Button } from '../button'
+import { FileUploadButton } from '../file-upload-button'
+import { Col } from '../layout/col'
+import { Modal } from '../layout/modal'
+import { Row } from '../layout/row'
+import { Tabs } from '../layout/tabs'
+
+const MODIFIERS =
+  '8k, beautiful, illustration, trending on art station, picture of the day, epic composition'
+
+export function ImageModal(props: {
+  editor: Editor | null
+  // TODO: Type this correctly?
+  upload: any
+  open: boolean
+  setOpen: (open: boolean) => void
+}) {
+  const { upload, open, setOpen } = props
+  return (
+    <Modal open={open} setOpen={setOpen}>
+      <Col className="gap-2 rounded bg-white p-6">
+        <Tabs
+          tabs={[
+            {
+              title: 'Upload file',
+              content: (
+                <FileUploadButton
+                  onFiles={(files) => {
+                    setOpen(false)
+                    upload.mutate(files)
+                  }}
+                  className="relative block w-full rounded-lg border-2 border-dashed border-gray-300 p-12 text-center hover:border-gray-400 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2"
+                >
+                  <UploadIcon className="mx-auto h-12 w-12 text-gray-400" />
+                  <span className="mt-2 block text-sm font-medium text-gray-400">
+                    Upload an image file
+                  </span>
+                </FileUploadButton>
+              ),
+            },
+            {
+              title: 'Dream',
+              content: <DreamTab {...props} />,
+            },
+          ]}
+        />
+      </Col>
+    </Modal>
+  )
+}
+
+// Note: this is currently tied to a DreamStudio API key tied to akrolsmir@gmail.com,
+// and injected on Vercel.
+const API_KEY = process.env.NEXT_PUBLIC_DREAM_KEY
+
+function DreamTab(props: {
+  editor: Editor | null
+  open: boolean
+  setOpen: (open: boolean) => void
+}) {
+  const { editor, setOpen } = props
+  const [input, setInput] = useState('')
+  const [isDreaming, setIsDreaming] = useState(false)
+  const [imageUrl, setImageUrl] = useState('')
+  const imageCode = `<img src="${imageUrl}" alt="${input}" />`
+
+  if (!API_KEY) {
+    return (
+      <AlertBox
+        title="Missing API Key"
+        text="An API key from https://beta.dreamstudio.ai/ is needed to dream; add it to your web/.env.local"
+      />
+    )
+  }
+
+  async function dream() {
+    setIsDreaming(true)
+    const data = {
+      prompt: input + ', ' + MODIFIERS,
+      apiKey: API_KEY,
+    }
+    const response = await fetch(`/api/v0/dream`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify(data),
+    })
+    const json = await response.json()
+    setImageUrl(json.url)
+    setIsDreaming(false)
+  }
+
+  return (
+    <Col className="gap-2">
+      <Row className="gap-2">
+        <input
+          autoFocus
+          type="text"
+          name="embed"
+          id="embed"
+          className="block w-full rounded-md border-gray-300 shadow-sm placeholder:text-gray-300 focus:border-indigo-500 focus:ring-indigo-500 sm:text-sm"
+          placeholder="A crane playing poker on a green table"
+          value={input}
+          onChange={(e) => setInput(e.target.value)}
+          autoComplete="off"
+        />
+        <Button
+          className="whitespace-nowrap"
+          onClick={dream}
+          loading={isDreaming}
+        >
+          Dream
+          {/* TODO: Charge M$5 with ({formatMoney(5)}) */}
+        </Button>
+      </Row>
+      {isDreaming && (
+        <div className="text-sm">This may take ~10 seconds...</div>
+      )}
+      {/* TODO: Allow the user to choose their own modifiers */}
+      <div className="pt-2 text-sm text-gray-400">Modifiers: {MODIFIERS}</div>
+
+      {/* Show the current imageUrl */}
+      {/* TODO: Keep the other generated images, so the user can play with different attempts. */}
+      {imageUrl && (
+        <>
+          <img src={imageUrl} alt="Image" />
+          <Row className="gap-2">
+            <Button
+              disabled={isDreaming}
+              onClick={() => {
+                if (editor) {
+                  editor.chain().insertContent(imageCode).run()
+                  setInput('')
+                  setOpen(false)
+                }
+              }}
+            >
+              Add image
+            </Button>
+            <Button
+              color="gray"
+              onClick={() => {
+                setInput('')
+                setOpen(false)
+              }}
+            >
+              Cancel
+            </Button>
+          </Row>
+        </>
+      )}
+    </Col>
+  )
+}
diff --git a/web/package.json b/web/package.json
index f64b79e8..5f39d174 100644
--- a/web/package.json
+++ b/web/package.json
@@ -63,6 +63,7 @@
     "react-masonry-css": "1.0.16",
     "react-query": "3.39.0",
     "react-twitter-embed": "4.0.4",
+    "stability-client": "1.5.0",
     "string-similarity": "^4.0.4",
     "tippy.js": "6.3.7"
   },
diff --git a/web/pages/api/v0/dream.ts b/web/pages/api/v0/dream.ts
new file mode 100644
index 00000000..df7905f3
--- /dev/null
+++ b/web/pages/api/v0/dream.ts
@@ -0,0 +1,90 @@
+import { getDownloadURL, ref, uploadBytesResumable } from 'firebase/storage'
+import { nanoid } from 'nanoid'
+import { NextApiRequest, NextApiResponse } from 'next'
+import { generateAsync } from 'stability-client'
+import { storage } from 'web/lib/firebase/init'
+import {
+  CORS_ORIGIN_MANIFOLD,
+  CORS_ORIGIN_LOCALHOST,
+} from 'common/envs/constants'
+import { applyCorsHeaders } from 'web/lib/api/cors'
+
+export const config = { api: { bodyParser: true } }
+
+// Highly experimental. Proxy for https://github.com/vpzomtrrfrt/stability-client
+export default async function route(req: NextApiRequest, res: NextApiResponse) {
+  await applyCorsHeaders(req, res, {
+    origin: [CORS_ORIGIN_MANIFOLD, CORS_ORIGIN_LOCALHOST],
+    methods: 'POST',
+  })
+
+  // Check that prompt and apiKey are included in the body
+  if (!req.body.prompt) {
+    res.status(400).json({ message: 'Missing prompt' })
+    return
+  }
+  if (!req.body.apiKey) {
+    res.status(400).json({ message: 'Missing apiKey' })
+    return
+  }
+  /** Optional params:
+  outDir: string
+  debug: boolean
+  requestId: string
+  samples: number
+  engine: string
+  host: string
+  seed: number
+  width: number
+  height: number
+  diffusion: keyof typeof diffusionMap
+  steps: number
+  cfgScale: number
+  noStore: boolean
+  imagePrompt: {mime: string; content: Buffer} | null
+  stepSchedule: {start?: number; end?: number}
+  */
+
+  try {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    const { _dreamResponse, images } = await generateAsync({
+      ...req.body,
+      // Don't actually write to disk, because we're going to upload it to Firestore
+      noStore: true,
+    })
+    const buffer: Buffer = images[0].buffer
+    const url = await upload(buffer)
+
+    res.status(200).json({ url })
+  } catch (e) {
+    console.error(e)
+    res.status(500).json({ message: `Error running code: ${e}` })
+  }
+}
+
+// Loosely copied from web/lib/firebase/storage.ts
+const ONE_YEAR_SECS = 60 * 60 * 24 * 365
+
+async function upload(buffer: Buffer) {
+  const filename = `${nanoid(10)}.png`
+  const storageRef = ref(storage, `dream/${filename}`)
+
+  function promisifiedUploadBytes(...args: any[]) {
+    return new Promise((resolve, reject) => {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      const task = uploadBytesResumable(...args)
+      task.on(
+        'state_changed',
+        null,
+        (e: Error) => reject(e),
+        () => getDownloadURL(task.snapshot.ref).then(resolve)
+      )
+    })
+  }
+  return promisifiedUploadBytes(storageRef, buffer, {
+    cacheControl: `public, max-age=${ONE_YEAR_SECS}`,
+    contentType: 'image/png',
+  })
+}
diff --git a/yarn.lock b/yarn.lock
index 4d57c590..64a5be0e 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2400,6 +2400,18 @@
   resolved "https://registry.yarnpkg.com/@humanwhocodes/object-schema/-/object-schema-1.2.1.tgz#b520529ec21d8e5945a1851dfd1c32e94e39ff45"
   integrity sha512-ZnQMnLV4e7hDlUvw8H+U8ASL02SS2Gn6+9Ac3wGGLIe7+je2AeAOxPY+izIPJDfFDb7eDjev0Us8MO1iFRN8hA==
 
+"@improbable-eng/grpc-web-node-http-transport@^0.15.0":
+  version "0.15.0"
+  resolved "https://registry.yarnpkg.com/@improbable-eng/grpc-web-node-http-transport/-/grpc-web-node-http-transport-0.15.0.tgz#5a064472ef43489cbd075a91fb831c2abeb09d68"
+  integrity sha512-HLgJfVolGGpjc9DWPhmMmXJx8YGzkek7jcCFO1YYkSOoO81MWRZentPOd/JiKiZuU08wtc4BG+WNuGzsQB5jZA==
+
+"@improbable-eng/grpc-web@^0.15.0":
+  version "0.15.0"
+  resolved "https://registry.yarnpkg.com/@improbable-eng/grpc-web/-/grpc-web-0.15.0.tgz#3e47e9fdd90381a74abd4b7d26e67422a2a04bef"
+  integrity sha512-ERft9/0/8CmYalqOVnJnpdDry28q+j+nAlFFARdjyxXDJ+Mhgv9+F600QC8BR9ygOfrXRlAk6CvST2j+JCpQPg==
+  dependencies:
+    browser-headers "^0.4.1"
+
 "@jridgewell/gen-mapping@^0.1.0":
   version "0.1.1"
   resolved "https://registry.yarnpkg.com/@jridgewell/gen-mapping/-/gen-mapping-0.1.1.tgz#e5d2e450306a9491e3bd77e323e38d7aff315996"
@@ -4546,6 +4558,11 @@ broadcast-channel@^3.4.1:
     rimraf "3.0.2"
     unload "2.2.0"
 
+browser-headers@^0.4.1:
+  version "0.4.1"
+  resolved "https://registry.yarnpkg.com/browser-headers/-/browser-headers-0.4.1.tgz#4308a7ad3b240f4203dbb45acedb38dc2d65dd02"
+  integrity sha512-CA9hsySZVo9371qEHjHZtYxV2cFtVj5Wj/ZHi8ooEsrtm4vOnl9Y9HmyYWk9q+05d7K3rdoAE0j3MVEFVvtQtg==
+
 browser-image-compression@2.0.0:
   version "2.0.0"
   resolved "https://registry.yarnpkg.com/browser-image-compression/-/browser-image-compression-2.0.0.tgz#f421381a76d474d4da7dcd82810daf595b09bef6"
@@ -4943,6 +4960,11 @@ commander@^8.0.0, commander@^8.3.0:
   resolved "https://registry.yarnpkg.com/commander/-/commander-8.3.0.tgz#4837ea1b2da67b9c616a67afbb0fafee567bca66"
   integrity sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==
 
+commander@^9.4.0:
+  version "9.4.1"
+  resolved "https://registry.yarnpkg.com/commander/-/commander-9.4.1.tgz#d1dd8f2ce6faf93147295c0df13c7c21141cfbdd"
+  integrity sha512-5EEkTNyHNGFPD2H+c/dXXfQZYa/scCKasxWcXJaWnNJ99pnQN9Vnmqow+p+PlFPE63Q6mThaZws1T+HxfpgtPw==
+
 commondir@^1.0.1:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/commondir/-/commondir-1.0.1.tgz#ddd800da0c66127393cca5950ea968a3aaf1253b"
@@ -5821,6 +5843,11 @@ dot-prop@^5.2.0:
   dependencies:
     is-obj "^2.0.0"
 
+dotenv@^16.0.2:
+  version "16.0.3"
+  resolved "https://registry.yarnpkg.com/dotenv/-/dotenv-16.0.3.tgz#115aec42bac5053db3c456db30cc243a5a836a07"
+  integrity sha512-7GO6HghkA5fYG9TYnNxi14/7K9f5occMlp3zXAuSxn7CKCxt9xbNWG7yF8hTCSUchlfWSe3uLmlPfigevRItzQ==
+
 duplexer3@^0.1.4:
   version "0.1.4"
   resolved "https://registry.yarnpkg.com/duplexer3/-/duplexer3-0.1.4.tgz#ee01dd1cac0ed3cbc7fdbea37dc0a8f1ce002ce2"
@@ -7051,6 +7078,11 @@ google-p12-pem@^3.1.3:
   dependencies:
     node-forge "^1.3.1"
 
+google-protobuf@^3.21.0:
+  version "3.21.2"
+  resolved "https://registry.yarnpkg.com/google-protobuf/-/google-protobuf-3.21.2.tgz#4580a2bea8bbb291ee579d1fefb14d6fa3070ea4"
+  integrity sha512-3MSOYFO5U9mPGikIYCzK0SaThypfGgS6bHqrUGXG3DPHCrb+txNqeEcns1W0lkGfk0rCyNXm7xB9rMxnCiZOoA==
+
 got@^9.6.0:
   version "9.6.0"
   resolved "https://registry.yarnpkg.com/got/-/got-9.6.0.tgz#edf45e7d67f99545705de1f7bbeeeb121765ed85"
@@ -8748,6 +8780,11 @@ mkdirp@0.3.0:
   resolved "https://registry.yarnpkg.com/mkdirp/-/mkdirp-0.3.0.tgz#1bbf5ab1ba827af23575143490426455f481fe1e"
   integrity sha1-G79asbqCevI1dRQ0kEJkVfSB/h4=
 
+mkdirp@^1.0.4:
+  version "1.0.4"
+  resolved "https://registry.yarnpkg.com/mkdirp/-/mkdirp-1.0.4.tgz#3eb5ed62622756d79a5f0e2a221dfebad75c2f7e"
+  integrity sha512-vVqVZQyf3WLx2Shd0qJ9xuvqgAyKPLAiqITEtqW0oIUjzo3PePDd6fW9iFz30ef7Ysp/oiWqbhszeGWW2T6Gzw==
+
 module-alias@2.2.2:
   version "2.2.2"
   resolved "https://registry.yarnpkg.com/module-alias/-/module-alias-2.2.2.tgz#151cdcecc24e25739ff0aa6e51e1c5716974c0e0"
@@ -10822,6 +10859,13 @@ rxjs@^6.6.3:
   dependencies:
     tslib "^1.9.0"
 
+rxjs@^7.5.2:
+  version "7.5.7"
+  resolved "https://registry.yarnpkg.com/rxjs/-/rxjs-7.5.7.tgz#2ec0d57fdc89ece220d2e702730ae8f1e49def39"
+  integrity sha512-z9MzKh/UcOqB3i20H6rtrlaE/CgjLOvheWK/9ILrbhROGTweAi1BaFsTT9FbwZi5Trr1qNRs+MXkhmR06awzQA==
+  dependencies:
+    tslib "^2.1.0"
+
 rxjs@^7.5.4:
   version "7.5.5"
   resolved "https://registry.yarnpkg.com/rxjs/-/rxjs-7.5.5.tgz#2ebad89af0f560f460ad5cc4213219e1f7dd4e9f"
@@ -11272,6 +11316,22 @@ sprintf-js@~1.0.2:
   resolved "https://registry.yarnpkg.com/sprintf-js/-/sprintf-js-1.0.3.tgz#04e6926f662895354f3dd015203633b857297e2c"
   integrity sha1-BOaSb2YolTVPPdAVIDYzuFcpfiw=
 
+stability-client@1.5.0:
+  version "1.5.0"
+  resolved "https://registry.yarnpkg.com/stability-client/-/stability-client-1.5.0.tgz#f221420a297c808f209c469a0df8fa2401f8f6ae"
+  integrity sha512-hXuDK6QW/msf50pu8M4L1hTCYG4w5ZrhLgxirigUrSZEARupIPT88O7qz4YSmUbbp9nKlzS8UJ6NjYUH7qy45w==
+  dependencies:
+    "@improbable-eng/grpc-web" "^0.15.0"
+    "@improbable-eng/grpc-web-node-http-transport" "^0.15.0"
+    commander "^9.4.0"
+    dotenv "^16.0.2"
+    google-protobuf "^3.21.0"
+    mime "^3.0.0"
+    mkdirp "^1.0.4"
+    read-pkg-up "^7.0.1"
+    typed-emitter "^2.1.0"
+    uuid4 "^2.0.3"
+
 stable@^0.1.8:
   version "0.1.8"
   resolved "https://registry.yarnpkg.com/stable/-/stable-0.1.8.tgz#836eb3c8382fe2936feaf544631017ce7d47a3cf"
@@ -11810,6 +11870,13 @@ type-is@~1.6.18:
     media-typer "0.3.0"
     mime-types "~2.1.24"
 
+typed-emitter@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/typed-emitter/-/typed-emitter-2.1.0.tgz#ca78e3d8ef1476f228f548d62e04e3d4d3fd77fb"
+  integrity sha512-g/KzbYKbH5C2vPkaXGu8DJlHrGKHLsM25Zg9WuC9pMGfuvT+X25tZQWo5fK1BjBm8+UrVE9LDCvaY0CQk+fXDA==
+  optionalDependencies:
+    rxjs "^7.5.2"
+
 typedarray-to-buffer@^3.1.5:
   version "3.1.5"
   resolved "https://registry.yarnpkg.com/typedarray-to-buffer/-/typedarray-to-buffer-3.1.5.tgz#a97ee7a9ff42691b9f783ff1bc5112fe3fca9080"
@@ -12108,6 +12175,11 @@ utils-merge@1.0.1:
   resolved "https://registry.yarnpkg.com/utils-merge/-/utils-merge-1.0.1.tgz#9f95710f50a267947b2ccc124741c1028427e713"
   integrity sha1-n5VxD1CiZ5R7LMwSR0HBAoQn5xM=
 
+uuid4@^2.0.3:
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/uuid4/-/uuid4-2.0.3.tgz#241e5dfe1704a79c52e2aa40e7e581a5e7b01ab4"
+  integrity sha512-CTpAkEVXMNJl2ojgtpLXHgz23dh8z81u6/HEPiQFOvBc/c2pde6TVHmH4uwY0d/GLF3tb7+VDAj4+2eJaQSdZQ==
+
 uuid@^8.0.0, uuid@^8.3.2:
   version "8.3.2"
   resolved "https://registry.yarnpkg.com/uuid/-/uuid-8.3.2.tgz#80d5b5ced271bb9af6c445f21a1a04c606cefbe2"