From 7424742ed2c85b9a5be95e56d28c810d0694a0d4 Mon Sep 17 00:00:00 2001
From: Webifi <john@webifi.com>
Date: Wed, 14 Jun 2023 00:34:24 -0500
Subject: [PATCH] Add DALL-E image generation

---
 package-lock.json                     |   7 ++
 package.json                          |   1 +
 src/lib/ApiUtil.svelte                |   2 +
 src/lib/ChatCompletionResponse.svelte |  26 ++++++-
 src/lib/ChatRequest.svelte            | 108 +++++++++++++++++++++-----
 src/lib/EditMessage.svelte            |  53 +++++++++++--
 src/lib/Export.svelte                 |  13 +++-
 src/lib/ImageStore.svelte             |  71 +++++++++++++++++
 src/lib/Models.svelte                 |  21 ++++-
 src/lib/Settings.svelte               |  19 +++++
 src/lib/Storage.svelte                |  36 +++++++--
 src/lib/Types.svelte                  |  32 +++++++-
 12 files changed, 347 insertions(+), 42 deletions(-)
 create mode 100644 src/lib/ImageStore.svelte

diff --git a/package-lock.json b/package-lock.json
index 2f715e4..e93028b 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -22,6 +22,7 @@
         "bulma": "^0.9.4",
         "bulma-prefers-dark": "^0.1.0-beta.1",
         "copy-to-clipboard": "^3.3.3",
+        "dexie": "^4.0.1-alpha.20",
         "eslint-config-standard-with-typescript": "^35.0.0",
         "eslint-plugin-svelte3": "^4.0.0",
         "flourite": "^1.2.3",
@@ -1601,6 +1602,12 @@
         "node": ">=8"
       }
     },
