feat: change image queueing logic to enhance UX

Refs: #59 Signed-off-by: julianbollig <julian.bollig@tngtech.com>
TNG · Feb 25, 2025 · 2c09112 · 2c09112
1 parent 87efe18
commit 2c09112
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 24 deletions.
diff --git a/WebUI/src/assets/js/store/comfyUi.ts b/WebUI/src/assets/js/store/comfyUi.ts
@@ -240,11 +240,17 @@ export const useComfyUi = defineStore(
             }
           } else {
             const msg = JSON.parse(event.data)
+            let currentImage: Media
             switch (msg.type) {
               case 'status':
                 break
               case 'progress':
                 imageGeneration.currentState = 'generating'
+                currentImage = queuedImages[generateIdx]
+                imageGeneration.updateImage({
+                  ...currentImage,
+                  state: 'generating',
+                })
                 imageGeneration.stepText = `${i18nState.COM_GENERATING} ${msg.data.value}/${msg.data.max}`
                 console.log('progress', { data: msg.data })
                 break
@@ -268,7 +274,11 @@ export const useComfyUi = defineStore(
                   subfolder: string
                   format: string
                 } = msg.data?.output?.gifs?.find((i: { type: string }) => i.type === 'output')
-                const currentImage = queuedImages[generateIdx]
+                if (!imageFromOutput && !gifFromOutput) {
+                  console.log('nothing to update')
+                  break
+                }
+                currentImage = queuedImages[generateIdx]
                 let newImage: Media
                 if (imageFromOutput) {
                   newImage = {
@@ -384,7 +394,7 @@ export const useComfyUi = defineStore(
       }
     }
 
-    async function generate() {
+    async function generate(imageIds: string[]) {
       console.log('generateWithComfy')
       if (imageGeneration.activeWorkflow.backend !== 'comfyui') {
         console.warn('The selected workflow is not a comfyui workflow')
@@ -424,14 +434,13 @@ export const useComfyUi = defineStore(
           ...findKeysByClassType(mutableWorkflow, 'Unet Loader (GGUF)'),
           ...findKeysByClassType(mutableWorkflow, 'DualCLIPLoader (GGUF)'),
         ]
-
         queuedImages = Array.from({ length: imageGeneration.batchSize }, (_, i) => {
           const seed = baseSeed + i
           const settings = imageGeneration.getGenerationParameters()
           settings.seed = seed
 
           return {
-            id: window.crypto.randomUUID(),
+            id: imageIds[i],
             imageUrl:
               'data:image/svg+xml,%3C%3Fxml%20version%3D%221.0%22%20encoding%3D%22UTF-8%22%3F%3E%3Csvg%20width%3D%2224px%22%20height%3D%2224px%22%20viewBox%3D%220%200%2024%2024%22%20stroke-width%3D%221.5%22%20fill%3D%22none%22%20xmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20color%3D%22%23000000%22%3E%3Cpath%20d%3D%22M12%2012C15.866%2012%2019%208.86599%2019%205H5C5%208.86599%208.13401%2012%2012%2012ZM12%2012C15.866%2012%2019%2015.134%2019%2019H5C5%2015.134%208.13401%2012%2012%2012Z%22%20stroke%3D%22%23000000%22%20stroke-width%3D%221.5%22%20stroke-linecap%3D%22round%22%20stroke-linejoin%3D%22round%22%3E%3C%2Fpath%3E%3Cpath%20d%3D%22M5%202L12%202L19%202%22%20stroke%3D%22%23000000%22%20stroke-width%3D%221.5%22%20stroke-linecap%3D%22round%22%20stroke-linejoin%3D%22round%22%3E%3C%2Fpath%3E%3Cpath%20d%3D%22M5%2022H12L19%2022%22%20stroke%3D%22%23000000%22%20stroke-width%3D%221.5%22%20stroke-linecap%3D%22round%22%20stroke-linejoin%3D%22round%22%3E%3C%2Fpath%3E%3C%2Fsvg%3E',
             state: 'queued',

diff --git a/WebUI/src/assets/js/store/imageGeneration.ts b/WebUI/src/assets/js/store/imageGeneration.ts
@@ -9,11 +9,6 @@ import { useModels } from './models'
 import { useBackendServices } from './backendServices'
 import { useGlobalSetup } from './globalSetup'
 
-export type RefImage = {
-  type: string
-  image: string
-}
-
 export type GenerateState =
   | 'no_start'
   | 'input_image'
@@ -25,10 +20,6 @@ export type GenerateState =
   | 'image_out'
   | 'error'
 
-export type ImageInfoParameter = {
-  [key: string]: string | number | boolean | RefImage
-}
-
 export type GenerationSettings = Partial<
   Settings & {
     workflow: string
@@ -688,7 +679,7 @@ export const useImageGeneration = defineStore(
       getSavedOrDefault('inpaintModel')
     }
 
-    async function updateImage(newImage: Media) {
+    function updateImage(newImage: Media) {
       console.log('updating image', newImage)
       const existingImageIndex = generatedImages.value.findIndex((img) => img.id === newImage.id)
       if (existingImageIndex !== -1) {
@@ -811,16 +802,25 @@ export const useImageGeneration = defineStore(
 
     async function generate() {
       generatedImages.value = generatedImages.value.filter((item) => item.state === 'done')
+      const imageIds: string[] = Array.from({ length: batchSize.value }, () => crypto.randomUUID())
+      imageIds.forEach((imageId) => {
+        updateImage({
+          id: imageId,
+          state: 'queued',
+          settings: {},
+          imageUrl: '',
+        })
+      })
       currentState.value = 'no_start'
       stepText.value = i18nState.COM_GENERATING
       const inferenceBackendService: BackendServiceName =
         backend.value === 'comfyui' ? 'comfyui-backend' : 'ai-backend'
       await backendServices.resetLastUsedInferenceBackend(inferenceBackendService)
       backendServices.updateLastUsedBackend(inferenceBackendService)
       if (activeWorkflow.value.backend === 'default') {
-        await stableDiffusion.generate()
+        await stableDiffusion.generate(imageIds)
       } else {
-        await comfyUi.generate()
+        await comfyUi.generate(imageIds)
       }
     }
 

diff --git a/WebUI/src/assets/js/store/stableDiffusion.ts b/WebUI/src/assets/js/store/stableDiffusion.ts
@@ -110,15 +110,15 @@ export const useStableDiffusion = defineStore(
       })
     }
 
-    async function generate() {
+    async function generate(imageIds: string[]) {
       if (imageGeneration.processing) {
         return
       }
       try {
         imageGeneration.processing = true
         await checkModel()
-        queuedImages = Array.from({ length: imageGeneration.batchSize }, () => ({
-          id: window.crypto.randomUUID(),
+        queuedImages = Array.from({ length: imageGeneration.batchSize }, (_, i) => ({
+          id: imageIds[i],
           imageUrl:
             'data:image/svg+xml,%3C%3Fxml%20version%3D%221.0%22%20encoding%3D%22UTF-8%22%3F%3E%3Csvg%20width%3D%2224px%22%20height%3D%2224px%22%20viewBox%3D%220%200%2024%2024%22%20stroke-width%3D%221.5%22%20fill%3D%22none%22%20xmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20color%3D%22%23000000%22%3E%3Cpath%20d%3D%22M12%2012C15.866%2012%2019%208.86599%2019%205H5C5%208.86599%208.13401%2012%2012%2012ZM12%2012C15.866%2012%2019%2015.134%2019%2019H5C5%2015.134%208.13401%2012%2012%2012Z%22%20stroke%3D%22%23000000%22%20stroke-width%3D%221.5%22%20stroke-linecap%3D%22round%22%20stroke-linejoin%3D%22round%22%3E%3C%2Fpath%3E%3Cpath%20d%3D%22M5%202L12%202L19%202%22%20stroke%3D%22%23000000%22%20stroke-width%3D%221.5%22%20stroke-linecap%3D%22round%22%20stroke-linejoin%3D%22round%22%3E%3C%2Fpath%3E%3Cpath%20d%3D%22M5%2022H12L19%2022%22%20stroke%3D%22%23000000%22%20stroke-width%3D%221.5%22%20stroke-linecap%3D%22round%22%20stroke-linejoin%3D%22round%22%3E%3C%2Fpath%3E%3C%2Fsvg%3E',
           state: 'queued',
@@ -193,7 +193,7 @@ export const useStableDiffusion = defineStore(
           currentImage.settings.seed = data.params.seed
           currentImage.imageUrl = data.image
 
-          await imageGeneration.updateImage(currentImage)
+          imageGeneration.updateImage(currentImage)
           break
 
         case 'step_end':
@@ -204,7 +204,7 @@ export const useStableDiffusion = defineStore(
           if (data.image) {
             currentImage.imageUrl = data.image
           }
-          await imageGeneration.updateImage(currentImage)
+          imageGeneration.updateImage(currentImage)
           break
         case 'load_model':
           imageGeneration.currentState = 'load_model'

diff --git a/WebUI/src/views/Create.vue b/WebUI/src/views/Create.vue
@@ -9,7 +9,7 @@
         <div class="image-preview-panel">
           <!-- eslint-disable vue/require-v-for-key -->
           <div
-            v-for="image in imageGeneration.generatedImages"
+            v-for="image in imageGeneration.generatedImages.filter((i) => i.state !== 'queued')"
             class="image-preview-item flex items-center justify-center"
             :class="{ active: selectedImageId === image.id }"
             @click="selectedImageId = image.id"
@@ -191,9 +191,11 @@ const currentImage: ComputedRef<Media | null> = computed(() => {
   return imageGeneration.generatedImages.find((image) => image.id === selectedImageId.value) ?? null
 })
 watch(
-  () => imageGeneration.generatedImages.length,
+  () => imageGeneration.generatedImages.filter((i) => i.state !== 'queued').length,
   () => {
-    const numberOfImages = imageGeneration.generatedImages.length
+    const numberOfImages = imageGeneration.generatedImages.filter(
+      (i) => i.state !== 'queued',
+    ).length
     if (numberOfImages > 0) {
       selectedImageId.value = imageGeneration.generatedImages[numberOfImages - 1].id
     } else {