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

…Diffusion Signed-off-by: julianbollig <[email protected]>
TNG · Feb 4, 2025 · 9a7c392 · 9a7c392
1 parent 22a4781
commit 9a7c392
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 18 deletions.
diff --git a/WebUI/src/assets/js/store/imageGeneration.ts b/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(),
@@ -416,7 +422,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
@@ -584,6 +590,7 @@ export const useImageGeneration = defineStore(
       saveToSettingsPerWorkflow('inpaintModel')
     })
 
+    const generatedImages = ref<generatedImage[]>([])
     const imageUrls = ref<string[]>([])
     const currentState = ref<SDGenerateState>('no_start')
     const stepText = ref('')
@@ -627,6 +634,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()
@@ -737,9 +758,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()
       }
     }
 
@@ -766,6 +788,7 @@ export const useImageGeneration = defineStore(
       activeWorkflow,
       processing,
       prompt,
+      generatedImages,
       imageUrls,
       currentState,
       stepText,
@@ -794,6 +817,7 @@ export const useImageGeneration = defineStore(
       loadWorkflowsFromIntel,
       getMissingModels,
       updateDestImage,
+      updateImage,
       generate,
       stopGeneration,
       reset,

diff --git a/WebUI/src/assets/js/store/stableDiffusion.ts b/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

diff --git a/WebUI/src/env.d.ts b/WebUI/src/env.d.ts
@@ -239,7 +239,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 = {

diff --git a/WebUI/src/views/Create.vue b/WebUI/src/views/Create.vue
@@ -1,6 +1,7 @@
 <template>
   <div id="createPanel" class="h-full flex flex-col p-4">
     <div class="image-panel flex-auto flex p-4">
+      <!--      Sidebar for display of generated images-->
       <div v-if="imageGeneration.imageUrls.length > 0" class="image-preview-panel flex-none">
         <div
           v-for="(image, i) in imageGeneration.imageUrls"
@@ -13,16 +14,19 @@
           </div>
         </div>
       </div>
+      <!--      Main box in where image generation process and finished image is displayed-->
       <div class="flex-auto relative flex items-center justify-center">
         <div
           class="flex justify-center items-center w-768px h-512px relative bg-color-image-bg rounded-lg border border-white/30"
         >
+          <!--          shows the currently selected image in the main box-->
           <img
             v-for="(image, i) in imageGeneration.imageUrls"
             :src="image"
             class="p-1 max-w-768px max-h-512px"
             v-show="imageGeneration.previewIdx == i"
           />
+          <!--          shows loading stage in main box-->
           <div
             v-show="
               imageGeneration.generateIdx == imageGeneration.previewIdx &&
@@ -47,10 +51,11 @@
               <span class="text-2xl">{{ imageGeneration.stepText }}</span>
             </div>
           </div>
+          <!--          side options on what to do with the generated image-->
           <div
             v-show="
               imageGeneration.previewIdx != -1 &&
-              imageGeneration.generateIdx > imageGeneration.previewIdx
+              imageGeneration.generateIdx > imageGeneration.previewIdx //I don't know what that does
             "
             class="absolute bottom-0 -right-8 box-content flex flex-col items-center justify-center gap-2"
           >
@@ -62,7 +67,9 @@
               <span class="svg-icon text-white i-transfer w-4 h-4"></span>
             </button>
             <button
-              v-if="imageGeneration.activeWorkflow.backend === 'default'"
+              v-if="
+                imageGeneration.activeWorkflow.backend === 'default' //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"
@@ -84,12 +91,13 @@
               <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"
             >
               <span class="svg-icon text-white i-folder w-4 h-4"></span>
             </button>
+            <!-- why is this commented out? delete -->
             <!-- <button @click="regenerate" :title="languages.COM_REGENERATE"
                             class="bg-color-image-tool-button rounded-sm w-6 h-6 flex items-center justify-center">
                             <span class="svg-icon text-white i-refresh w-4 h-4"></span>
@@ -104,13 +112,15 @@
             </button>
           </div>
         </div>
+        <!--        extra panel displaying the image properties-->
         <paint-info
           :params="infoParams"
           v-show="showParams"
           @close="showParams = false"
         ></paint-info>
       </div>
     </div>
+    <!--    Generating prompts -->
     <div class="h-32 gap-3 flex-none flex items-center">
       <textarea
         class="rounded-xl border border-color-spilter flex-auto h-full resize-none"
@@ -174,7 +184,7 @@ const emits = defineEmits<{
 
 async function generateImage() {
   await ensureModelsAreAvailable()
-  reset()
+  reset() // change if one wants to save images in the queue
   const inferenceBackendService: BackendServiceName =
     imageGeneration.backend === 'comfyui' ? 'comfyui-backend' : 'ai-backend'
   await globalSetup.resetLastUsedInferenceBackend(inferenceBackendService)
@@ -203,7 +213,7 @@ function openImage() {
   window.electronAPI.openImageWithSystem(path)
 }
 
-function selecteImage() {
+function selectedImage() {
   window.electronAPI.selecteImage(imageGeneration.imageUrls[imageGeneration.previewIdx])
 }
 
@@ -229,5 +239,6 @@ function swithPreview(i: number) {
 function showParamsDialog() {
   showParams.value = true
   infoParams.value = stableDiffusion.generateParams[imageGeneration.previewIdx]
+  console.log(infoParams.value)
 }
 </script>