#58 - created generateImages type and implemented an array for stable…

…Diffusion

Signed-off-by: julianbollig <julian.bollig@tngtech.com>

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

@@ -26,6 +26,12 @@ export type StableDiffusionSettings = {
   safetyCheck: boolean
 }
 
+export type generatedImage = {
+  id: number
+  imageUrl: string
+  infoParams: KVObject | undefined
+}
+
 const SettingsSchema = z.object({
   imageModel: z.string(),
   inpaintModel: z.string(),
@@ -407,7 +413,7 @@ export const useImageGeneration = defineStore(
     const seed = ref<number>(generalDefaultSettings.seed)
     const imagePreview = ref<boolean>(generalDefaultSettings.imagePreview)
     const safeCheck = ref<boolean>(generalDefaultSettings.safeCheck)
-    const batchSize = ref<number>(globalDefaultSettings.batchSize) // TODO this should be imageCount instead, as we only support batchSize 1 due to memory constraints
+    const batchSize = ref<number>(globalDefaultSettings.batchSize)
 
     const resetActiveWorkflowSettings = () => {
       prompt.value = generalDefaultSettings.prompt
@@ -575,6 +581,7 @@ export const useImageGeneration = defineStore(
       saveToSettingsPerWorkflow('inpaintModel')
     })
 
+    const generatedImages = ref<generatedImage[]>([])
     const imageUrls = ref<string[]>([])
     const currentState = ref<SDGenerateState>('no_start')
     const stepText = ref('')
@@ -618,6 +625,20 @@ export const useImageGeneration = defineStore(
       }
     }
 
+    async function updateImage(
+      index: number,
+      image: string,
+      infoParams: KVObject | undefined = undefined,
+    ) {
+      const newImage: generatedImage = { id: index, imageUrl: image, infoParams: infoParams }
+      const existingImageIndex = generatedImages.value.findIndex((img) => img.id === newImage.id)
+      if (existingImageIndex !== -1) {
+        generatedImages.value.splice(existingImageIndex, 1, newImage)
+      } else {
+        generatedImages.value.push(newImage)
+      }
+    }
+
     async function loadWorkflowsFromIntel() {
       const syncResponse = await window.electronAPI.updateWorkflowsFromIntelRepo()
       await loadWorkflowsFromJson()
@@ -732,9 +753,10 @@ export const useImageGeneration = defineStore(
       previewIdx.value = 0
       stepText.value = i18nState.COM_GENERATING
       if (activeWorkflow.value.backend === 'default') {
-        stableDiffusion.generate()
+        await stableDiffusion.generate()
+        console.log(generatedImages.value)
       } else {
-        comfyUi.generate()
+        await comfyUi.generate()
       }
     }
 
@@ -761,6 +783,7 @@ export const useImageGeneration = defineStore(
       activeWorkflow,
       processing,
       prompt,
+      generatedImages,
       imageUrls,
       currentState,
       stepText,
@@ -789,6 +812,7 @@ export const useImageGeneration = defineStore(
       loadWorkflowsFromIntel,
       getMissingModels,
       updateDestImage,
+      updateImage,
       generate,
       stopGeneration,
       reset,

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

@@ -34,7 +34,7 @@ export const useStableDiffusion = defineStore(
     const i18nState = useI18N().state
     const models = useModels()
 
-    let abortContooler: AbortController | null
+    let abortController: AbortController | null
     const generateParams = ref(new Array<KVObject>())
 
     async function generate() {
@@ -123,17 +123,24 @@ export const useStableDiffusion = defineStore(
           if (!data.safe_check_pass) {
             data.image = '/src/assets/image/nsfw_result_detected.png'
           }
+          //old
           await imageGeneration.updateDestImage(data.index, data.image)
           generateParams.value.push(data.params)
-          imageGeneration.generateIdx++
+          imageGeneration.generateIdx++ //necessary?
+          //new
+          await imageGeneration.updateImage(data.index, data.image, data.params)
           break
         case 'step_end':
           imageGeneration.currentState = 'generating'
           imageGeneration.stepText = `${i18nState.COM_GENERATING} ${data.step}/${data.total_step}`
           if (data.image) {
+            //old
             await imageGeneration.updateDestImage(data.index, data.image)
+            //new
+            await imageGeneration.updateImage(data.index, data.image)
           }
           if (data.step == 0) {
+            //necessary?
             imageGeneration.previewIdx = data.index
           }
           break
@@ -162,7 +169,7 @@ export const useStableDiffusion = defineStore(
             case 'runtime_error':
               toast.error(i18nState.ERROR_RUNTIME_ERROR)
               break
-            case 'unknow_exception':
+            case 'unknown_exception':
               toast.error(i18nState.ERROR_GENERATE_UNKONW_EXCEPTION)
               break
           }
@@ -173,13 +180,13 @@ export const useStableDiffusion = defineStore(
     async function sendGenerate(defaultBackendParams: BackendParams) {
       try {
         imageGeneration.processing = true
-        if (!abortContooler) {
-          abortContooler = new AbortController()
+        if (!abortController) {
+          abortController = new AbortController()
         }
         const response = await fetch(`${useGlobalSetup().apiHost}/api/sd/generate`, {
           method: 'POST',
           body: util.convertToFormData(defaultBackendParams),
-          signal: abortContooler.signal,
+          signal: abortController.signal,
         })
         const reader = response.body!.getReader()
         await new SSEProcessor(reader, dataProcess, finishGenerate).start()
@@ -192,9 +199,9 @@ export const useStableDiffusion = defineStore(
       if (imageGeneration.processing && !imageGeneration.stopping) {
         imageGeneration.stopping = true
         await fetch(`${globalSetup.apiHost}/api/sd/stopGenerate`)
-        if (abortContooler) {
-          abortContooler.abort()
-          abortContooler = null
+        if (abortController) {
+          abortController.abort()
+          abortController = null
         }
         imageGeneration.processing = false
         imageGeneration.stopping = false

WebUI/src/env.d.ts

@@ -259,7 +259,7 @@ type NotEnoughDiskSpaceExceptionCallback = {
 
 type ErrorOutCallback = {
   type: 'error'
-  err_type: 'runtime_error' | 'download_exception' | 'unknow_exception'
+  err_type: 'runtime_error' | 'download_exception' | 'unknown_exception'
 }
 
 type DownloadModelProgressCallback = {

WebUI/src/views/Create.vue

@@ -88,7 +88,7 @@
               <span class="svg-icon text-white i-copy w-4 h-4"></span>
             </button>
             <button
-              @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"
             >
@@ -207,7 +207,7 @@ function openImage() {
   window.electronAPI.openImageWithSystem(path)
 }
 
-function selecteImage() {
+function selectedImage() {
   window.electronAPI.selecteImage(imageGeneration.imageUrls[imageGeneration.previewIdx])
 }