From af378262f44a964a6a6ef2e76f03704820ec28cc Mon Sep 17 00:00:00 2001 From: "Alex \"mcmonkey\" Goodwin" <4000772+mcmonkey4eva@users.noreply.github.com> Date: Fri, 23 Aug 2024 06:43:20 -0700 Subject: [PATCH] Model downloader dialog (#569) * API core for model downloader * initial basic dialog for missing models * app.ts handling for missing models * don't explode if getModels is a 404 * actually track downloads in progress * overall pile of improvements to the missing models view * minor fixes * add setting to disable missing models warning * temporarily remove 'models' entry from default graph to avoid missing model dialog causing issues. Also because ckpt autodownloading shouldn't be allowed * swap the url to a title * add model directory to display * match settingStore commit * check setting before scanning models list ie avoid redundant calcs when setting is disabled anyway --- .../dialog/content/MissingModelsWarning.vue | 262 ++++++++++++++++++ src/scripts/api.ts | 47 ++++ src/scripts/app.ts | 55 +++- src/scripts/defaultGraph.ts | 11 +- src/services/dialogService.ts | 12 + src/stores/settingStore.ts | 7 + src/types/apiTypes.ts | 10 + tests-ui/globalSetup.ts | 3 +- tests-ui/utils/setup.ts | 1 + 9 files changed, 396 insertions(+), 12 deletions(-) create mode 100644 src/components/dialog/content/MissingModelsWarning.vue diff --git a/src/components/dialog/content/MissingModelsWarning.vue b/src/components/dialog/content/MissingModelsWarning.vue new file mode 100644 index 000000000..3a226ad32 --- /dev/null +++ b/src/components/dialog/content/MissingModelsWarning.vue @@ -0,0 +1,262 @@ + + + + + + + diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 817670893..11da47abb 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -1,5 +1,6 @@ import { ComfyWorkflowJSON } from '@/types/comfyWorkflow' import { + DownloadModelStatus, HistoryTaskItem, PendingTaskItem, RunningTaskItem, @@ -216,6 +217,11 @@ class ComfyApi extends EventTarget { new CustomEvent('execution_cached', { detail: msg.data }) ) break + case 'download_progress': + this.dispatchEvent( + new CustomEvent('download_progress', { detail: msg.data }) + ) + break default: if (this.#registered.has(msg.type)) { this.dispatchEvent( @@ -319,6 +325,47 @@ class ComfyApi extends EventTarget { return await res.json() } + /** + * Gets a list of models in the specified folder + * @param {string} folder The folder to list models from, such as 'checkpoints' + * @returns The list of model filenames within the specified folder + */ + async getModels(folder: string) { + const res = await this.fetchApi(`/models/${folder}`) + if (res.status === 404) { + return null + } + return await res.json() + } + + /** + * Tells the server to download a model from the specified URL to the specified directory and filename + * @param {string} url The URL to download the model from + * @param {string} model_directory The main directory (eg 'checkpoints') to save the model to + * @param {string} model_filename The filename to save the model as + * @param {number} progress_interval The interval in seconds at which to report download progress (via 'download_progress' event) + */ + async internalDownloadModel( + url: string, + model_directory: string, + model_filename: string, + progress_interval: number + ): Promise { + const res = await this.fetchApi('/internal/models/download', { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify({ + url, + model_directory, + model_filename, + progress_interval + }) + }) + return await res.json() + } + /** * Loads a list of items (queue or history) * @param {string} type The type of items to load, queue or history diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 8b44c5f07..4f6942b46 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -45,7 +45,8 @@ import { Vector2 } from '@comfyorg/litegraph' import _ from 'lodash' import { showExecutionErrorDialog, - showLoadWorkflowWarning + showLoadWorkflowWarning, + showMissingModelsWarning } from '@/services/dialogService' import { useSettingStore } from '@/stores/settingStore' import { useToastStore } from '@/stores/toastStore' @@ -134,6 +135,7 @@ export class ComfyApp { bodyBottom: HTMLElement canvasContainer: HTMLElement menu: ComfyAppMenu + modelsInFolderCache: Record constructor() { this.vueAppReady = false @@ -148,6 +150,7 @@ export class ComfyApp { parent: document.body }) this.menu = new ComfyAppMenu(this) + this.modelsInFolderCache = {} /** * List of extensions that are registered with the app @@ -2175,6 +2178,22 @@ export class ComfyApp { }) } + showMissingModelsError(missingModels) { + if ( + this.vueAppReady && + useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning') + ) { + showMissingModelsWarning({ + missingModels, + maximizable: true + }) + } + + this.logging.addEntry('Comfy.App', 'warn', { + MissingModels: missingModels + }) + } + async changeWorkflow(callback, workflow = null) { try { this.workflowManager.activeWorkflow?.changeTracker?.store() @@ -2233,10 +2252,12 @@ export class ComfyApp { } const missingNodeTypes = [] + const missingModels = [] await this.#invokeExtensionsAsync( 'beforeConfigureGraph', graphData, missingNodeTypes + // TODO: missingModels ) for (let n of graphData.nodes) { // Patch T2IAdapterLoader to ControlNetLoader since they are the same node now @@ -2251,6 +2272,19 @@ export class ComfyApp { n.type = sanitizeNodeName(n.type) } } + if (graphData.models && useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')) { + for (let m of graphData.models) { + const models_available = await this.getModelsInFolderCached(m.directory) + if (models_available === null) { + // @ts-expect-error + m.directory_invalid = true + missingModels.push(m) + } + else if (!models_available.includes(m.name)) { + missingModels.push(m) + } + } + } try { this.graph.configure(graphData) @@ -2366,9 +2400,13 @@ export class ComfyApp { this.#invokeExtensions('loadedGraphNode', node) } + // TODO: Properly handle if both nodes and models are missing (sequential dialogs?) if (missingNodeTypes.length) { this.showMissingNodesError(missingNodeTypes) } + if (missingModels.length) { + this.showMissingModelsError(missingModels) + } await this.#invokeExtensionsAsync('afterConfigureGraph', missingNodeTypes) requestAnimationFrame(() => { this.graph.setDirtyCanvas(true, true) @@ -2837,6 +2875,19 @@ export class ComfyApp { app.graph.arrange() } + /** + * Gets the list of model names in a folder, using a temporary local cache + */ + async getModelsInFolderCached(folder: string): Promise { + if (folder in this.modelsInFolderCache) { + return this.modelsInFolderCache[folder] + } + // TODO: needs a lock to avoid overlapping calls + const models = await api.getModels(folder) + this.modelsInFolderCache[folder] = models + return models + } + /** * Registers a Comfy web extension with the app * @param {ComfyExtension} extension @@ -2862,6 +2913,8 @@ export class ComfyApp { } if (this.vueAppReady) useToastStore().add(requestToastMessage) + this.modelsInFolderCache = {} + const defs = await api.getNodeDefs() for (const nodeId in defs) { diff --git a/src/scripts/defaultGraph.ts b/src/scripts/defaultGraph.ts index 753a18ce7..ebf7c5b10 100644 --- a/src/scripts/defaultGraph.ts +++ b/src/scripts/defaultGraph.ts @@ -133,14 +133,5 @@ export const defaultGraph: ComfyWorkflowJSON = { groups: [], config: {}, extra: {}, - version: 0.4, - models: [ - { - name: 'v1-5-pruned-emaonly.ckpt', - url: 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt', - hash: 'cc6cb27103417325ff94f52b7a5d2dde45a7515b25c255d8e396c90014281516', - hash_type: 'SHA256', - directory: 'checkpoints' - } - ] + version: 0.4 } diff --git a/src/services/dialogService.ts b/src/services/dialogService.ts index 86cb222c9..a25c105e4 100644 --- a/src/services/dialogService.ts +++ b/src/services/dialogService.ts @@ -3,6 +3,7 @@ // about importing primevue components. import { useDialogStore } from '@/stores/dialogStore' import LoadWorkflowWarning from '@/components/dialog/content/LoadWorkflowWarning.vue' +import MissingModelsWarning from '@/components/dialog/content/MissingModelsWarning.vue' import SettingDialogContent from '@/components/dialog/content/SettingDialogContent.vue' import SettingDialogHeader from '@/components/dialog/header/SettingDialogHeader.vue' import type { ExecutionErrorWsMessage } from '@/types/apiTypes' @@ -20,6 +21,17 @@ export function showLoadWorkflowWarning(props: { }) } +export function showMissingModelsWarning(props: { + missingModels: any[] + [key: string]: any +}) { + const dialogStore = useDialogStore() + dialogStore.showDialog({ + component: MissingModelsWarning, + props + }) +} + export function showSettingsDialog() { useDialogStore().showDialog({ headerComponent: SettingDialogHeader, diff --git a/src/stores/settingStore.ts b/src/stores/settingStore.ts index 7e82ceb62..4b4b3c76b 100644 --- a/src/stores/settingStore.ts +++ b/src/stores/settingStore.ts @@ -173,6 +173,13 @@ export const useSettingStore = defineStore('setting', { defaultValue: true }) + app.ui.settings.addSetting({ + id: 'Comfy.Workflow.ShowMissingModelsWarning', + name: 'Show missing models warning', + type: 'boolean', + defaultValue: true + }) + app.ui.settings.addSetting({ id: 'Comfy.Graph.ZoomSpeed', name: 'Canvas zoom speed', diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 0b6b5f117..e3fdcfb6d 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -73,6 +73,13 @@ const zExecutionErrorWsMessage = zExecutionWsMessageBase.extend({ current_outputs: z.any() }) +const zDownloadModelStatus = z.object({ + status: z.string(), + progress_percentage: z.number(), + message: z.string(), + already_existed: z.boolean(), +}) + export type StatusWsMessageStatus = z.infer export type StatusWsMessage = z.infer export type ProgressWsMessage = z.infer @@ -87,6 +94,8 @@ export type ExecutionInterruptedWsMessage = z.infer< typeof zExecutionInterruptedWsMessage > export type ExecutionErrorWsMessage = z.infer + +export type DownloadModelStatus = z.infer // End of ws messages const zPromptInputItem = z.object({ @@ -403,6 +412,7 @@ const zSettings = z.record(z.any()).and( 'Comfy.ConfirmClear': z.boolean(), 'Comfy.DevMode': z.boolean(), 'Comfy.Workflow.ShowMissingNodesWarning': z.boolean(), + 'Comfy.Workflow.ShowMissingModelsWarning': z.boolean(), 'Comfy.DisableFloatRounding': z.boolean(), 'Comfy.DisableSliders': z.boolean(), 'Comfy.DOMClippingEnabled': z.boolean(), diff --git a/tests-ui/globalSetup.ts b/tests-ui/globalSetup.ts index 439563e1b..16c6d9c22 100644 --- a/tests-ui/globalSetup.ts +++ b/tests-ui/globalSetup.ts @@ -15,7 +15,8 @@ module.exports = async function () { jest.mock('@/services/dialogService', () => { return { - showLoadWorkflowWarning: jest.fn() + showLoadWorkflowWarning: jest.fn(), + showMissingModelsWarning: jest.fn() } }) } diff --git a/tests-ui/utils/setup.ts b/tests-ui/utils/setup.ts index 808ff51ae..fe414b210 100644 --- a/tests-ui/utils/setup.ts +++ b/tests-ui/utils/setup.ts @@ -73,6 +73,7 @@ export function mockApi(config: APIConfig = {}) { userConfig.users[username + '!'] = username return { status: 200, json: () => username + '!' } }), + getModels: jest.fn(() => []), getUserConfig: jest.fn( () => userConfig ?? { storage: 'browser', migrated: false } ),