diff --git a/__init__.py b/__init__.py index 13190c1..c5a0770 100644 --- a/__init__.py +++ b/__init__.py @@ -231,8 +231,8 @@ async def download_model_info(request): """ post = await utils.get_request_body(request) try: - scan_mode = post.get("scanMode", "supplement") - services.download_model_info(scan_mode) + scan_mode = post.get("scanMode", "diff") + await services.download_model_info(scan_mode) return web.json_response({"success": True}) except Exception as e: error_msg = f"Download model info failed: {str(e)}" @@ -282,6 +282,20 @@ async def read_download_preview(request): return web.FileResponse(preview_path) +@routes.post("/model-manager/migrate") +async def migrate_legacy_information(request): + """ + Migrate legacy information. + """ + try: + await services.migrate_legacy_information() + return web.json_response({"success": True}) + except Exception as e: + error_msg = f"Download model info failed: {str(e)}" + utils.print_error(error_msg) + return web.json_response({"success": False, "error": error_msg}) + + WEB_DIRECTORY = "web" NODE_CLASS_MAPPINGS = {} __all__ = ["WEB_DIRECTORY", "NODE_CLASS_MAPPINGS"] diff --git a/py/services.py b/py/services.py index 82ba020..92a6cbe 100644 --- a/py/services.py +++ b/py/services.py @@ -6,7 +6,6 @@ from . import utils from . import download from . import searcher -from . import thread def scan_models(): @@ -141,8 +140,8 @@ def fetch_model_info(model_page: str): return result -def try_to_scan_model_and_download_information(scan_mode: str): - utils.print_debug(f"scan_mode: {scan_mode}") +async def download_model_info(scan_mode: str): + utils.print_debug(f"Download model info for {scan_mode}") model_base_paths = config.model_base_paths for model_type in model_base_paths: @@ -217,8 +216,89 @@ def try_to_scan_model_and_download_information(scan_mode: str): utils.print_debug("Completed scan model information.") -def download_model_info(scan_mode: str): - utils.print_debug(f"Download model info for {scan_mode}") - thread.start_threading( - try_to_scan_model_and_download_information, args=(scan_mode,) - ) +async def migrate_legacy_information(): + import json + import yaml + from PIL import Image + + model_base_paths = config.model_base_paths + for model_type in model_base_paths: + + folders, extensions = folder_paths.folder_names_and_paths[model_type] + for path_index, base_path in enumerate(folders): + files = utils.recursive_search_files(base_path) + + models = folder_paths.filter_files_extensions(files, extensions) + + for fullname in models: + fullname = utils.normalize_path(fullname) + + abs_model_path = utils.join_path(base_path, fullname) + + base_file_name = os.path.splitext(abs_model_path)[0] + + utils.print_debug(f"migrate legacy info for {abs_model_path}") + + preview_path = utils.join_path( + os.path.dirname(abs_model_path), + utils.get_model_preview_name(abs_model_path), + ) + new_preview_path = f"{base_file_name}.webp" + + if os.path.isfile(preview_path) and preview_path != new_preview_path: + with Image.open(preview_path) as image: + image.save(new_preview_path, format="WEBP") + os.remove(preview_path) + + description_path = f"{base_file_name}.md" + + metadata_info = { + "website": "Civitai", + } + + url_info_path = f"{base_file_name}.url" + if os.path.isfile(url_info_path): + with open(url_info_path, "r", encoding="utf-8") as f: + for line in f: + if line.startswith("URL="): + model_page_url = line[len("URL=") :].strip() + metadata_info.update({"modelPage": model_page_url}) + + json_info_path = f"{base_file_name}.json" + if os.path.isfile(json_info_path): + with open(json_info_path, "r", encoding="utf-8") as f: + version = json.load(f) + metadata_info.update( + { + "baseModel": version.get("baseModel"), + "preview": [i["url"] for i in version["images"]], + } + ) + + description_parts: list[str] = [ + "---", + yaml.dump(metadata_info).strip(), + "---", + "", + ] + + text_info_path = f"{base_file_name}.txt" + if os.path.isfile(text_info_path): + with open(text_info_path, "r", encoding="utf-8") as f: + description_parts.append(f.read()) + + description_path = f"{base_file_name}.md" + + if os.path.isfile(text_info_path): + with open(description_path, "w", encoding="utf-8", newline="") as f: + f.write("\n".join(description_parts)) + + def try_to_remove_file(file_path): + if os.path.isfile(file_path): + os.remove(file_path) + + try_to_remove_file(url_info_path) + try_to_remove_file(text_info_path) + try_to_remove_file(json_info_path) + + utils.print_debug("Completed migrate model information.") diff --git a/py/thread.py b/py/thread.py index d6506a8..cabeb54 100644 --- a/py/thread.py +++ b/py/thread.py @@ -55,17 +55,3 @@ def _worker(self): with self._lock: self.workers_count -= 1 - - -def start_threading(target, *, group=None, name=None, args=(), kwargs=None): - utils.print_debug(f"args {args}") - t = threading.Thread( - target=target, - group=group, - name=name, - args=args, - kwargs=kwargs, - daemon=True, - ) - t.start() - return t diff --git a/src/hooks/config.ts b/src/hooks/config.ts index 358b148..0674d3d 100644 --- a/src/hooks/config.ts +++ b/src/hooks/config.ts @@ -1,9 +1,10 @@ -import { useRequest } from 'hooks/request' +import { request, useRequest } from 'hooks/request' import { defineStore } from 'hooks/store' -import { app } from 'scripts/comfyAPI' +import { $el, app, ComfyDialog } from 'scripts/comfyAPI' import { onMounted, onUnmounted, ref } from 'vue' +import { useToast } from './toast' -export const useConfig = defineStore('config', () => { +export const useConfig = defineStore('config', (store) => { const mobileDeviceBreakPoint = 759 const isMobile = ref(window.innerWidth < mobileDeviceBreakPoint) @@ -36,7 +37,7 @@ export const useConfig = defineStore('config', () => { refresh, } - useAddConfigSettings() + useAddConfigSettings(store) return config }) @@ -49,7 +50,41 @@ declare module 'hooks/store' { } } -function useAddConfigSettings() { +function useAddConfigSettings(store: import('hooks/store').StoreProvider) { + const { toast } = useToast() + + const confirm = (opts: { + message?: string + accept?: () => void + reject?: () => void + }) => { + const dialog = new ComfyDialog('div', []) + + dialog.show( + $el('div', [ + $el('p', { textContent: opts.message }), + $el('div.flex.gap-4', [ + $el('button.flex-1', { + textContent: 'Cancel', + onclick: () => { + opts.reject?.() + dialog.close() + document.body.removeChild(dialog.element) + }, + }), + $el('button.flex-1', { + textContent: 'Continue', + onclick: () => { + opts.accept?.() + dialog.close() + document.body.removeChild(dialog.element) + }, + }), + ]), + ]), + ) + } + onMounted(() => { // API keys app.ui?.settings.addSetting({ @@ -65,5 +100,144 @@ function useAddConfigSettings() { type: 'text', defaultValue: undefined, }) + + // Migrate + app.ui?.settings.addSetting({ + id: 'ModelManager.Migrate.Migrate', + name: 'Migrate information from cdb-boop/main', + defaultValue: '', + type: () => { + return $el('button.p-button.p-component.p-button-secondary', { + textContent: 'Migrate', + onclick: () => { + confirm({ + message: [ + 'This operation will delete old files and override current files if it exists.', + // 'This may take a while and generate MANY server requests!', + 'Continue?', + ].join('\n'), + accept: () => { + store.loading.loading.value = true + request('/migrate', { + method: 'POST', + }) + .then(() => { + toast.add({ + severity: 'success', + summary: 'Complete migration', + life: 2000, + }) + store.models.refresh() + }) + .catch((err) => { + toast.add({ + severity: 'error', + summary: 'Error', + detail: err.message ?? 'Failed to migrate information', + life: 15000, + }) + }) + .finally(() => { + store.loading.loading.value = false + }) + }, + }) + }, + }) + }, + }) + + // Scan information + app.ui?.settings.addSetting({ + id: 'ModelManager.ScanFiles.Full', + name: "Override all models' information and preview", + defaultValue: '', + type: () => { + return $el('button.p-button.p-component.p-button-secondary', { + textContent: 'Full Scan', + onclick: () => { + confirm({ + message: [ + 'This operation will override current files.', + 'This may take a while and generate MANY server requests!', + 'USE AT YOUR OWN RISK! Continue?', + ].join('\n'), + accept: () => { + store.loading.loading.value = true + request('/model-info/scan', { + method: 'POST', + body: JSON.stringify({ scanMode: 'full' }), + }) + .then(() => { + toast.add({ + severity: 'success', + summary: 'Complete download information', + life: 2000, + }) + store.models.refresh() + }) + .catch((err) => { + toast.add({ + severity: 'error', + summary: 'Error', + detail: err.message ?? 'Failed to download information', + life: 15000, + }) + }) + .finally(() => { + store.loading.loading.value = false + }) + }, + }) + }, + }) + }, + }) + + app.ui?.settings.addSetting({ + id: 'ModelManager.ScanFiles.Incremental', + name: 'Download missing information or preview', + defaultValue: '', + type: () => { + return $el('button.p-button.p-component.p-button-secondary', { + textContent: 'Diff Scan', + onclick: () => { + confirm({ + message: [ + 'Download missing information or preview.', + 'This may take a while and generate MANY server requests!', + 'USE AT YOUR OWN RISK! Continue?', + ].join('\n'), + accept: () => { + store.loading.loading.value = true + request('/model-info/scan', { + method: 'POST', + body: JSON.stringify({ scanMode: 'diff' }), + }) + .then(() => { + toast.add({ + severity: 'success', + summary: 'Complete download information', + life: 2000, + }) + store.models.refresh() + }) + .catch((err) => { + toast.add({ + severity: 'error', + summary: 'Error', + detail: err.message ?? 'Failed to download information', + life: 15000, + }) + }) + .finally(() => { + store.loading.loading.value = false + }) + }, + }) + }, + }) + }, + }) }) } diff --git a/src/hooks/loading.ts b/src/hooks/loading.ts index 349b7d2..2a66005 100644 --- a/src/hooks/loading.ts +++ b/src/hooks/loading.ts @@ -31,6 +31,12 @@ export const useGlobalLoading = defineStore('loading', () => { return { loading } }) +declare module 'hooks/store' { + interface StoreProvider { + loading: ReturnType + } +} + export const useLoading = () => { const timer = ref() diff --git a/src/scripts/comfyAPI.ts b/src/scripts/comfyAPI.ts index 61394ff..1f5445b 100644 --- a/src/scripts/comfyAPI.ts +++ b/src/scripts/comfyAPI.ts @@ -5,3 +5,4 @@ export const $el = window.comfyAPI.ui.$el export const ComfyApp = window.comfyAPI.app.ComfyApp export const ComfyButton = window.comfyAPI.button.ComfyButton +export const ComfyDialog = window.comfyAPI.dialog.ComfyDialog diff --git a/src/types/global.d.ts b/src/types/global.d.ts index 9a4e612..f3fb9ee 100644 --- a/src/types/global.d.ts +++ b/src/types/global.d.ts @@ -112,6 +112,7 @@ declare namespace ComfyAPI { settings: ComfySettingsDialog menuHamburger?: HTMLDivElement menuContainer?: HTMLDivElement + dialog: dialog.ComfyDialog } type SettingInputType = @@ -197,6 +198,15 @@ declare namespace ComfyAPI { constructor(...buttons: (HTMLElement | ComfyButton)[]): ComfyButtonGroup } } + + namespace dialog { + class ComfyDialog { + constructor(type = 'div', buttons: HTMLElement[] = null) + element: HTMLElement + close(): void + show(html: string | HTMLElement | HTMLElement[]): void + } + } } declare namespace lightGraph {