From 4564d1c69ca66041ca3c6aa8350483ca325307f9 Mon Sep 17 00:00:00 2001 From: Arslan Ablikim Date: Thu, 18 Jan 2024 00:59:57 +0800 Subject: [PATCH 1/2] update queue to show queue in the modal --- service/model_manager/model_installer.py | 86 ++++++++++--------- .../install-models/InatallModelsModal.tsx | 46 ++++++---- 2 files changed, 75 insertions(+), 57 deletions(-) diff --git a/service/model_manager/model_installer.py b/service/model_manager/model_installer.py index ec31cc8..55ff50e 100644 --- a/service/model_manager/model_installer.py +++ b/service/model_manager/model_installer.py @@ -6,9 +6,9 @@ import sys from aiohttp import web import folder_paths import urllib.request +import threading import server -import server -from .missing_models import find_missing_models +import time comfy_path = os.path.dirname(folder_paths.__file__) def download_url_with_wget(url, save_path): @@ -56,32 +56,29 @@ def parse_wget_output(output): match = re.search(r'\d+%|\d+K \d+Kb/s', output) return match.group(0) if match else "" -import queue -import threading - -# Create a Queue object -download_queue = queue.Queue() +# Create a dictionary to store download tasks +download_tasks = [] +# Create a lock object to synchronize access to the dictionary +download_tasks_lock = threading.Lock() # Create a thread that will execute the download tasks download_thread = None -# Create a Condition object to share the download progress -download_progress = threading.Condition() -download_progress.value = {'progress': "", 'error': None} def download_worker(): - while not download_queue.empty(): - # Get a download task from the queue - (url, save_path) = download_queue.get() + while True: + task = None + with download_tasks_lock: + # Get a task from the list + if download_tasks: + task = download_tasks.pop(0) - # Execute the download task and update the download progress - download_url_with_agent(url, save_path) - - # Indicate that the task is done - download_queue.task_done() - - # If all tasks are done, reset the progress and exit - with download_progress: - download_progress.value['progress'] = "" - download_progress.notify_all() + if task is not None: + # Execute the download task and update the download progress + download_url_with_agent(task['url'], task['save_path']) + + else: + # No more tasks, break the loop + send_ws('download_progress', []) + break @server.PromptServer.instance.routes.post("/model_manager/install_model") async def install_model(request): @@ -91,7 +88,9 @@ async def install_model(request): url = json_data['url'] save_path = get_model_path(json_data) - download_queue.put((url, save_path)) + with download_tasks_lock: + # Add the task to the list + download_tasks.append({ 'url': url, 'save_path': save_path }) # If the previous thread is not active, start a new one if download_thread is None or not download_thread.is_alive(): @@ -101,16 +100,6 @@ async def install_model(request): return web.Response(text=f"Downloading {url} to {save_path} ...") -@server.PromptServer.instance.routes.get("/model_manager/download_progress") -def get_download_progress(request): - with download_progress: - if download_progress.value['error'] is not None: - error = download_progress.value['error'] - download_progress.value['error'] = None - return web.Response(text=error) - else: - return web.Response(text=download_progress.value['progress']) - @server.PromptServer.instance.routes.post("/model_manager/install_model_stream") async def install_model_stream(request): json_data = await request.json() @@ -170,9 +159,7 @@ def download_url_with_agent(url, save_path, progress_callback=None): downloaded += len(chunk) progress = (downloaded / file_size) * 100 print(f'\rProgress: {progress:.2f}%', end='') - with download_progress: - download_progress.value['progress'] = f'Downloading {os.path.basename(save_path)}: {progress:.2f}%' # Update the download progress - download_progress.notify_all() + send_download_status({'save_path': save_path, 'progress': progress}) # if progress_callback: # progress_callback(progress) @@ -182,9 +169,7 @@ def download_url_with_agent(url, save_path, progress_callback=None): except Exception as e: print(f"\nDownload error: {url} / {e}", file=sys.stderr) - with download_progress: - download_progress.value['error'] = f"\nDownload error: {url} / {e}" # Update the download progress - download_progress.notify_all() + send_ws('download_error', f"{url} / {e}") if os.path.exists(temp_save_path): os.remove(temp_save_path) # Clean up the temporary file in case of failure return False @@ -234,4 +219,23 @@ def get_model_dir(data): def get_model_path(data): base_model = get_model_dir(data) - return os.path.join(base_model, data['filename']) \ No newline at end of file + return os.path.join(base_model, data['filename']) + + +last_call_time = 0 +def send_download_status(data): + global last_call_time + current_time = time.time() + if current_time - last_call_time < 0.1: + return + last_call_time = time.time() + with download_tasks_lock: + progress_list = [data] + for task in download_tasks: + progress_list.append({'save_path': task['save_path'], 'progress': 0}) + send_ws('download_progress', progress_list) + +loop = asyncio.get_event_loop() +def send_ws(event, data): + asyncio.run_coroutine_threadsafe(server.PromptServer.instance.send(event, data) , loop) + \ No newline at end of file diff --git a/ui/src/model-manager/install-models/InatallModelsModal.tsx b/ui/src/model-manager/install-models/InatallModelsModal.tsx index 2e5ba1c..953136c 100644 --- a/ui/src/model-manager/install-models/InatallModelsModal.tsx +++ b/ui/src/model-manager/install-models/InatallModelsModal.tsx @@ -14,15 +14,20 @@ import { Heading, Checkbox, Spinner, + useToast, + Progress, + Stack } from "@chakra-ui/react"; -import { useCallback, useEffect, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import { IconX } from "@tabler/icons-react"; import { CivitiModel, CivitiModelFileVersion } from "../types"; import { installModelsApi } from "../api/modelsApi"; import ModelCard from "./ModelCard"; import InstallProgress from "./InstallProgress"; import InstallModelSearchBar from "./InstallModelSearchBar"; -import { useToast } from "@chakra-ui/react"; + + +type Queue = { save_path: string, progress: number }; type CivitModelQueryParams = { types?: MODEL_TYPE; @@ -71,6 +76,7 @@ export default function InatallModelsModal({ const toast = useToast(); const [installing, setInstalling] = useState([]); const [searchQuery, setSearchQuery] = useState(""); + const [queue, setQueue] = useState([]); const loadData = useCallback(async () => { setLoading(true); const params: CivitModelQueryParams = { @@ -112,7 +118,7 @@ export default function InatallModelsModal({ } toast({ title: - "Installing...Please check the progress in your python server console", + "Installing...", description: file.name, status: "info", duration: 4000, @@ -131,20 +137,20 @@ export default function InatallModelsModal({ loadData(); }, [searchQuery, modelType]); - async function refreshHistory() { - const res = await api.fetchApi('/model_manager/download_progress'); - const ret = await res.text(); - console.log(ret); - window.dispatchEvent( - new CustomEvent("model_install_message", { - detail: ret, - }) - ); - } - useEffect(() => { - const autoRefreshInterval = setInterval(refreshHistory, 1000); - return () => clearInterval(autoRefreshInterval); + api.addEventListener("download_progress", (e: { detail: Queue[] }) => { + setQueue(e.detail); + }); + api.addEventListener("download_error", (e: { detail: string }) => { + toast({ + title: + "Download Error", + description: e.detail, + status: "error", + duration: 4000, + isClosable: true, + }); + }); }, []); const isAllSelected = @@ -223,6 +229,14 @@ export default function InatallModelsModal({ ); })} + + {queue.map(({ save_path, progress }) => ( + + {save_path.replace(/^.*[\\/]/, '')} + + + ))} + From 1292913b289354c458d09d9eee56a2607168173f Mon Sep 17 00:00:00 2001 From: Arslan Ablikim Date: Thu, 18 Jan 2024 02:35:33 +0800 Subject: [PATCH 2/2] add custom url download --- .../install-models/ChooseFolder.tsx | 64 +++++ .../install-models/InatallModelsModal.tsx | 233 ++++++++++-------- .../install-models/InstallProgress.tsx | 37 --- 3 files changed, 196 insertions(+), 138 deletions(-) create mode 100644 ui/src/model-manager/install-models/ChooseFolder.tsx delete mode 100644 ui/src/model-manager/install-models/InstallProgress.tsx diff --git a/ui/src/model-manager/install-models/ChooseFolder.tsx b/ui/src/model-manager/install-models/ChooseFolder.tsx new file mode 100644 index 0000000..250645d --- /dev/null +++ b/ui/src/model-manager/install-models/ChooseFolder.tsx @@ -0,0 +1,64 @@ +import { Button, AlertDialog, AlertDialogOverlay, AlertDialogContent, AlertDialogHeader, AlertDialogBody, AlertDialogFooter, Select } from "@chakra-ui/react"; +import { useRef, useState } from "react"; + +const allFolderPaths = [ + "checkpoints", + "configs", + "loras", + "vae", + "clip", + "unet", + "clip_vision", + "style_models", + "embeddings", + "diffusers", + "vae_approx", + "controlnet", + "gligen", + "upscale_models", + "custom_nodes", + "hypernetworks", +]; +interface ChooseFolderProps { + isOpen: boolean; + onClose: () => void; + selectFolder: (folderPath: string) => void; +} +export default function ChooseFolder({ isOpen, onClose, selectFolder }: ChooseFolderProps) { + const [folderPath, setFolderPath] = useState(""); + + const cancelRef = useRef(null); + + return ( + + + + + Choose Folder + + + + + + + + + + + + + + ); +} \ No newline at end of file diff --git a/ui/src/model-manager/install-models/InatallModelsModal.tsx b/ui/src/model-manager/install-models/InatallModelsModal.tsx index 953136c..cba3547 100644 --- a/ui/src/model-manager/install-models/InatallModelsModal.tsx +++ b/ui/src/model-manager/install-models/InatallModelsModal.tsx @@ -16,15 +16,16 @@ import { Spinner, useToast, Progress, - Stack + Stack, + useDisclosure } from "@chakra-ui/react"; import { useCallback, useEffect, useRef, useState } from "react"; import { IconX } from "@tabler/icons-react"; import { CivitiModel, CivitiModelFileVersion } from "../types"; import { installModelsApi } from "../api/modelsApi"; import ModelCard from "./ModelCard"; -import InstallProgress from "./InstallProgress"; import InstallModelSearchBar from "./InstallModelSearchBar"; +import ChooseFolder from "./ChooseFolder"; type Queue = { save_path: string, progress: number }; @@ -77,6 +78,8 @@ export default function InatallModelsModal({ const [installing, setInstalling] = useState([]); const [searchQuery, setSearchQuery] = useState(""); const [queue, setQueue] = useState([]); + const { isOpen, onOpen, onClose } = useDisclosure(); + const file = useRef(); const loadData = useCallback(async () => { setLoading(true); const params: CivitModelQueryParams = { @@ -98,40 +101,57 @@ export default function InatallModelsModal({ setModels(json.items); setLoading(false); }, [searchQuery, modelType]); - const onClickInstallModel = ( - file: CivitiModelFileVersion, - model: CivitiModel - ) => { - if (file.downloadUrl == null || file.name == null) { - console.error("file.downloadUrl or file.name is null"); + + const downloadModels = (folderPath: string) => { + if (!file.current?.downloadUrl) { + console.error("file.downloadUrl is null"); return; } - let folderPath: string | null = - MODEL_TYPE_TO_FOLDER_MAPPING[model.type as MODEL_TYPE]; - if (folderPath == null) { - folderPath = prompt( - "What's the folder path under /ComfyUI/models you want to save the model? " - ); - } - if (folderPath == null) { - return; + if (!file.current.name) { + file.current.name = file.current.downloadUrl.split("/").pop(); + if (!file.current.name) { + console.error("file.downloadUrl is malformed"); + return; + } } toast({ title: "Installing...", - description: file.name, + description: file.current.name, status: "info", duration: 4000, isClosable: true, }); - file.name != null && setInstalling((cur) => [...cur, file.name ?? ""]); + file.current.name != null && setInstalling((cur) => [...cur, file.current?.name ?? ""]); installModelsApi({ - filename: file.name, - name: file.name, + filename: file.current.name, + name: file.current.name, save_path: folderPath, - url: file.downloadUrl, + url: file.current.downloadUrl, }); + onClose(); }; + const onClickInstallModel = ( + _file: CivitiModelFileVersion, + model: CivitiModel + ) => { + let folderPath: string | null = + MODEL_TYPE_TO_FOLDER_MAPPING[model.type as MODEL_TYPE]; + file.current = _file; + if (folderPath == null) { + onOpen(); + } else { + downloadModels(folderPath); + } + }; + const customUrlDownload = () => { + const downloadUrl = prompt("Enter the URL to download"); + if (!downloadUrl) { + return; + } + file.current = ({ id: 0, downloadUrl }); + onOpen(); + } useEffect(() => { loadData(); @@ -157,88 +177,99 @@ export default function InatallModelsModal({ models.length > 0 && selectedID.length === models.length; return ( - - - - - - - Models - - - - - - - {ALL_MODEL_TYPES.map((type) => { - return ( - - ); - })} - - {isSelecting && ( - - All - {selectedID.length} Selected - + + + + + + + Models + + + - )} - {loading && ( - - )} - - - - - {models?.map((model) => { - return ( - + + {ALL_MODEL_TYPES.map((type) => { + return ( + + ); + })} + + {isSelecting && ( + + All + {selectedID.length} Selected + } + onClick={() => setIsSelecting(false)} + aria-label="cancel" /> - ); - })} - - - {queue.map(({ save_path, progress }) => ( - - {save_path.replace(/^.*[\\/]/, '')} - - ))} - - - - + )} + {loading && ( + + )} + + + + + {models?.map((model) => { + return ( + + ); + })} + + + {queue.map(({ save_path, progress }) => ( + + {save_path.replace(/^.*[\\/]/, '')} + + + ))} + + + + + + ); } + diff --git a/ui/src/model-manager/install-models/InstallProgress.tsx b/ui/src/model-manager/install-models/InstallProgress.tsx deleted file mode 100644 index f0d8d00..0000000 --- a/ui/src/model-manager/install-models/InstallProgress.tsx +++ /dev/null @@ -1,37 +0,0 @@ -import { useEffect, useState } from "react"; -import { Button, HStack, Text } from "@chakra-ui/react"; -import { setCancelInstall } from "../api/modelsApi"; - -export default function InstallProgress() { - const [progress, setProgress] = useState(""); - const handleModelInstallMessage = (event: Event) => { - const text = (event as CustomEvent).detail; - const lines = text.split("\n"); - const lastLine = lines.reverse().find((line: string) => line.trim() !== ""); - setProgress(lastLine ?? ""); - }; - - useEffect(() => { - // Attach the event listener - window.addEventListener("model_install_message", handleModelInstallMessage); - - // Cleanup function to remove the event listener - return () => { - window.removeEventListener( - "model_install_message", - handleModelInstallMessage - ); - }; - }, []); - - return ( - - - {progress} - - {/* {progress !== "" && ( - - )} */} - - ); -}