#58 - change infoParam output in StableDiffusion to a mapping table, …

…allow for infoParam in ComfyUi

Signed-off-by: julianbollig <[email protected]>

WebUI/electron/main.ts

@@ -523,7 +523,7 @@ function initEventHandle() {
     shell.openPath(imagePath)
   })
 
-  ipcMain.on('selecteImage', (_event, url: string) => {
+  ipcMain.on('selectedImage', (_event, url: string) => {
     const imagePath = getImagePathFromUrl(url)
     if (!imagePath) return
 

WebUI/electron/preload.ts

@@ -61,7 +61,7 @@ contextBridge.exposeInMainWorld('electronAPI', {
   getDownloadedGGUFLLMs: () => ipcRenderer.invoke('getDownloadedGGUFLLMs'),
   getDownloadedEmbeddingModels: () => ipcRenderer.invoke('getDownloadedEmbeddingModels'),
   openImageWithSystem: (url: string) => ipcRenderer.send('openImageWithSystem', url),
-  selecteImage: (url: string) => ipcRenderer.send('selecteImage', url),
+  selectedImage: (url: string) => ipcRenderer.send('selectedImage', url),
   setFullScreen: (enable: boolean) => ipcRenderer.send('setFullScreen', enable),
   onDebugLog: (callback: (data: { level: string; source: string; message: string }) => void) =>
     ipcRenderer.on('debugLog', (_event, value) => callback(value)),

WebUI/src/assets/js/store/comfyUi.ts

@@ -20,6 +20,7 @@ export const useComfyUi = defineStore(
     const websocket = ref<WebSocket | null>(null)
     const clientId = '12345'
     const loaderNodes = ref<string[]>([])
+    const pendingPrompts = ref<ComfyUiPrompt[]>([])
     const generateIdx = ref<number>(0)
 
     const backendServices = useBackendServices()
@@ -191,16 +192,19 @@ export const useComfyUi = defineStore(
                 }
                 break
               case 'executed':
-                const images: { filename: string; type: string; subfolder: string }[] =
-                  msg.data?.output?.images?.filter((i: { type: string }) => i.type === 'output')
-                images.forEach((image) => {
-                  imageGeneration.updateImage(
-                    generateIdx.value,
-                    `${comfyBaseUrl.value}/view?filename=${image.filename}&type=${image.type}&subfolder=${image.subfolder ?? ''}`,
-                    false,
-                  )
-                  generateIdx.value++
-                })
+                const newImage: { filename: string; type: string; subfolder: string } =
+                  msg.data?.output?.images?.find((i: { type: string }) => i.type === 'output')
+                const promptIndex: number = pendingPrompts.value.findIndex(
+                  (item: ComfyUiPrompt) => item.promptId === msg.data?.prompt_id,
+                )
+                imageGeneration.updateImage(
+                  generateIdx.value,
+                  `${comfyBaseUrl.value}/view?filename=${newImage.filename}&type=${newImage.type}&subfolder=${newImage.subfolder ?? ''}`,
+                  false,
+                  pendingPrompts.value[promptIndex].prompt,
+                )
+                pendingPrompts.value.splice(promptIndex, 1)
+                generateIdx.value++
                 console.log('executed', { detail: msg.data })
                 break
               case 'execution_start':
@@ -329,7 +333,6 @@ export const useComfyUi = defineStore(
           ...findKeysByClassType(mutableWorkflow, 'DualCLIPLoader (GGUF)'),
         ]
 
-        generateIdx.value = 0
         for (let i = 0; i < imageGeneration.batchSize; i++) {
           modifySettingInWorkflow(mutableWorkflow, 'seed', `${(seed + i).toFixed(0)}`)
 
@@ -348,7 +351,15 @@ export const useComfyUi = defineStore(
               `ComfyUI Backend responded with ${result.status}: ${await result.text()}`,
             )
           }
+          const response = await result.json()
+          pendingPrompts.value.push({
+            promptId: response.prompt_id,
+            prompt: mutableWorkflow,
+          })
         }
+        console.log('################')
+        console.log(pendingPrompts.value)
+        console.log('################')
       } catch (ex) {
         console.error('Error generating image', ex)
         toast.error('Backend could not generate image.')
@@ -374,9 +385,16 @@ export const useComfyUi = defineStore(
       })
     }
 
+    function reset() {
+      loaderNodes.value.length = 0
+      pendingPrompts.value.length = 0
+      generateIdx.value = 0
+    }
+
     return {
       generate,
       stop,
+      reset,
     }
   },
   {
@@ -397,17 +415,25 @@ const settingToComfyInputsName = {
   scheduler: ['scheduler'],
   batchSize: ['batch_size'],
 } satisfies Partial<Record<Setting, string[]>>
+
 type ComfySetting = keyof typeof settingToComfyInputsName
+type ComfyUiPrompt = {
+  promptId: string
+  prompt: ComfyUIApiWorkflow
+}
+
 const findKeysByTitle = (workflow: ComfyUIApiWorkflow, title: ComfySetting | 'loader' | string) =>
   Object.entries(workflow)
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     .filter(([_key, value]) => (value as any)?.['_meta']?.title === title)
     .map(([key, _value]) => key)
+
 const findKeysByClassType = (workflow: ComfyUIApiWorkflow, classType: string) =>
   Object.entries(workflow)
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     .filter(([_key, value]) => (value as any)?.['class_type'] === classType)
     .map(([key, _value]) => key)
+
 const findKeysByInputsName = (workflow: ComfyUIApiWorkflow, setting: ComfySetting) => {
   for (const inputName of settingToComfyInputsName[setting]) {
     if (inputName === 'text') continue
@@ -419,6 +445,7 @@ const findKeysByInputsName = (workflow: ComfyUIApiWorkflow, setting: ComfySettin
   }
   return []
 }
+
 const getInputNameBySettingAndKey = (
   workflow: ComfyUIApiWorkflow,
   key: string,

WebUI/src/assets/js/store/imageGeneration.ts

@@ -6,26 +6,6 @@ import { useI18N } from './i18n'
 import * as Const from '../const'
 import { useGlobalSetup } from './globalSetup'
 import * as toast from '@/assets/js/toast.ts'
-import { preview } from 'vite'
-
-export type StableDiffusionSettings = {
-  resolution: 'standard' | 'hd' | 'manual' // ~ modelSettings.resolution 0, 1, 3
-  quality: 'standard' | 'high' | 'fast' // ~ modelSettings.quality 0, 1, 2
-  imageModel: string
-  inpaintModel: string
-  negativePrompt: string
-  batchSize: number // ~ modelSettings.generateNumber
-  pickerResolution?: string
-  width: number
-  height: number
-  guidanceScale: number
-  inferenceSteps: number
-  seed: number
-  lora: string | null
-  scheduler: string | null
-  imagePreview: boolean
-  safetyCheck: boolean
-}
 
 export type generatedImage = {
   id: number
@@ -596,7 +576,7 @@ export const useImageGeneration = defineStore(
     const generatedImages = ref<generatedImage[]>([])
     const currentState = ref<SDGenerateState>('no_start')
     const stepText = ref('')
-    const previewIdx = ref(0)
+    const previewIdx = ref<string | number>(0)
 
     function loadSettingsForActiveWorkflow() {
       console.log('loading settings for', activeWorkflowName.value)
@@ -784,6 +764,7 @@ export const useImageGeneration = defineStore(
       currentState.value = 'no_start'
       stepText.value = ''
       previewIdx.value = 0
+      comfyUi.reset()
     }
 
     loadWorkflowsFromJson()

WebUI/src/assets/js/store/stableDiffusion.ts

@@ -7,6 +7,7 @@ import * as util from '../util'
 import { SSEProcessor } from '../sseProcessor'
 import { useI18N } from './i18n'
 import * as toast from '../toast'
+import { mapModeToText } from '@/lib/utils.ts'
 
 type BackendParams = {
   mode: number
@@ -121,7 +122,12 @@ export const useStableDiffusion = defineStore(
           if (!data.safe_check_pass) {
             data.image = '/src/assets/image/nsfw_result_detected.png'
           }
-          await imageGeneration.updateImage(data.index, data.image, false, data.params)
+          await imageGeneration.updateImage(
+            data.index,
+            data.image,
+            false,
+            createInfoParamTable(data.params),
+          )
           break
         case 'step_end':
           imageGeneration.currentState = 'generating'
@@ -194,6 +200,25 @@ export const useStableDiffusion = defineStore(
       }
     }
 
+    function createInfoParamTable(infoParams: KVObject) {
+      const infoParamsTable: KVObject = {
+        resolution: infoParams.width + 'x' + infoParams.height,
+        size: infoParams.size,
+        Device: infoParams.device,
+        prompt: infoParams.prompt,
+        model_name: infoParams.model_name,
+        mode: mapModeToText(infoParams.mode),
+        negative_prompt: infoParams.negative_prompt,
+        inference_steps: infoParams.inference_steps,
+        guidance_scale: infoParams.guidance_scale,
+        seed: infoParams.seed,
+        scheduler: infoParams.scheduler,
+        lora: infoParams.lora,
+        safe_check: infoParams.safe_check.toBoolean(),
+      }
+      return infoParamsTable
+    }
+
     return {
       generate,
       stop,

WebUI/src/components/InfoTable.vue

@@ -0,0 +1,28 @@
+<template>
+  <div v-if="props.params" class="info-params absolute px-5 pt-8 pb-5 text-white">
+    <button
+      class="w-5 h-5 svg-icon i-close absolute right-2 top-2"
+      @click="emits('close')"
+    ></button>
+    <div class="params-list">
+      <ul class="border border-color-spilter">
+        <li
+          v-for="(value, key) in props.params"
+          class="last:border-none border-b border-color-spilter flex items-start"
+        >
+          <span class="text-base font-bold px-4 items-stretch w-36 flex-none">{{ key }}</span>
+          <span class="px-4 flex-auto break-word">{{ value }}</span>
+        </li>
+      </ul>
+    </div>
+  </div>
+</template>
+<script setup lang="ts">
+const props = defineProps<{
+  params: KVObject | undefined
+}>()
+
+const emits = defineEmits<{
+  (e: 'close'): void
+}>()
+</script>

WebUI/src/lib/utils.ts

@@ -63,3 +63,21 @@ export function mapServiceNameToDisplayName(serviceName: string) {
       return serviceName
   }
 }
+
+export function mapModeToText(value: number) {
+  const i18nState = useI18N().state
+  switch (value) {
+    case 0:
+      return i18nState.TAB_CREATE
+    case 1:
+      return i18nState.ENHANCE_UPSCALE
+    case 2:
+      return i18nState.ENHANCE_IMAGE_PROMPT
+    case 3:
+      return i18nState.ENHANCE_INPAINT
+    case 4:
+      return i18nState.ENHANCE_OUTPAINT
+    default:
+      return 'unknown'
+  }
+}

WebUI/src/views/Create.vue

@@ -56,9 +56,6 @@
               <span class="svg-icon text-white i-transfer w-4 h-4"></span>
             </button>
             <button
-              v-if="
-                imageGeneration.activeWorkflow.backend === 'default' //ToDo: change after comfyui integration
-              "
               @click="showParamsDialog"
               :title="languages.COM_OPEN_PARAMS"
               class="bg-color-image-tool-button rounded-sm w-6 h-6 flex items-center justify-center"
@@ -103,11 +100,11 @@
             </button>
           </div>
         </div>
-        <paint-info
+        <info-table
           v-show="showInfoParams"
           :params="currentImage?.infoParams"
           @close="showInfoParams = false"
-        ></paint-info>
+        ></info-table>
       </div>
     </div>
     <div class="h-32 gap-3 flex-none flex items-center">
@@ -144,7 +141,7 @@ import { useI18N } from '@/assets/js/store/i18n'
 import * as toast from '@/assets/js/toast'
 import * as util from '@/assets/js/util'
 import LoadingBar from '../components/LoadingBar.vue'
-import PaintInfo from '@/components/PaintInfo.vue'
+import InfoTable from '@/components/InfoTable.vue'
 import { generatedImage, useImageGeneration } from '@/assets/js/store/imageGeneration'
 import { useGlobalSetup } from '@/assets/js/store/globalSetup.ts'
 

WebUI/src/views/Enhance.vue

@@ -162,7 +162,7 @@
             </button>
             <button
               v-show="previewIdx != -1 && previewIdx <= generateFinishIdx"
-              @click="selecteImage"
+              @click="selectedImage"
               :title="languages.COM_OPEN_LOCATION"
               class="bg-color-image-tool-button rounded-sm w-6 h-6 flex items-center justify-center"
             >
@@ -321,7 +321,7 @@ import * as toast from '@/assets/js/toast'
 import { SSEProcessor } from '@/assets/js/sseProcessor'
 import LoadingBar from '../components/LoadingBar.vue'
 import VerticalSlideBar from '@/components/VerticalSlideBar.vue'
-import PaintInfo from '@/components/PaintInfo.vue'
+import PaintInfo from '@/components/InfoTable.vue'
 import { useGlobalSetup } from '@/assets/js/store/globalSetup'
 import InpaintMask from '../components/InpaintMask.vue'
 import * as Const from '@/assets/js/const'
@@ -691,7 +691,7 @@ function copyImage() {
   util.copyImage(destImg.value[previewIdx.value])
 }
 
-async function selecteImage() {
+async function selectedImage() {
   const url = destImg.value[previewIdx.value]
   window.electronAPI.selectedImage(url)
 }