diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/Fields/ModelImageUpload.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/Fields/ModelImageUpload.tsx index 0e28c1802cb..929d16285df 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/Fields/ModelImageUpload.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/Fields/ModelImageUpload.tsx @@ -31,8 +31,15 @@ type Props = { const ModelImageUpload = ({ model_key, model_image }: Props) => { const [image, setImage] = useState(model_image || null); + const [prevModelImage, setPrevModelImage] = useState(model_image); const { t } = useTranslation(); + // Sync local state when the model_image prop changes (e.g. switching models) without a cascading effect. + if (model_image !== prevModelImage) { + setPrevModelImage(model_image); + setImage(model_image || null); + } + const [updateModelImage, request] = useUpdateModelImageMutation(); const [deleteModelImage] = useDeleteModelImageMutation(); diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsExportButton.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsExportButton.tsx index 66aeda4551e..916c0f94d94 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsExportButton.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsExportButton.tsx @@ -1,51 +1,33 @@ import { IconButton } from '@invoke-ai/ui-library'; import { toast } from 'features/toast/toast'; -import { memo, useCallback, useMemo } from 'react'; +import { memo, useCallback } from 'react'; import { useTranslation } from 'react-i18next'; import { PiDownloadSimpleBold } from 'react-icons/pi'; import type { AnyModelConfigWithExternal } from 'services/api/types'; +import { buildExportData, fetchImageAsDataUrl, sanitizeFilename } from './modelSettingsIO'; + type Props = { modelConfig: AnyModelConfigWithExternal; }; -const buildExportData = (modelConfig: AnyModelConfigWithExternal): Record => { - const data: Record = {}; - - if ( - 'default_settings' in modelConfig && - modelConfig.default_settings !== undefined && - modelConfig.default_settings !== null - ) { - data.default_settings = modelConfig.default_settings; - } - - if ( - 'trigger_phrases' in modelConfig && - modelConfig.trigger_phrases !== undefined && - modelConfig.trigger_phrases !== null - ) { - data.trigger_phrases = modelConfig.trigger_phrases; - } - - if ('cpu_only' in modelConfig && modelConfig.cpu_only !== null) { - data.cpu_only = modelConfig.cpu_only; - } - - return data; -}; - -const sanitizeFilename = (name: string): string => { - return name.replace(/[<>:"/\\|?*]/g, '_'); -}; - export const ModelSettingsExportButton = memo(({ modelConfig }: Props) => { const { t } = useTranslation(); - const hasExportableData = useMemo(() => Object.keys(buildExportData(modelConfig)).length > 0, [modelConfig]); - - const handleExport = useCallback(() => { + const handleExport = useCallback(async () => { const data = buildExportData(modelConfig); + + if ( + 'cover_image' in modelConfig && + typeof modelConfig.cover_image === 'string' && + modelConfig.cover_image.length > 0 + ) { + const dataUrl = await fetchImageAsDataUrl(modelConfig.cover_image); + if (dataUrl) { + data.cover_image = dataUrl; + } + } + const json = JSON.stringify(data, null, 2); const blob = new Blob([json], { type: 'application/json' }); const url = URL.createObjectURL(blob); @@ -73,7 +55,6 @@ export const ModelSettingsExportButton = memo(({ modelConfig }: Props) => { aria-label={t('modelManager.exportSettings')} tooltip={t('modelManager.exportSettings')} onClick={handleExport} - isDisabled={!hasExportableData} /> ); }); diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsImportButton.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsImportButton.tsx index eeaaabbe730..2353b21f782 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsImportButton.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelSettingsImportButton.tsx @@ -4,40 +4,10 @@ import type { ChangeEvent } from 'react'; import { memo, useCallback, useRef } from 'react'; import { useTranslation } from 'react-i18next'; import { PiUploadSimpleBold } from 'react-icons/pi'; -import { useUpdateModelMutation } from 'services/api/endpoints/models'; +import { useUpdateModelImageMutation, useUpdateModelMutation } from 'services/api/endpoints/models'; import type { AnyModelConfigWithExternal } from 'services/api/types'; -const validateImportData = (data: unknown): data is Record => { - if (typeof data !== 'object' || data === null || Array.isArray(data)) { - return false; - } - - const obj = data as Record; - - if ('trigger_phrases' in obj && obj.trigger_phrases !== undefined) { - if (!Array.isArray(obj.trigger_phrases) || !obj.trigger_phrases.every((p) => typeof p === 'string')) { - return false; - } - } - - if ('default_settings' in obj && obj.default_settings !== undefined) { - if ( - typeof obj.default_settings !== 'object' || - obj.default_settings === null || - Array.isArray(obj.default_settings) - ) { - return false; - } - } - - if ('cpu_only' in obj && obj.cpu_only !== undefined) { - if (typeof obj.cpu_only !== 'boolean') { - return false; - } - } - - return true; -}; +import { dataUrlToFile, isImageDataUrl, validateImportData } from './modelSettingsIO'; type Props = { modelConfig: AnyModelConfigWithExternal; @@ -47,13 +17,21 @@ export const ModelSettingsImportButton = memo(({ modelConfig }: Props) => { const { t } = useTranslation(); const fileInputRef = useRef(null); const [updateModel] = useUpdateModelMutation(); + const [updateModelImage] = useUpdateModelImageMutation(); const applySettings = useCallback( async (data: Record) => { const body: Record = {}; const skippedFields: string[] = []; - const importableFields = ['default_settings', 'trigger_phrases', 'cpu_only'] as const; + const importableFields = [ + 'name', + 'description', + 'source_url', + 'default_settings', + 'trigger_phrases', + 'cpu_only', + ] as const; for (const field of importableFields) { if (!(field in data) || data[field] === undefined || data[field] === null) { @@ -66,7 +44,12 @@ export const ModelSettingsImportButton = memo(({ modelConfig }: Props) => { } } - if (Object.keys(body).length === 0) { + const coverImageDataUrl = + 'cover_image' in data && typeof data.cover_image === 'string' && isImageDataUrl(data.cover_image) + ? data.cover_image + : null; + + if (Object.keys(body).length === 0 && !coverImageDataUrl) { if (skippedFields.length > 0) { toast({ id: 'SETTINGS_IMPORT_INCOMPATIBLE', @@ -77,35 +60,62 @@ export const ModelSettingsImportButton = memo(({ modelConfig }: Props) => { return; } - await updateModel({ - key: modelConfig.key, - body, - }) - .unwrap() - .then(() => { - if (skippedFields.length > 0) { - toast({ - id: 'SETTINGS_IMPORTED', - title: t('modelManager.settingsImportedPartial', { fields: skippedFields.join(', ') }), - status: 'warning', - }); - } else { - toast({ - id: 'SETTINGS_IMPORTED', - title: t('modelManager.settingsImported'), - status: 'success', - }); - } - }) - .catch((_error) => { + let appliedAnything = false; + if (Object.keys(body).length > 0) { + try { + await updateModel({ + key: modelConfig.key, + body, + }).unwrap(); + appliedAnything = true; + } catch { toast({ id: 'SETTINGS_IMPORT_FAILED', title: t('modelManager.settingsImportFailed'), status: 'error', }); + return; + } + } + + if (coverImageDataUrl) { + const imageFile = dataUrlToFile(coverImageDataUrl, `${modelConfig.key}.png`); + if (!imageFile) { + skippedFields.push('cover_image'); + } else { + try { + await updateModelImage({ key: modelConfig.key, image: imageFile }).unwrap(); + appliedAnything = true; + } catch { + skippedFields.push('cover_image'); + } + } + } + + if (!appliedAnything) { + toast({ + id: 'SETTINGS_IMPORT_FAILED', + title: t('modelManager.settingsImportFailed'), + status: 'error', }); + return; + } + + if (skippedFields.length > 0) { + toast({ + id: 'SETTINGS_IMPORTED', + title: t('modelManager.settingsImportedPartial', { fields: skippedFields.join(', ') }), + status: 'warning', + }); + } else { + toast({ + id: 'SETTINGS_IMPORTED', + title: t('modelManager.settingsImported'), + status: 'success', + }); + } }, - [modelConfig, updateModel, t] + [modelConfig, updateModel, updateModelImage, t] ); const handleFileChange = useCallback( diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/modelSettingsIO.test.ts b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/modelSettingsIO.test.ts new file mode 100644 index 00000000000..ca3db90366d --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/modelSettingsIO.test.ts @@ -0,0 +1,262 @@ +import type { AnyModelConfigWithExternal } from 'services/api/types'; +import { describe, expect, it } from 'vitest'; + +import { + buildExportData, + dataUrlToFile, + isImageDataUrl, + isSafeUrl, + sanitizeFilename, + validateImportData, +} from './modelSettingsIO'; + +const makeConfig = (overrides: Record = {}): AnyModelConfigWithExternal => { + return { + key: 'test-key', + hash: 'abc123', + path: '/models/test.safetensors', + file_size: 0, + name: 'Test Model', + description: null, + source: '/models/test.safetensors', + source_type: 'path', + source_api_response: null, + source_url: null, + cover_image: null, + base: 'sd-1', + type: 'main', + format: 'checkpoint', + default_settings: null, + trigger_phrases: null, + cpu_only: null, + ...overrides, + } as unknown as AnyModelConfigWithExternal; +}; + +describe('sanitizeFilename', () => { + it('replaces filesystem-unsafe characters with underscores', () => { + expect(sanitizeFilename('foobaz')).toBe('foo_bar_baz'); + expect(sanitizeFilename('a/b\\c:d|e?f*g"h')).toBe('a_b_c_d_e_f_g_h'); + }); + + it('leaves safe filenames untouched', () => { + expect(sanitizeFilename('My Model v2.1')).toBe('My Model v2.1'); + }); +}); + +describe('isSafeUrl', () => { + it('accepts http and https URLs', () => { + expect(isSafeUrl('https://example.com')).toBe(true); + expect(isSafeUrl('http://example.com')).toBe(true); + }); + + it('rejects other schemes', () => { + expect(isSafeUrl('javascript:alert(1)')).toBe(false); + expect(isSafeUrl('data:text/html,foo')).toBe(false); + expect(isSafeUrl('ftp://example.com')).toBe(false); + expect(isSafeUrl('example.com')).toBe(false); + expect(isSafeUrl('')).toBe(false); + }); +}); + +describe('isImageDataUrl', () => { + it('accepts data URLs with image MIME types', () => { + expect(isImageDataUrl('data:image/png;base64,iVBORw0KGgo=')).toBe(true); + expect(isImageDataUrl('data:image/webp;base64,UklGRg==')).toBe(true); + expect(isImageDataUrl('data:image/jpeg;base64,/9j/4AAQ')).toBe(true); + expect(isImageDataUrl('data:image/svg+xml;base64,PHN2Zw==')).toBe(true); + }); + + it('rejects non-image data URLs', () => { + expect(isImageDataUrl('data:text/plain;base64,aGVsbG8=')).toBe(false); + expect(isImageDataUrl('data:application/json;base64,e30=')).toBe(false); + }); + + it('rejects non-base64 image data URLs', () => { + expect(isImageDataUrl('data:image/png,iVBORw0KGgo=')).toBe(false); + }); + + it('rejects malformed inputs', () => { + expect(isImageDataUrl('https://example.com/img.png')).toBe(false); + expect(isImageDataUrl('not a url at all')).toBe(false); + expect(isImageDataUrl('')).toBe(false); + }); +}); + +describe('dataUrlToFile', () => { + it('decodes a valid base64 image data URL into a File', () => { + // 1x1 transparent PNG + const dataUrl = + 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII='; + const file = dataUrlToFile(dataUrl, 'thumb.png'); + expect(file).not.toBeNull(); + expect(file?.name).toBe('thumb.png'); + expect(file?.type).toBe('image/png'); + expect(file?.size).toBeGreaterThan(0); + }); + + it('returns null when the prefix is missing', () => { + expect(dataUrlToFile('iVBORw0KGgo=', 'thumb.png')).toBeNull(); + }); + + it('returns null when the data URL has no base64 payload', () => { + expect(dataUrlToFile('data:image/png;base64,', 'thumb.png')).toBeNull(); + }); + + it('returns null on invalid base64 content', () => { + expect(dataUrlToFile('data:image/png;base64,!!!not-base64!!!', 'thumb.png')).toBeNull(); + }); +}); + +describe('buildExportData', () => { + it('returns an empty object when no exportable fields are set', () => { + const config = makeConfig({ name: '' }); + expect(buildExportData(config)).toEqual({}); + }); + + it('includes name, description, and source_url when present', () => { + const config = makeConfig({ + name: 'My LoRA', + description: 'A cool LoRA', + source_url: 'https://civitai.com/models/12345', + }); + expect(buildExportData(config)).toEqual({ + name: 'My LoRA', + description: 'A cool LoRA', + source_url: 'https://civitai.com/models/12345', + }); + }); + + it('omits empty string metadata fields', () => { + const config = makeConfig({ + name: 'My LoRA', + description: '', + source_url: '', + }); + expect(buildExportData(config)).toEqual({ name: 'My LoRA' }); + }); + + it('includes trigger_phrases, default_settings, and cpu_only when present', () => { + const triggerPhrases = ['foo', 'bar']; + const defaultSettings = { steps: 30 }; + const config = makeConfig({ + name: '', + trigger_phrases: triggerPhrases, + default_settings: defaultSettings, + cpu_only: true, + }); + expect(buildExportData(config)).toEqual({ + trigger_phrases: triggerPhrases, + default_settings: defaultSettings, + cpu_only: true, + }); + }); + + it('includes cpu_only: false (only null is omitted)', () => { + const config = makeConfig({ + name: '', + cpu_only: false, + }); + expect(buildExportData(config)).toEqual({ cpu_only: false }); + }); + + it('does not include cover_image (handled separately, async)', () => { + const config = makeConfig({ + name: 'X', + cover_image: 'https://example.com/img.png', + }); + expect(buildExportData(config)).toEqual({ name: 'X' }); + }); +}); + +describe('validateImportData', () => { + it('accepts an empty object', () => { + expect(validateImportData({})).toBe(true); + }); + + it('rejects non-object inputs', () => { + expect(validateImportData(null)).toBe(false); + expect(validateImportData(undefined)).toBe(false); + expect(validateImportData('string')).toBe(false); + expect(validateImportData(42)).toBe(false); + expect(validateImportData([])).toBe(false); + }); + + it('accepts valid metadata fields', () => { + expect( + validateImportData({ + name: 'My Model', + description: 'a model', + source_url: 'https://example.com', + }) + ).toBe(true); + }); + + it('accepts null metadata fields', () => { + expect(validateImportData({ name: null, description: null, source_url: null })).toBe(true); + }); + + it('rejects non-string name and description', () => { + expect(validateImportData({ name: 123 })).toBe(false); + expect(validateImportData({ description: { not: 'a string' } })).toBe(false); + }); + + it('rejects source_url that is not http(s)', () => { + expect(validateImportData({ source_url: 'javascript:alert(1)' })).toBe(false); + expect(validateImportData({ source_url: 'ftp://example.com' })).toBe(false); + expect(validateImportData({ source_url: 'example.com' })).toBe(false); + }); + + it('accepts an empty source_url string', () => { + expect(validateImportData({ source_url: '' })).toBe(true); + }); + + it('accepts a valid image data URL for cover_image', () => { + expect(validateImportData({ cover_image: 'data:image/png;base64,iVBORw0KGgo=' })).toBe(true); + }); + + it('rejects non-image cover_image values', () => { + expect(validateImportData({ cover_image: 'https://example.com/img.png' })).toBe(false); + expect(validateImportData({ cover_image: 'data:text/plain;base64,aGk=' })).toBe(false); + expect(validateImportData({ cover_image: 42 })).toBe(false); + }); + + it('validates trigger_phrases as an array of strings', () => { + expect(validateImportData({ trigger_phrases: ['a', 'b'] })).toBe(true); + expect(validateImportData({ trigger_phrases: [] })).toBe(true); + expect(validateImportData({ trigger_phrases: ['a', 1] })).toBe(false); + expect(validateImportData({ trigger_phrases: 'not-an-array' })).toBe(false); + }); + + it('validates default_settings as a plain object', () => { + expect(validateImportData({ default_settings: { steps: 30 } })).toBe(true); + expect(validateImportData({ default_settings: {} })).toBe(true); + expect(validateImportData({ default_settings: [] })).toBe(false); + expect(validateImportData({ default_settings: 'nope' })).toBe(false); + }); + + it('validates cpu_only as a boolean', () => { + expect(validateImportData({ cpu_only: true })).toBe(true); + expect(validateImportData({ cpu_only: false })).toBe(true); + expect(validateImportData({ cpu_only: 'true' })).toBe(false); + expect(validateImportData({ cpu_only: 1 })).toBe(false); + }); + + it('accepts a fully populated valid export', () => { + expect( + validateImportData({ + name: 'My Model', + description: 'desc', + source_url: 'https://civitai.com/models/1', + cover_image: 'data:image/webp;base64,UklGRg==', + trigger_phrases: ['trigger'], + default_settings: { steps: 30 }, + cpu_only: false, + }) + ).toBe(true); + }); + + it('ignores unknown fields', () => { + expect(validateImportData({ name: 'X', someUnknownField: 'whatever' })).toBe(true); + }); +}); diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/modelSettingsIO.ts b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/modelSettingsIO.ts new file mode 100644 index 00000000000..cf45f517911 --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/modelSettingsIO.ts @@ -0,0 +1,157 @@ +import type { AnyModelConfigWithExternal } from 'services/api/types'; + +export const sanitizeFilename = (name: string): string => { + return name.replace(/[<>:"/\\|?*]/g, '_'); +}; + +export const isSafeUrl = (url: string): boolean => { + return url.startsWith('https://') || url.startsWith('http://'); +}; + +export const isImageDataUrl = (value: string): boolean => { + return /^data:image\/[a-zA-Z0-9.+-]+;base64,/.test(value); +}; + +export const dataUrlToFile = (dataUrl: string, filename: string): File | null => { + const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/); + if (!match) { + return null; + } + const mime = match[1]; + const b64 = match[2]; + if (!mime || !b64) { + return null; + } + try { + const binary = atob(b64); + const bytes = new Uint8Array(binary.length); + for (let i = 0; i < binary.length; i++) { + bytes[i] = binary.charCodeAt(i); + } + return new File([bytes], filename, { type: mime }); + } catch { + return null; + } +}; + +export const buildExportData = (modelConfig: AnyModelConfigWithExternal): Record => { + const data: Record = {}; + + if ('name' in modelConfig && typeof modelConfig.name === 'string' && modelConfig.name.length > 0) { + data.name = modelConfig.name; + } + + if ( + 'description' in modelConfig && + typeof modelConfig.description === 'string' && + modelConfig.description.length > 0 + ) { + data.description = modelConfig.description; + } + + if ('source_url' in modelConfig && typeof modelConfig.source_url === 'string' && modelConfig.source_url.length > 0) { + data.source_url = modelConfig.source_url; + } + + if ( + 'default_settings' in modelConfig && + modelConfig.default_settings !== undefined && + modelConfig.default_settings !== null + ) { + data.default_settings = modelConfig.default_settings; + } + + if ( + 'trigger_phrases' in modelConfig && + modelConfig.trigger_phrases !== undefined && + modelConfig.trigger_phrases !== null + ) { + data.trigger_phrases = modelConfig.trigger_phrases; + } + + if ('cpu_only' in modelConfig && modelConfig.cpu_only !== null) { + data.cpu_only = modelConfig.cpu_only; + } + + return data; +}; + +export const fetchImageAsDataUrl = async (url: string): Promise => { + try { + const response = await fetch(url); + if (!response.ok) { + return null; + } + const blob = await response.blob(); + if (!blob.type.startsWith('image/')) { + return null; + } + return await new Promise((resolve) => { + const reader = new FileReader(); + reader.onload = () => resolve(typeof reader.result === 'string' ? reader.result : null); + reader.onerror = () => resolve(null); + reader.readAsDataURL(blob); + }); + } catch { + return null; + } +}; + +export const validateImportData = (data: unknown): data is Record => { + if (typeof data !== 'object' || data === null || Array.isArray(data)) { + return false; + } + + const obj = data as Record; + + if ('name' in obj && obj.name !== undefined && obj.name !== null) { + if (typeof obj.name !== 'string') { + return false; + } + } + + if ('description' in obj && obj.description !== undefined && obj.description !== null) { + if (typeof obj.description !== 'string') { + return false; + } + } + + if ('source_url' in obj && obj.source_url !== undefined && obj.source_url !== null) { + if (typeof obj.source_url !== 'string') { + return false; + } + if (obj.source_url.length > 0 && !isSafeUrl(obj.source_url)) { + return false; + } + } + + if ('cover_image' in obj && obj.cover_image !== undefined && obj.cover_image !== null) { + if (typeof obj.cover_image !== 'string' || !isImageDataUrl(obj.cover_image)) { + return false; + } + } + + if ('trigger_phrases' in obj && obj.trigger_phrases !== undefined) { + if (!Array.isArray(obj.trigger_phrases) || !obj.trigger_phrases.every((p) => typeof p === 'string')) { + return false; + } + } + + if ('default_settings' in obj && obj.default_settings !== undefined) { + if ( + typeof obj.default_settings !== 'object' || + obj.default_settings === null || + Array.isArray(obj.default_settings) + ) { + return false; + } + } + + if ('cpu_only' in obj && obj.cpu_only !== undefined) { + if (typeof obj.cpu_only !== 'boolean') { + return false; + } + } + + return true; +};