#58 - fix info table for stable Diffusion

Signed-off-by: julianbollig <[email protected]>
TNG · Feb 6, 2025 · 51700da · 51700da
1 parent 3033e2b
commit 51700da
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 19 deletions.
diff --git a/WebUI/src/assets/js/store/stableDiffusion.ts b/WebUI/src/assets/js/store/stableDiffusion.ts
@@ -27,6 +27,11 @@ type BackendParams = {
   safe_check: boolean
 }
 
+type ImageInfoParams = {
+  size: string
+  model_name: string
+} & Partial<BackendParams>
+
 export const useStableDiffusion = defineStore(
   'stableDiffusion',
   () => {
@@ -36,6 +41,7 @@ export const useStableDiffusion = defineStore(
     const models = useModels()
 
     let abortController: AbortController | null
+    const defaultBackendParams = ref<BackendParams>()
 
     async function generate() {
       if (imageGeneration.processing) {
@@ -44,7 +50,7 @@ export const useStableDiffusion = defineStore(
       try {
         imageGeneration.processing = true
         await checkModel()
-        const defaultBackendParams = {
+        defaultBackendParams.value = {
           mode: 0,
           device: globalSetup.modelSettings.graphics,
           prompt: imageGeneration.prompt,
@@ -53,15 +59,15 @@ export const useStableDiffusion = defineStore(
           generate_number: imageGeneration.batchSize,
           inference_steps: imageGeneration.inferenceSteps,
           guidance_scale: imageGeneration.guidanceScale,
-          seed: imageGeneration.seed,
+          seed: imageGeneration.seed === -1 ? Math.random() * 1000000 : imageGeneration.seed,
           height: imageGeneration.height,
           width: imageGeneration.width,
           lora: imageGeneration.lora,
           scheduler: imageGeneration.scheduler,
           image_preview: imageGeneration.imagePreview,
           safe_check: imageGeneration.safeCheck,
         }
-        await sendGenerate(defaultBackendParams)
+        await sendGenerate()
       } catch (_error: unknown) {
       } finally {
         imageGeneration.processing = false
@@ -122,13 +128,16 @@ 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,
-            createInfoParamTable(data.params),
-          )
+          const infoParams: KVObject | undefined =
+            defaultBackendParams.value &&
+            createInfoParamTable({
+              ...defaultBackendParams.value,
+              size: String(data.params.size),
+              model_name: String(data.params.model_name),
+            })
+          await imageGeneration.updateImage(data.index, data.image, false, infoParams)
           break
+
         case 'step_end':
           imageGeneration.currentState = 'generating'
           imageGeneration.stepText = `${i18nState.COM_GENERATING} ${data.step}/${data.total_step}`
@@ -169,15 +178,16 @@ export const useStableDiffusion = defineStore(
       }
     }
 
-    async function sendGenerate(defaultBackendParams: BackendParams) {
+    async function sendGenerate() {
       try {
         imageGeneration.processing = true
         if (!abortController) {
           abortController = new AbortController()
         }
+
         const response = await fetch(`${useGlobalSetup().apiHost}/api/sd/generate`, {
           method: 'POST',
-          body: util.convertToFormData(defaultBackendParams),
+          body: util.convertToFormData(defaultBackendParams.value),
           signal: abortController.signal,
         })
         const reader = response.body!.getReader()
@@ -200,7 +210,7 @@ export const useStableDiffusion = defineStore(
       }
     }
 
-    function createInfoParamTable(infoParams: KVObject) {
+    function createInfoParamTable(infoParams: ImageInfoParams) {
       const infoParamsTable: KVObject = {
         resolution: infoParams.width + 'x' + infoParams.height,
         size: infoParams.size,
@@ -214,7 +224,7 @@ export const useStableDiffusion = defineStore(
         seed: infoParams.seed,
         scheduler: infoParams.scheduler,
         lora: infoParams.lora,
-        safe_check: infoParams.safe_check.toBoolean(),
+        safe_check: infoParams.safe_check,
       }
       return infoParamsTable
     }

diff --git a/WebUI/src/env.d.ts b/WebUI/src/env.d.ts
@@ -128,6 +128,10 @@ type StringKV = {
   [key: string]: string
 }
 
+type StringOrNumberKV = {
+  [key: string]: string | number
+}
+
 type WebSettings = {
   graphics: { name: string; index: number }[]
   schedulers: string[]
@@ -205,8 +209,8 @@ type LLMOutTextCallback = {
 
 type SDOutCallback =
   | LoadModelCallback
-  | LoadModelcomponentsCallback
-  | SDOutImagelCallback
+  | LoadModelComponentsCallback
+  | SDOutImageCallback
   | SDStepEndCallback
   | ErrorOutCallback
   | NotEnoughDiskSpaceExceptionCallback
@@ -229,17 +233,17 @@ type LoadModelCallback = {
   event: 'start' | 'finish'
 }
 
-type LoadModelcomponentsCallback = {
+type LoadModelComponentsCallback = {
   type: 'load_model_components'
   event: 'start' | 'finish'
 }
 
-type SDOutImagelCallback = {
+type SDOutImageCallback = {
   type: 'image_out'
   index: number
   image: string
   safe_check_pass: boolean
-  params: KVObject
+  params: StringOrNumberKV
 }
 
 type SDStepEndCallback = {

diff --git a/WebUI/src/lib/utils.ts b/WebUI/src/lib/utils.ts
@@ -64,7 +64,7 @@ export function mapServiceNameToDisplayName(serviceName: string) {
   }
 }
 
-export function mapModeToText(value: number) {
+export function mapModeToText(value: number | undefined) {
   const i18nState = useI18N().state
   switch (value) {
     case 0: