Generate images from StableDiffusion

web/.gitignore

@@ -2,4 +2,5 @@
 .next
 node_modules
 out
-tsconfig.tsbuildinfo
+tsconfig.tsbuildinfo
+.env*

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

web/components/editor/image-modal.tsx

@@ -0,0 +1,159 @@
+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>
+  )
+}
+
+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 url = `/api/v0/dream`
+    const data = {
+      prompt: input + ', ' + MODIFIERS,
+      apiKey: API_KEY,
+    }
+    const headers = {
+      'Content-Type': 'application/json',
+    }
+    const response = await fetch(url, {
+      method: 'POST',
+      headers,
+      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>
+  )
+}

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"
   },

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',
+  })
+
+  // const body = JSON.parse(req.body)
+  // 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) {
+    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',
+  })
+}