diff --git a/browser_tests/tests/dialog.spec.ts b/browser_tests/tests/dialog.spec.ts index 993281c64c..bf2de274b7 100644 --- a/browser_tests/tests/dialog.spec.ts +++ b/browser_tests/tests/dialog.spec.ts @@ -104,15 +104,13 @@ test.describe('Missing models warning', () => { }) => { await comfyPage.workflow.loadWorkflow('missing/missing_models') - const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models') - await expect(missingModelsWarning).toBeVisible() + const dialogTitle = comfyPage.page.getByText( + 'This workflow is missing models' + ) + await expect(dialogTitle).toBeVisible() - const downloadButton = missingModelsWarning.getByText('Download') - await expect(downloadButton).toBeVisible() - - // Check that the copy URL button is also visible for Desktop environment - const copyUrlButton = missingModelsWarning.getByText('Copy URL') - await expect(copyUrlButton).toBeVisible() + const downloadAllButton = comfyPage.page.getByText('Download all') + await expect(downloadAllButton).toBeVisible() }) test('Should display a warning when missing models are found in node properties', async ({ @@ -123,15 +121,13 @@ test.describe('Missing models warning', () => { 'missing/missing_models_from_node_properties' ) - const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models') - await expect(missingModelsWarning).toBeVisible() + const dialogTitle = comfyPage.page.getByText( + 'This workflow is missing models' + ) + await expect(dialogTitle).toBeVisible() - const downloadButton = missingModelsWarning.getByText('Download') - await expect(downloadButton).toBeVisible() - - // Check that the copy URL button is also visible for Desktop environment - const copyUrlButton = missingModelsWarning.getByText('Copy URL') - await expect(copyUrlButton).toBeVisible() + const downloadAllButton = comfyPage.page.getByText('Download all') + await expect(downloadAllButton).toBeVisible() }) test('Should not display a warning when no missing models are found', async ({ @@ -172,8 +168,10 @@ test.describe('Missing models warning', () => { await comfyPage.workflow.loadWorkflow('missing/missing_models') - const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models') - await expect(missingModelsWarning).not.toBeVisible() + const dialogTitle = comfyPage.page.getByText( + 'This workflow is missing models' + ) + await expect(dialogTitle).not.toBeVisible() }) test('Should not display warning when model metadata exists but widget values have changed', async ({ @@ -186,8 +184,10 @@ test.describe('Missing models warning', () => { ) // The missing models warning should NOT appear - const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models') - await expect(missingModelsWarning).not.toBeVisible() + const dialogTitle = comfyPage.page.getByText( + 'This workflow is missing models' + ) + await expect(dialogTitle).not.toBeVisible() }) // Flaky test after parallelization @@ -199,13 +199,15 @@ test.describe('Missing models warning', () => { // https://github.com/Comfy-Org/ComfyUI_devtools/blob/main/__init__.py await comfyPage.workflow.loadWorkflow('missing/missing_models') - const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models') - await expect(missingModelsWarning).toBeVisible() + const dialogTitle = comfyPage.page.getByText( + 'This workflow is missing models' + ) + await expect(dialogTitle).toBeVisible() - const downloadButton = comfyPage.page.getByText('Download') - await expect(downloadButton).toBeVisible() + const downloadAllButton = comfyPage.page.getByText('Download all') + await expect(downloadAllButton).toBeVisible() const downloadPromise = comfyPage.page.waitForEvent('download') - await downloadButton.click() + await downloadAllButton.click() const download = await downloadPromise expect(download.suggestedFilename()).toBe('fake_model.safetensors') @@ -229,13 +231,14 @@ test.describe('Missing models warning', () => { test('Should disable warning dialog when checkbox is checked', async ({ comfyPage }) => { - await checkbox.click() const changeSettingPromise = comfyPage.page.waitForRequest( '**/api/settings/Comfy.Workflow.ShowMissingModelsWarning' ) - await closeButton.click() + await checkbox.click() await changeSettingPromise + await closeButton.click() + const settingValue = await comfyPage.settings.getSetting( 'Comfy.Workflow.ShowMissingModelsWarning' ) diff --git a/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/dropped-workflow-url-hidream-dev-example-png-chromium-linux.png b/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/dropped-workflow-url-hidream-dev-example-png-chromium-linux.png index a40febb866..1535c58e45 100644 Binary files a/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/dropped-workflow-url-hidream-dev-example-png-chromium-linux.png and b/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/dropped-workflow-url-hidream-dev-example-png-chromium-linux.png differ diff --git a/browser_tests/tests/viewport.spec.ts-snapshots/viewport-fits-when-saved-offscreen-chromium-linux.png b/browser_tests/tests/viewport.spec.ts-snapshots/viewport-fits-when-saved-offscreen-chromium-linux.png index 4970a12023..937fb6f839 100644 Binary files a/browser_tests/tests/viewport.spec.ts-snapshots/viewport-fits-when-saved-offscreen-chromium-linux.png and b/browser_tests/tests/viewport.spec.ts-snapshots/viewport-fits-when-saved-offscreen-chromium-linux.png differ diff --git a/src/components/common/ElectronFileDownload.vue b/src/components/common/ElectronFileDownload.vue deleted file mode 100644 index eb86541afa..0000000000 --- a/src/components/common/ElectronFileDownload.vue +++ /dev/null @@ -1,148 +0,0 @@ - - - - diff --git a/src/components/common/FileDownload.vue b/src/components/common/FileDownload.vue deleted file mode 100644 index af92752b74..0000000000 --- a/src/components/common/FileDownload.vue +++ /dev/null @@ -1,69 +0,0 @@ - - - - diff --git a/src/components/dialog/content/MissingModelsContent.vue b/src/components/dialog/content/MissingModelsContent.vue new file mode 100644 index 0000000000..09fbeecb96 --- /dev/null +++ b/src/components/dialog/content/MissingModelsContent.vue @@ -0,0 +1,173 @@ + + + diff --git a/src/components/dialog/content/MissingModelsFooter.vue b/src/components/dialog/content/MissingModelsFooter.vue new file mode 100644 index 0000000000..dc6c734d9d --- /dev/null +++ b/src/components/dialog/content/MissingModelsFooter.vue @@ -0,0 +1,103 @@ + + + diff --git a/src/components/dialog/content/MissingModelsHeader.vue b/src/components/dialog/content/MissingModelsHeader.vue new file mode 100644 index 0000000000..df46375fec --- /dev/null +++ b/src/components/dialog/content/MissingModelsHeader.vue @@ -0,0 +1,10 @@ + diff --git a/src/components/dialog/content/MissingModelsWarning.vue b/src/components/dialog/content/MissingModelsWarning.vue deleted file mode 100644 index d7403743bd..0000000000 --- a/src/components/dialog/content/MissingModelsWarning.vue +++ /dev/null @@ -1,177 +0,0 @@ - - - - - diff --git a/src/components/dialog/content/missingModelsUtils.ts b/src/components/dialog/content/missingModelsUtils.ts new file mode 100644 index 0000000000..b5246cfea8 --- /dev/null +++ b/src/components/dialog/content/missingModelsUtils.ts @@ -0,0 +1,83 @@ +import { isDesktop } from '@/platform/distribution/types' +import { useElectronDownloadStore } from '@/stores/electronDownloadStore' + +const ALLOWED_SOURCES = [ + 'https://civitai.com/', + 'https://huggingface.co/', + 'http://localhost:' +] as const + +const ALLOWED_SUFFIXES = [ + '.safetensors', + '.sft', + '.ckpt', + '.pth', + '.pt' +] as const + +const WHITE_LISTED_URLS: ReadonlySet = new Set([ + 'https://huggingface.co/stabilityai/stable-zero123/resolve/main/stable_zero123.ckpt', + 'https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd14v1.pth?download=true', + 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth' +]) + +const DIRECTORY_BADGE_MAP = { + vae: 'VAE', + diffusion_models: 'DIFFUSION', + text_encoders: 'TEXT ENCODER', + loras: 'LORA', + checkpoints: 'CHECKPOINT' +} as const + +export interface ModelWithUrl { + name: string + url: string + directory: string +} + +export function isModelDownloadable(model: ModelWithUrl): boolean { + if (WHITE_LISTED_URLS.has(model.url)) return true + if (!ALLOWED_SOURCES.some((source) => model.url.startsWith(source))) + return false + if (!ALLOWED_SUFFIXES.some((suffix) => model.name.endsWith(suffix))) + return false + return true +} + +export function hasValidDirectory( + model: ModelWithUrl, + paths: Record +): boolean { + return !!paths[model.directory] +} + +export function getBadgeLabel(directory: string): string { + if (directory in DIRECTORY_BADGE_MAP) { + return DIRECTORY_BADGE_MAP[directory as keyof typeof DIRECTORY_BADGE_MAP] + } + return directory.toUpperCase() +} + +export function downloadModel( + model: ModelWithUrl, + paths: Record +): void { + if (!isDesktop) { + const link = document.createElement('a') + link.href = model.url + link.download = model.name + link.target = '_blank' + link.rel = 'noopener noreferrer' + link.click() + return + } + + const modelPaths = paths[model.directory] + if (modelPaths?.[0]) { + void useElectronDownloadStore().start({ + url: model.url, + savePath: modelPaths[0], + filename: model.name + }) + } +} diff --git a/src/composables/useCivitaiModel.ts b/src/composables/useCivitaiModel.ts deleted file mode 100644 index be15f828c7..0000000000 --- a/src/composables/useCivitaiModel.ts +++ /dev/null @@ -1,95 +0,0 @@ -import { useAsyncState } from '@vueuse/core' -import { computed } from 'vue' - -type ModelType = - | 'Checkpoint' - | 'TextualInversion' - | 'Hypernetwork' - | 'AestheticGradient' - | 'LORA' - | 'Controlnet' - | 'Poses' - -interface CivitaiFileMetadata { - fp?: 'fp16' | 'fp32' - size?: 'full' | 'pruned' - format?: 'SafeTensor' | 'PickleTensor' | 'Other' -} - -interface CivitaiModelFile { - name: string - id: number - sizeKB: number - type: string - downloadUrl: string - metadata: CivitaiFileMetadata -} - -interface CivitaiModel { - name: string - type: ModelType -} - -interface CivitaiModelVersionResponse { - id: number - name: string - model: CivitaiModel - modelId: number - files: CivitaiModelFile[] - [key: string]: unknown -} - -/** - * Composable to manage Civitai model - * @param url - The URL of the Civitai model, where the model ID is the last part of the URL's pathname - * @see https://developer.civitai.com/docs/api/public-rest - * @example - * const { fileSize, isLoading, error, modelData } = - * useCivitaiModel('https://civitai.com/api/download/models/16576?type=Model&format=SafeTensor&size=full&fp=fp16') - */ -export function useCivitaiModel(url: string) { - const createModelVersionUrl = (modelId: string): string => - `https://civitai.com/api/v1/model-versions/${modelId}` - - const extractModelIdFromUrl = (): string | null => { - const urlObj = new URL(url) - return urlObj.pathname.split('/').pop() || null - } - - const fetchModelData = - async (): Promise => { - const modelId = extractModelIdFromUrl() - if (!modelId) return null - - const apiUrl = createModelVersionUrl(modelId) - const res = await fetch(apiUrl) - return res.json() - } - - const findMatchingFileSize = (): number | null => { - const matchingFile = modelData.value?.files?.find( - (file) => file.downloadUrl && url.startsWith(file.downloadUrl) - ) - - return matchingFile?.sizeKB ? matchingFile.sizeKB << 10 : null - } - - const { - state: modelData, - isLoading, - error - } = useAsyncState(fetchModelData, null, { - immediate: true - }) - - const fileSize = computed(() => - !isLoading.value ? findMatchingFileSize() : null - ) - - return { - fileSize, - isLoading, - error, - modelData - } -} diff --git a/src/composables/useDownload.ts b/src/composables/useDownload.ts deleted file mode 100644 index 2b1d8924a7..0000000000 --- a/src/composables/useDownload.ts +++ /dev/null @@ -1,67 +0,0 @@ -import { whenever } from '@vueuse/core' -import { onMounted, ref } from 'vue' - -import { useCivitaiModel } from '@/composables/useCivitaiModel' -import { downloadUrlToHfRepoUrl, isCivitaiModelUrl } from '@/utils/formatUtil' - -export function useDownload(url: string, fileName?: string) { - const fileSize = ref(null) - const error = ref(null) - - const setFileSize = (size: number) => { - fileSize.value = size - } - - const fetchFileSize = async () => { - try { - const response = await fetch(url, { method: 'HEAD' }) - if (!response.ok) throw new Error('Failed to fetch file size') - - const size = response.headers.get('content-length') - if (size) { - setFileSize(parseInt(size)) - } else { - console.error('"content-length" header not found') - return null - } - } catch (e) { - console.error('Error fetching file size:', e) - error.value = e instanceof Error ? e : new Error(String(e)) - return null - } - } - - /** - * Trigger browser download - */ - const triggerBrowserDownload = () => { - const link = document.createElement('a') - if (url.includes('huggingface.co') && error.value) { - // If model is a gated HF model, send user to the repo page so they can sign in first - link.href = downloadUrlToHfRepoUrl(url) - } else { - link.href = url - link.download = fileName || url.split('/').pop() || 'download' - } - link.target = '_blank' // Opens in new tab if download attribute is not supported - link.rel = 'noopener noreferrer' // Security best practice for _blank links - link.click() - } - - onMounted(() => { - if (isCivitaiModelUrl(url)) { - const { fileSize: civitaiSize, error: civitaiErr } = useCivitaiModel(url) - whenever(civitaiSize, setFileSize) - // Try falling back to normal fetch if using Civitai API fails - whenever(civitaiErr, fetchFileSize, { once: true }) - } else { - // Fetch file size in the background - void fetchFileSize() - } - }) - - return { - triggerBrowserDownload, - fileSize - } -} diff --git a/src/composables/useMissingModelsDialog.ts b/src/composables/useMissingModelsDialog.ts index d68a73d357..c2e482532d 100644 --- a/src/composables/useMissingModelsDialog.ts +++ b/src/composables/useMissingModelsDialog.ts @@ -1,6 +1,8 @@ import type { ComponentAttrs } from 'vue-component-type-helpers' -import MissingModelsWarning from '@/components/dialog/content/MissingModelsWarning.vue' +import MissingModelsContent from '@/components/dialog/content/MissingModelsContent.vue' +import MissingModelsFooter from '@/components/dialog/content/MissingModelsFooter.vue' +import MissingModelsHeader from '@/components/dialog/content/MissingModelsHeader.vue' import { useDialogService } from '@/services/dialogService' import { useDialogStore } from '@/stores/dialogStore' @@ -14,11 +16,14 @@ export function useMissingModelsDialog() { dialogStore.closeDialog({ key: DIALOG_KEY }) } - function show(props: ComponentAttrs) { + function show(props: ComponentAttrs) { showSmallLayoutDialog({ key: DIALOG_KEY, - component: MissingModelsWarning, - props + headerComponent: MissingModelsHeader, + footerComponent: MissingModelsFooter, + component: MissingModelsContent, + props, + footerProps: props }) } diff --git a/src/locales/en/main.json b/src/locales/en/main.json index aa494011b0..d914c4165f 100644 --- a/src/locales/en/main.json +++ b/src/locales/en/main.json @@ -1749,8 +1749,15 @@ "doNotAskAgain": "Don't show this again", "reEnableInSettings": "Re-enable in {link}", "reEnableInSettingsLink": "Settings", - "missingModels": "Missing Models", - "missingModelsMessage": "When loading the graph, the following models were not found" + "title": "This workflow is missing models", + "description": "This workflow requires models you haven't downloaded yet.", + "totalSize": "Total download size:", + "downloadAll": "Download all", + "downloadAvailable": "Download available", + "gotIt": "Ok, got it", + "footerDescription": "Download and place these models in the correct folder.\nNodes with missing models are highlighted red on the canvas.", + "customModelsWarning": "Some of these are custom models that we don't recognize.", + "customModelsInstruction": "You'll need to find and download them manually. Search for them online (try Civitai or HuggingFace) or contact the original workflow provider." }, "versionMismatchWarning": { "title": "Version Compatibility Warning",