+    "node_modules/dexie": {
+      "version": "4.0.1-alpha.20",
+      "resolved": "https://registry.npmjs.org/dexie/-/dexie-4.0.1-alpha.20.tgz",
+      "integrity": "sha512-q/nMsCQiTWTmnw11aseJLfAsGQ/9t05sjWltgw1/r2TbfnIkmfjdTt8ATSIwmtKXuSznEZ5czazvL5LO5rR+6w==",
+      "dev": true
+    },
     "node_modules/dir-glob": {
       "version": "3.0.1",
       "resolved": "https://registry.npmjs.org/dir-glob/-/dir-glob-3.0.1.tgz",
diff --git a/package.json b/package.json
index b08903e..2eba101 100644
--- a/package.json
+++ b/package.json
@@ -28,6 +28,7 @@
     "bulma": "^0.9.4",
     "bulma-prefers-dark": "^0.1.0-beta.1",
     "copy-to-clipboard": "^3.3.3",
+    "dexie": "^4.0.1-alpha.20",
     "eslint-config-standard-with-typescript": "^35.0.0",
     "eslint-plugin-svelte3": "^4.0.0",
     "flourite": "^1.2.3",
diff --git a/src/lib/ApiUtil.svelte b/src/lib/ApiUtil.svelte
index 7beff67..01b41ff 100644
--- a/src/lib/ApiUtil.svelte
+++ b/src/lib/ApiUtil.svelte
@@ -2,9 +2,11 @@
   // This makes it possible to override the OpenAI API base URL in the .env file
   const apiBase = import.meta.env.VITE_API_BASE || 'https://api.openai.com'
   const endpointCompletions = import.meta.env.VITE_ENDPOINT_COMPLETIONS || '/v1/chat/completions'
+  const endpointGenerations = import.meta.env.VITE_ENDPOINT_GENERATIONS || '/v1/images/generations'
   const endpointModels = import.meta.env.VITE_ENDPOINT_MODELS || '/v1/models'
 
   export const getApiBase = ():string => apiBase
   export const getEndpointCompletions = ():string => endpointCompletions
+  export const getEndpointGenerations = ():string => endpointGenerations
   export const getEndpointModels = ():string => endpointModels
 </script>
\ No newline at end of file
diff --git a/src/lib/ChatCompletionResponse.svelte b/src/lib/ChatCompletionResponse.svelte
index e2afb6d..986784a 100644
--- a/src/lib/ChatCompletionResponse.svelte
+++ b/src/lib/ChatCompletionResponse.svelte
@@ -1,7 +1,8 @@
 <script context="module" lang="ts">
+import { setImage } from './ImageStore.svelte'
 // TODO: Integrate API calls
 import { addMessage, getLatestKnownModel, saveChatStore, setLatestKnownModel, subtractRunningTotal, updateRunningTotal } from './Storage.svelte'
-import type { Chat, ChatCompletionOpts, Message, Model, Response, Usage } from './Types.svelte'
+import type { Chat, ChatCompletionOpts, ChatImage, Message, Model, Response, ResponseImage, Usage } from './Types.svelte'
 import { encode } from 'gpt-tokenizer'
 import { v4 as uuidv4 } from 'uuid'
 
@@ -64,6 +65,29 @@ export class ChatCompletionResponse {
     this.promptTokenCount = tokens
   }
 
+  async updateImageFromSyncResponse (response: ResponseImage, prompt: string, model: Model) {
+    this.setModel(model)
+    for (let i = 0; i < response.data.length; i++) {
+      const d = response.data[i]
+      const message = {
+        role: 'image',
+        uuid: uuidv4(),
+        content: prompt,
+        image: await setImage(this.chat.id, { b64image: d.b64_json } as ChatImage),
+        model,
+        usage: {
+          prompt_tokens: 0,
+          completion_tokens: 1,
+          total_tokens: 1
+        } as Usage
+      } as Message
+      this.messages[i] = message
+      if (this.opts.autoAddMessages) addMessage(this.chat.id, message)
+    }
+    this.notifyMessageChange()
+    this.finish()
+  }
+
   updateFromSyncResponse (response: Response) {
     this.setModel(response.model)
     response.choices.forEach((choice, i) => {
diff --git a/src/lib/ChatRequest.svelte b/src/lib/ChatRequest.svelte
index 7f46329..ae51002 100644
--- a/src/lib/ChatRequest.svelte
+++ b/src/lib/ChatRequest.svelte
@@ -2,12 +2,12 @@
     import { ChatCompletionResponse } from './ChatCompletionResponse.svelte'
     import { mergeProfileFields, prepareSummaryPrompt } from './Profiles.svelte'
     import { countMessageTokens, countPromptTokens, getModelMaxTokens } from './Stats.svelte'
-    import type { Chat, ChatCompletionOpts, ChatSettings, Message, Model, Request } from './Types.svelte'
+    import type { Chat, ChatCompletionOpts, ChatSettings, Message, Model, Request, RequestImageGeneration } from './Types.svelte'
     import { deleteMessage, getChatSettingValueNullDefault, insertMessages, saveChatStore, getApiKey, addError } from './Storage.svelte'
     import { scrollToBottom, scrollToMessage } from './Util.svelte'
     import { getRequestSettingList, defaultModel } from './Settings.svelte'
     import { EventStreamContentType, fetchEventSource } from '@microsoft/fetch-event-source'
-    import { getApiBase, getEndpointCompletions } from './ApiUtil.svelte'
+    import { getApiBase, getEndpointCompletions, getEndpointGenerations } from './ApiUtil.svelte'
     import { v4 as uuidv4 } from 'uuid'
 
 export class ChatRequest {
@@ -25,7 +25,67 @@ export class ChatRequest {
       setChat (chat: Chat) {
         this.chat = chat
       }
+
+      // Common error handler
+      async handleError (response) {
+        let errorResponse
+        try {
+          const errObj = await response.json()
+          errorResponse = errObj?.error?.message || errObj?.error?.code
+          if (!errorResponse && response.choices && response.choices[0]) {
+            errorResponse = response.choices[0]?.message?.content
+          }
+          errorResponse = errorResponse || 'Unexpected Response'
+        } catch (e) {
+          errorResponse = 'Unknown Response'
+        }
+        throw new Error(`${response.status} - ${errorResponse}`)
+      }
     
+      async imageRequest (message: Message, prompt: string, count:number, messages: Message[], opts: ChatCompletionOpts, overrides: ChatSettings = {} as ChatSettings): Promise<ChatCompletionResponse> {
+        const _this = this
+        count = count || 1
+        _this.updating = true
+        _this.updatingMessage = 'Generating Image...'
+        const signal = _this.controller.signal
+        const size = this.chat.settings.imageGenerationSize
+        const request: RequestImageGeneration = {
+          prompt,
+          response_format: 'b64_json',
+          size,
+          n: count
+        }
+        const fetchOptions = {
+          method: 'POST',
+          headers: {
+            Authorization: `Bearer ${getApiKey()}`,
+            'Content-Type': 'application/json'
+          },
+          body: JSON.stringify(request),
+          signal
+        }
+        const chatResponse = new ChatCompletionResponse(opts)
+
+        try {
+          const response = await fetch(getApiBase() + getEndpointGenerations(), fetchOptions)
+          if (!response.ok) {
+            await _this.handleError(response)
+          } else {
+            const json = await response.json()
+            // Remove updating indicator
+            _this.updating = false
+            _this.updatingMessage = ''
+            // console.log('image json', json, json?.data[0])
+            chatResponse.updateImageFromSyncResponse(json, prompt, 'dall-e-' + size)
+          }
+        } catch (e) {
+          chatResponse.updateFromError(e)
+          throw e
+        }
+        message.suppress = true
+        return chatResponse
+      }
+
       /**
        * Send API request
        * @param messages
@@ -38,11 +98,33 @@ export class ChatRequest {
         const chat = _this.chat
         const chatSettings = _this.chat.settings
         const chatId = chat.id
+        const imagePromptDetect = /^\s*(please|can\s+you|will\s+you)*\s*(give|generate|create|show|build|design)\s+(me)*\s*(an|a|set|a\s+set\s+of)*\s*([0-9]+|one|two|three|four)*\s+(image|photo|picture|pic)s*\s*(for\s+me)*\s*(of|[^a-z0-9]+|about|that\s+has|showing|with|having|depicting)\s+[^a-z0-9]*(.*)$/i
         opts.chat = chat
         _this.updating = true
+
+        const lastMessage = messages[messages.length - 1]
+
+        if (chatSettings.imageGenerationSize && !opts.didSummary && !opts.summaryRequest && lastMessage?.role === 'user') {
+          const im = lastMessage.content.match(imagePromptDetect)
+          if (im) {
+            // console.log('image prompt request', im)
+            let n = parseInt((im[5] || '').toLowerCase().trim()
+              .replace(/one/ig, '1')
+              .replace(/two/ig, '2')
+              .replace(/three/ig, '3')
+              .replace(/four/ig, '4')
+            )
+            if (isNaN(n)) n = 1
+            n = Math.min(Math.max(1, n), 4)
+            return await this.imageRequest(lastMessage, im[9], n, messages, opts, overrides)
+            // throw new Error('Image prompt:' + im[7])
+          }
+        }
     
         // Submit only the role and content of the messages, provide the previous messages as well for context
-        const messageFilter = (m:Message) => !m.suppress && m.role !== 'error' && m.content && !m.summarized
+        const messageFilter = (m:Message) => !m.suppress &&
+          ['user', 'assistant', 'system'].includes(m.role) &&
+          m.content && !m.summarized
         const filtered = messages.filter(messageFilter)
     
         // If we're doing continuous chat, do it
@@ -110,22 +192,6 @@ export class ChatRequest {
             signal
           }
 
-          // Common error handler
-          const handleError = async (response) => {
-            let errorResponse
-            try {
-              const errObj = await response.json()
-              errorResponse = errObj?.error?.message || errObj?.error?.code
-              if (!errorResponse && response.choices && response.choices[0]) {
-                errorResponse = response.choices[0]?.message?.content
-              }
-              errorResponse = errorResponse || 'Unexpected Response'
-            } catch (e) {
-              errorResponse = 'Unknown Response'
-            }
-            throw new Error(`${response.status} - ${errorResponse}`)
-          }
-
           // fetchEventSource doesn't seem to throw on abort,
           // so we deal with it ourselves
           const abortListener = (e:Event) => {
@@ -174,7 +240,7 @@ export class ChatRequest {
                 // everything's good
                 } else {
                 // client-side errors are usually non-retriable:
-                  await handleError(response)
+                  await _this.handleError(response)
                 }
               }
             }).catch(err => {
@@ -187,7 +253,7 @@ export class ChatRequest {
              */
             const response = await fetch(getApiBase() + getEndpointCompletions(), fetchOptions)
             if (!response.ok) {
-              await handleError(response)
+              await _this.handleError(response)
             } else {
               const json = await response.json()
               // Remove updating indicator
diff --git a/src/lib/EditMessage.svelte b/src/lib/EditMessage.svelte
index 836c42a..672bd56 100644
--- a/src/lib/EditMessage.svelte
+++ b/src/lib/EditMessage.svelte
@@ -6,10 +6,11 @@
   import SvelteMarkdown from 'svelte-markdown'
   import type { Message, Model, Chat } from './Types.svelte'
   import Fa from 'svelte-fa/src/fa.svelte'
-  import { faTrash, faDiagramPredecessor, faDiagramNext, faCircleCheck, faPaperPlane, faEye, faEyeSlash, faEllipsis } from '@fortawesome/free-solid-svg-icons/index'
+  import { faTrash, faDiagramPredecessor, faDiagramNext, faCircleCheck, faPaperPlane, faEye, faEyeSlash, faEllipsis, faDownload } from '@fortawesome/free-solid-svg-icons/index'
   import { errorNotice, scrollToMessage } from './Util.svelte'
   import { openModal } from 'svelte-modals'
   import PromptConfirm from './PromptConfirm.svelte'
+  import { getImage } from './ImageStore.svelte'
 
   export let message:Message
   export let chatId:number
@@ -22,6 +23,7 @@
   const isSystem = message.role === 'system'
   const isUser = message.role === 'user'
   const isAssistant = message.role === 'assistant'
+  const isImage = message.role === 'image'
 
   // Marked options
   const markdownOptions = {
@@ -34,9 +36,15 @@
   let editing = false
   let original:string
   let defaultModel:Model
+  let imageUrl:string
 
   onMount(() => {
     defaultModel = chatSettings.model
+    if (message?.image) {
+      getImage(message.image.id).then(i => {
+        imageUrl = 'data:image/png;base64, ' + i.b64image
+      })
+    }
   })
 
   const edit = () => {
@@ -168,17 +176,28 @@
     saveChatStore()
   }
 
+  const downloadImage = () => {
+    const filename = (message?.content || `${chat.name}-image-${message?.image?.id}`)
+      .replace(/([^a-z0-9- ]|\.)+/gi, '_').trim().slice(0, 80)
+    const a = document.createElement('a')
+    a.download = `${filename}.png`
+    a.href = imageUrl
+    document.body.appendChild(a)
+    a.click()
+    document.body.removeChild(a)
+  }
+
 </script>
 
 <article
   id="{'message-' + message.uuid}"
   class="message chat-message" 
   class:is-info={isUser}
-  class:is-success={isAssistant}
+  class:is-success={isAssistant || isImage}
   class:is-warning={isSystem}
   class:is-danger={isError}
   class:user-message={isUser || isSystem}
-  class:assistant-message={isError || isAssistant}
+  class:assistant-message={isError || isAssistant || isImage}
   class:summarized={message.summarized}
   class:suppress={message.suppress} 
   class:editing={editing}
@@ -192,6 +211,9 @@
         <div id={'edit-' + message.uuid} class="message-editor" bind:innerText={message.content} contenteditable
         on:input={update} on:blur={exit} />
       </form>
+        {#if imageUrl}
+          <img src={imageUrl} alt="">
+        {/if}
     {:else}
       <div 
         class="message-display" 
@@ -207,6 +229,9 @@
           options={markdownOptions} 
           renderers={{ code: Code, html: Code }}
         />
+        {#if imageUrl}
+          <img src={imageUrl} alt="">
+        {/if}
     </div>
     {/if}
     {#if isSystem}
@@ -261,7 +286,7 @@
       <a
         href={'#'}
         title="Delete this message"
-        class=" msg-delete button is-small"
+        class="msg-delete button is-small"
         on:click|preventDefault={() => {
           checkDelete()
         }}
@@ -273,11 +298,11 @@
       {/if}
       </a>
       {/if}
-      {#if !message.summarized && !isError}
+      {#if !isImage && !message.summarized && !isError}
         <a
           href={'#'}
           title="Truncate from here and send"
-          class=" msg-truncate button is-small"
+          class="msg-truncate button is-small"
           on:click|preventDefault={() => {
             checkTruncate()
           }}
@@ -289,11 +314,11 @@
         {/if}
         </a>
       {/if}
-      {#if !message.summarized && !isSystem && !isError}
+      {#if !isImage && !message.summarized && !isSystem && !isError}
         <a
           href={'#'}
           title={(message.suppress ? 'Uns' : 'S') + 'uppress message from submission'}
-          class=" msg-supress button is-small"
+          class="msg-supress button is-small"
           on:click|preventDefault={() => {
             setSuppress(!message.suppress)
           }}
@@ -305,6 +330,18 @@
         {/if}
         </a>
       {/if}
+      {#if imageUrl}
+        <a
+          href={'#'}
+          title="Download Image"
+          class="msg-image button is-small"
+          on:click|preventDefault={() => {
+            downloadImage()
+          }}
+        >
+        <span class="icon"><Fa icon={faDownload} /></span>
+        </a>
+      {/if}
       </div>
 
   </div>
diff --git a/src/lib/Export.svelte b/src/lib/Export.svelte
index 7050fd2..e7fa582 100644
--- a/src/lib/Export.svelte
+++ b/src/lib/Export.svelte
@@ -1,8 +1,9 @@
 <script context="module" lang="ts">
   import { get } from 'svelte/store'
   import type { Chat } from './Types.svelte'
-  import { chatsStorage } from './Storage.svelte'
+  import { chatsStorage, getChat } from './Storage.svelte'
   import { getExcludeFromProfile } from './Settings.svelte'
+  import { getImage } from './ImageStore.svelte'
 
   export const exportAsMarkdown = (chatId: number) => {
     const chats = get(chatsStorage)
@@ -27,9 +28,13 @@
     document.body.removeChild(a)
   }
 
-  export const exportChatAsJSON = (chatId: number) => {
-    const chats = get(chatsStorage)
-    const chat = chats.find((chat) => chat.id === chatId) as Chat
+  export const exportChatAsJSON = async (chatId: number) => {
+    const chat = JSON.parse(JSON.stringify(getChat(chatId))) as Chat
+    for (let i = 0; i < chat.messages.length; i++) {
+      // Pull images out of indexedDB store for JSON download
+      const m = chat.messages[i]
+      if (m.image) m.image = await getImage(m.image.id)
+    }
     const exportContent = JSON.stringify(chat)
     const blob = new Blob([exportContent], { type: 'text/json' })
     const url = URL.createObjectURL(blob)
diff --git a/src/lib/ImageStore.svelte b/src/lib/ImageStore.svelte
new file mode 100644
index 0000000..93091cc
--- /dev/null
+++ b/src/lib/ImageStore.svelte
@@ -0,0 +1,71 @@
+<script context="module" lang="ts">
+  import Dexie, { type Table } from 'dexie'
+  import type { ChatImage } from './Types.svelte'
+  import { v4 as uuidv4 } from 'uuid'
+
+  let _hasIndexedDb = !!window.indexedDB
+  const dbCheck = _hasIndexedDb && window.indexedDB.open('test')
+  if (_hasIndexedDb) dbCheck.onerror = () => { _hasIndexedDb = false }
+  
+  const imageCache: Record<string, ChatImage> = {}
+
+  class ChatImageStore extends Dexie {
+    images!: Table<ChatImage>
+    constructor () {
+      super('chatImageStore')
+      this.version(1).stores({
+        images: 'id' // Primary key and indexed props
+      })
+    }
+  }
+
+  const imageDb = new ChatImageStore()
+
+  export const hasIndexedDb = () => _hasIndexedDb
+
+  export const getImage = async (uuid:string): Promise<ChatImage> => {
+    let image = imageCache[uuid]
+    if (image || !_hasIndexedDb) return image
+    image = await imageDb.images.get(uuid) as any
+    imageCache[uuid] = image
+    return image
+  }
+
+  export const deleteImage = async (chatId:number, uuid:string): Promise<void> => {
+    const cached = imageCache[uuid]
+    if (cached) cached.chats = cached.chats?.filter(c => c !== chatId)
+    if (!cached?.chats?.length) delete imageCache[uuid]
+    if (_hasIndexedDb) {
+      const stored:ChatImage = await imageDb.images.get({ id: uuid }) as any
+      if (stored) stored.chats = stored.chats?.filter(c => c !== chatId)
+      if (!stored?.chats?.length) {
+        imageDb.images.delete(uuid)
+      } else if (stored) {
+        await setImage(chatId, stored)
+      }
+    }
+  }
+
+  export const setImage = async (chatId:number, image:ChatImage): Promise<ChatImage> => {
+    image.id = image.id || uuidv4()
+    let current: ChatImage
+    if (_hasIndexedDb) {
+      current = await imageDb.images.get({ id: image.id }) as any
+    } else {
+      current = imageCache[image.id]
+    }
+    current = current || image
+    current.chats = current.chats || []
+    if (!(chatId in current.chats)) current.chats.push(chatId)
+    imageCache[current.id] = current
+    if (_hasIndexedDb) {
+      imageDb.images.put(current, current.id)
+    }
+    const clone = JSON.parse(JSON.stringify(current))
+    // Return a copy without the payload so local storage doesn't get clobbered
+    delete clone.b64image
+    delete clone.chats
+    return clone
+  }
+
+</script>
\ No newline at end of file
diff --git a/src/lib/Models.svelte b/src/lib/Models.svelte
index 4d36ea2..a3953c8 100644
--- a/src/lib/Models.svelte
+++ b/src/lib/Models.svelte
@@ -31,6 +31,24 @@ const modelDetails : Record<string, ModelDetail> = {
       }
 }
 
+const imageModels : Record<string, ModelDetail> = {
+      'dall-e-1024x1024': {
+        prompt: 0.00,
+        completion: 0.020, // $0.020 per image
+        max: 1000 // 1000 char prompt, max
+      },
+      'dall-e-512x512': {
+        prompt: 0.00,
+        completion: 0.018, // $0.018 per image
+        max: 1000 // 1000 char prompt, max
+      },
+      'dall-e-256x256': {
+        prompt: 0.00,
+        completion: 0.016, // $0.016 per image
+        max: 1000 // 1000 char prompt, max
+      }
+}
+
 const unknownDetail = {
   prompt: 0,
   completion: 0,
@@ -51,11 +69,12 @@ export const supportedModels : Record<string, ModelDetail> = {
 }
 
 const lookupList = {
+  ...imageModels,
   ...modelDetails,
   ...supportedModels
 }
 
-export const supportedModelKeys = Object.keys(supportedModels)
+export const supportedModelKeys = Object.keys({ ...supportedModels, ...imageModels })
 
 const tpCache : Record<string, ModelDetail> = {}
 
diff --git a/src/lib/Settings.svelte b/src/lib/Settings.svelte
index b3dc4b9..d732417 100644
--- a/src/lib/Settings.svelte
+++ b/src/lib/Settings.svelte
@@ -86,6 +86,7 @@ const defaults:ChatSettings = {
   autoStartSession: false,
   trainingPrompts: [],
   hiddenPromptPrefix: '',
+  imageGenerationSize: '',
   // useResponseAlteration: false,
   // responseAlterations: [],
   isDirty: false
@@ -97,6 +98,12 @@ const excludeFromProfile = {
   isDirty: true
 }
 
+export const imageGenerationSizes = [
+  '1024x1024', '512x512', '256x256'
+]
+
+export const imageGenerationSizeTypes = ['', ...imageGenerationSizes]
+
 const profileSetting: ChatSetting & SettingSelect = {
       key: 'profile',
       name: 'Profile',
@@ -269,6 +276,18 @@ const summarySettings: ChatSetting[] = [
         placeholder: 'Enter a prompt that will be used to summarize past prompts here.',
         type: 'textarea',
         hide: (chatId) => getChatSettings(chatId).continuousChat !== 'summary'
+      },
+      {
+        key: 'imageGenerationSize',
+        name: 'Image Generation Size',
+        header: 'Image Generation',
+        headerClass: 'is-info',
+        title: 'Prompt an image with: show me an image of ...',
+        type: 'select',
+        options: [
+          { value: '', text: 'OFF - Disable Image Generation' },
+          ...imageGenerationSizes.map(s => { return { value: s, text: s } })
+        ]
       }
 ]
 
diff --git a/src/lib/Storage.svelte b/src/lib/Storage.svelte
index 80eafe0..967eca0 100644
--- a/src/lib/Storage.svelte
+++ b/src/lib/Storage.svelte
@@ -6,7 +6,10 @@
   import { v4 as uuidv4 } from 'uuid'
   import { getProfile, getProfiles, isStaticProfile, newNameForProfile, restartProfile } from './Profiles.svelte'
   import { errorNotice } from './Util.svelte'
+  import { deleteImage, setImage } from './ImageStore.svelte'
 
+  // TODO: move chatsStorage to indexedDB with localStorage as a fallback for private browsing.
+  //       Enough long chats will overflow localStorage.
   export const chatsStorage = persisted('chats', [] as Chat[])
   export const latestModelMap = persisted('latestModelMap', {} as Record<Model, Model>) // What was returned when a model was requested
   export const globalStorage = persisted('global', {} as GlobalSettings)
@@ -53,7 +56,7 @@
     return chatId
   }
 
-  export const addChatFromJSON = (json: string): number => {
+  export const addChatFromJSON = async (json: string): Promise<number> => {
     const chats = get(chatsStorage)
 
     // Find the max chatId
@@ -73,6 +76,10 @@
 
     chat.id = chatId
 
+    // Make sure images are moved to indexedDB store,
+    // else they would clobber local storage
+    await updateChatImages(chatId, chat)
+
     // Add a new chat
     chats.push(chat)
     chatsStorage.set(chats)
@@ -154,7 +161,10 @@
   }
 
   export const clearChats = () => {
-    chatsStorage.set([])
+    const chats = get(chatsStorage)
+    chats.forEach(c => deleteChat(c.id)) // make sure images are removed
+    // TODO: add a clear images option to make this faster
+    // chatsStorage.set([])
   }
   export const saveChatStore = () => {
     const chats = get(chatsStorage)
@@ -268,13 +278,16 @@
     const chat = chats.find((chat) => chat.id === chatId) as Chat
     const index = chat.messages.findIndex((m) => m.uuid === uuid)
     const message = getMessage(chat, uuid)
-    if (message && message.summarized) throw new Error('Unable to delete summarized message')
-    if (message && message.summary) throw new Error('Unable to directly delete message summary')
+    if (message?.summarized) throw new Error('Unable to delete summarized message')
+    if (message?.summary) throw new Error('Unable to directly delete message summary')
     // const found = chat.messages.filter((m) => m.uuid === uuid)
     if (index < 0) {
       console.error(`Unable to find and delete message with ID: ${uuid}`)
       return
     }
+    if (message?.image) {
+      deleteImage(chatId, message.image.id)
+    }
     // console.warn(`Deleting message with ID: ${uuid}`, found, index)
     chat.messages.splice(index, 1) // remove item
     chatsStorage.set(chats)
@@ -303,10 +316,21 @@
 
   export const deleteChat = (chatId: number) => {
     const chats = get(chatsStorage)
+    const chat = getChat(chatId)
+    chat?.messages?.forEach(m => {
+      if (m.image) deleteImage(chatId, m.image.id)
+    })
     chatsStorage.set(chats.filter((chat) => chat.id !== chatId))
   }
 
-  export const copyChat = (chatId: number) => {
+  export const updateChatImages = async (chatId: number, chat: Chat) => {
+    for (let i = 0; i < chat.messages.length; i++) {
+      const m = chat.messages[i]
+      if (m.image) m.image = await setImage(chatId, m.image)
+    }
+  }
+
+  export const copyChat = async (chatId: number) => {
     const chats = get(chatsStorage)
     const chat = chats.find((chat) => chat.id === chatId) as Chat
     const nameMap = chats.reduce((a, chat) => { a[chat.name] = chat; return a }, {})
@@ -323,6 +347,8 @@
     // Set new name
     chatCopy.name = cname
 
+    await updateChatImages(chatId, chatCopy)
+
     // Add a new chat
     chats.push(chatCopy)
   
diff --git a/src/lib/Types.svelte b/src/lib/Types.svelte
index 246eedf..f944556 100644
--- a/src/lib/Types.svelte
+++ b/src/lib/Types.svelte
@@ -1,8 +1,11 @@
 <script context="module" lang="ts">
-  import type { supportedModelKeys } from './Models.svelte'
+  import { supportedModelKeys } from './Models.svelte'
+  import { imageGenerationSizeTypes } from './Settings.svelte'
 
   export type Model = typeof supportedModelKeys[number];
 
+  export type ImageGenerationSizes = typeof imageGenerationSizeTypes[number];
+
   export type ModelDetail = {
     prompt: number;
     completion: number;
@@ -15,8 +18,14 @@
     total_tokens: number;
   };
 
+  export interface ChatImage {
+    id: string;
+    b64image: string;
+    chats: number[];
+  }
+
   export type Message = {
-    role: 'user' | 'assistant' | 'system' | 'error';
+    role: 'user' | 'assistant' | 'system' | 'error' | 'image';
     content: string;
     uuid: string;
     usage?: Usage;
@@ -27,6 +36,7 @@
     suppress?: boolean;
     finish_reason?: string;
     streaming?: boolean;
+    image?: ChatImage;
   };
 
   export type ResponseAlteration = {
@@ -35,6 +45,23 @@
     replace: string;
   }
 
+  export type ResponseImageDetail = {
+    url: string;
+    b64_json: string;
+  }
+
+  export type ResponseImage = {
+    created: number;
+    data: ResponseImageDetail[];
+  }
+
+  export type RequestImageGeneration = {
+    prompt: string;
+    n?: number;
+    size?: ImageGenerationSizes;
+    response_format?: keyof ResponseImageDetail;
+  }
+
   export type Request = {
     model: Model;
     messages?: Message[];
@@ -66,6 +93,7 @@
     systemPrompt: string;
     autoStartSession: boolean;
     hiddenPromptPrefix: string;
+    imageGenerationSize: ImageGenerationSizes;
     trainingPrompts?: Message[];
     useResponseAlteration?: boolean;
     responseAlterations?: ResponseAlteration[];