Skip to content

Commit

Permalink
#58 - refactored code with both backends where images are stored as o…
Browse files Browse the repository at this point in the history
…bjects rather than just urls

Signed-off-by: julianbollig <[email protected]>
  • Loading branch information
julianbollig committed Feb 4, 2025
1 parent 9a7c392 commit 2caea3c
Show file tree
Hide file tree
Showing 4 changed files with 44 additions and 65 deletions.
12 changes: 7 additions & 5 deletions WebUI/src/assets/js/store/comfyUi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export const useComfyUi = defineStore(
const websocket = ref<WebSocket | null>(null)
const clientId = '12345'
const loaderNodes = ref<string[]>([])
const generateIdx = ref<number>(0)

const backendServices = useBackendServices()
const comfyUiState = computed(() => {
Expand Down Expand Up @@ -163,8 +164,7 @@ export const useComfyUi = defineStore(
const imageUrl = URL.createObjectURL(imageBlob)
console.log('image url', imageUrl)
if (imageBlob) {
imageGeneration.previewIdx = imageGeneration.generateIdx
imageGeneration.updateDestImage(imageGeneration.generateIdx, imageUrl)
imageGeneration.updateImage(generateIdx.value, imageUrl, true)
}
break
default:
Expand Down Expand Up @@ -194,11 +194,12 @@ export const useComfyUi = defineStore(
const images: { filename: string; type: string; subfolder: string }[] =
msg.data?.output?.images?.filter((i: { type: string }) => i.type === 'output')
images.forEach((image) => {
imageGeneration.updateDestImage(
imageGeneration.generateIdx,
imageGeneration.updateImage(
generateIdx.value,
`${comfyBaseUrl.value}/view?filename=${image.filename}&type=${image.type}&subfolder=${image.subfolder ?? ''}`,
false,
)
imageGeneration.generateIdx++
generateIdx.value++
})
console.log('executed', { detail: msg.data })
break
Expand Down Expand Up @@ -328,6 +329,7 @@ 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)}`)

Expand Down
28 changes: 13 additions & 15 deletions WebUI/src/assets/js/store/imageGeneration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ export type StableDiffusionSettings = {
export type generatedImage = {
id: number
imageUrl: string
isLoading: boolean
infoParams: KVObject | undefined
}

Expand Down Expand Up @@ -595,7 +596,6 @@ export const useImageGeneration = defineStore(
const currentState = ref<SDGenerateState>('no_start')
const stepText = ref('')
const previewIdx = ref(0)
const generateIdx = ref(-999)

function loadSettingsForActiveWorkflow() {
console.log('loading settings for', activeWorkflowName.value)
Expand Down Expand Up @@ -626,25 +626,27 @@ export const useImageGeneration = defineStore(
getSavedOrDefault('inpaintModel')
}

async function updateDestImage(index: number, image: string) {
if (index + 1 > imageUrls.value.length) {
imageUrls.value.push(image)
} else {
imageUrls.value.splice(index, 1, image)
}
}

async function updateImage(
index: number,
image: string,
loading: boolean,
infoParams: KVObject | undefined = undefined,
) {
const newImage: generatedImage = { id: index, imageUrl: image, infoParams: infoParams }
const newImage: generatedImage = {
id: index,
imageUrl: image,
isLoading: loading,
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)
previewIdx.value = newImage.id
console.log('#################')
console.log(previewIdx)
console.log('#################')
}
}

Expand Down Expand Up @@ -754,12 +756,10 @@ export const useImageGeneration = defineStore(
}

async function generate() {
generateIdx.value = 0
previewIdx.value = 0
stepText.value = i18nState.COM_GENERATING
if (activeWorkflow.value.backend === 'default') {
await stableDiffusion.generate()
console.log(generatedImages.value)
} else {
await comfyUi.generate()
}
Expand All @@ -774,7 +774,7 @@ export const useImageGeneration = defineStore(
currentState.value = 'no_start'
stableDiffusion.generateParams.length = 0
imageUrls.value.length = 0
generateIdx.value = -999
generatedImages.value.length = 0 //new
previewIdx.value = -1
}

Expand All @@ -794,7 +794,6 @@ export const useImageGeneration = defineStore(
stepText,
stopping,
previewIdx,
generateIdx,
imageModel,
inpaintModel,
lora,
Expand All @@ -816,7 +815,6 @@ export const useImageGeneration = defineStore(
loadWorkflowsFromJson,
loadWorkflowsFromIntel,
getMissingModels,
updateDestImage,
updateImage,
generate,
stopGeneration,
Expand Down
17 changes: 2 additions & 15 deletions WebUI/src/assets/js/store/stableDiffusion.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,6 @@ export const useStableDiffusion = defineStore(
image_preview: imageGeneration.imagePreview,
safe_check: imageGeneration.safeCheck,
}

await sendGenerate(defaultBackendParams)
} catch (_error: unknown) {
} finally {
Expand Down Expand Up @@ -123,25 +122,13 @@ 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++ //necessary?
//new
await imageGeneration.updateImage(data.index, data.image, data.params)
await imageGeneration.updateImage(data.index, data.image, false)
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
await imageGeneration.updateImage(data.index, data.image, true)
}
break
case 'load_model':
Expand Down
52 changes: 22 additions & 30 deletions WebUI/src/views/Create.vue
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@
<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-if="imageGeneration.generatedImages.length > 0" class="image-preview-panel flex-none">
<div
v-for="(image, i) in imageGeneration.imageUrls"
v-for="image in imageGeneration.generatedImages"
class="image-preview-item flex items-center justify-center"
:class="{ active: imageGeneration.previewIdx == i }"
@click="swithPreview(i)"
:class="{ active: imageGeneration.previewIdx === image.id }"
@click="switchPreview(image.id)"
>
<div class="image-preview-item-bg">
<img :src="image" class="image-thumb" />
<img :src="image.imageUrl" class="image-thumb" />
</div>
</div>
</div>
Expand All @@ -21,17 +21,14 @@
>
<!-- shows the currently selected image in the main box-->
<img
v-for="(image, i) in imageGeneration.imageUrls"
:src="image"
v-for="image in imageGeneration.generatedImages"
:src="image.imageUrl"
class="p-1 max-w-768px max-h-512px"
v-show="imageGeneration.previewIdx == i"
v-show="imageGeneration.previewIdx === image.id"
/>
<!-- shows loading stage in main box-->
<div
v-show="
imageGeneration.generateIdx == imageGeneration.previewIdx &&
imageGeneration.processing
"
v-show="imageGeneration.processing && (!currentImage || currentImage?.isLoading)"
class="absolute left-0 top-0 w-full h-full bg-black/50 flex justify-center items-center"
>
<loading-bar
Expand All @@ -53,10 +50,7 @@
</div>
<!-- side options on what to do with the generated image-->
<div
v-show="
imageGeneration.previewIdx != -1 &&
imageGeneration.generateIdx > imageGeneration.previewIdx //I don't know what that does
"
v-show="currentImage && !currentImage.isLoading"
class="absolute bottom-0 -right-8 box-content flex flex-col items-center justify-center gap-2"
>
<button
Expand All @@ -68,7 +62,7 @@
</button>
<button
v-if="
imageGeneration.activeWorkflow.backend === 'default' //change after comfyui integration
imageGeneration.activeWorkflow.backend === 'default' //ToDo: change after comfyui integration
"
@click="showParamsDialog"
:title="languages.COM_OPEN_PARAMS"
Expand Down Expand Up @@ -97,11 +91,8 @@
>
<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>
</button> -->

<!-- change reset button in delete image-->
<button
v-show="!imageGeneration.processing"
@click="reset"
Expand Down Expand Up @@ -156,7 +147,7 @@ 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 { useImageGeneration } from '@/assets/js/store/imageGeneration'
import { generatedImage, useImageGeneration } from '@/assets/js/store/imageGeneration'
import { useStableDiffusion } from '@/assets/js/store/stableDiffusion'
import { useGlobalSetup } from '@/assets/js/store/globalSetup.ts'
Expand All @@ -171,7 +162,9 @@ const downloadModel = reactive({
})
const showParams = ref(false)
const infoParams = ref<KVObject>({})
const currentImage: ComputedRef<generatedImage | undefined> = computed(() => {
return imageGeneration.generatedImages.find((image) => image.id === imageGeneration.previewIdx)
})
const emits = defineEmits<{
(
e: 'showDownloadModelConfirm',
Expand Down Expand Up @@ -205,20 +198,19 @@ async function ensureModelsAreAvailable() {
}
function postImageToEnhance() {
emits('postImageToEnhance', imageGeneration.imageUrls[imageGeneration.previewIdx])
emits('postImageToEnhance', currentImage.value?.imageUrl ?? '')
}
function openImage() {
const path = imageGeneration.imageUrls[imageGeneration.previewIdx]
window.electronAPI.openImageWithSystem(path)
window.electronAPI.openImageWithSystem(currentImage.value?.imageUrl ?? '')
}
function selectedImage() {
window.electronAPI.selecteImage(imageGeneration.imageUrls[imageGeneration.previewIdx])
window.electronAPI.selecteImage(currentImage.value?.imageUrl ?? '')
}
function copyImage() {
util.copyImage(imageGeneration.imageUrls[imageGeneration.previewIdx])
util.copyImage(currentImage.value?.imageUrl ?? '')
toast.success(i18nState.COM_COPY_SUCCESS_TIP)
}
Expand All @@ -227,7 +219,7 @@ function reset() {
imageGeneration.reset()
}
function swithPreview(i: number) {
function switchPreview(i: number) {
imageGeneration.previewIdx = i
if (i > -1) {
infoParams.value = stableDiffusion.generateParams[i]
Expand Down

0 comments on commit 2caea3c

Please sign in to comment.