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/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 2e5ba1c..cba3547 100644
--- a/ui/src/model-manager/install-models/InatallModelsModal.tsx
+++ b/ui/src/model-manager/install-models/InatallModelsModal.tsx
@@ -14,15 +14,21 @@ import {
Heading,
Checkbox,
Spinner,
+ useToast,
+ Progress,
+ Stack,
+ useDisclosure
} 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";
+import ChooseFolder from "./ChooseFolder";
+
+
+type Queue = { save_path: string, progress: number };
type CivitModelQueryParams = {
types?: MODEL_TYPE;
@@ -71,6 +77,9 @@ export default function InatallModelsModal({
const toast = useToast();
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 = {
@@ -92,139 +101,175 @@ 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...Please check the progress in your python server console",
- description: file.name,
+ "Installing...",
+ 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();
}, [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 =
models.length > 0 && selectedID.length === models.length;
return (
-
-
-
-
-
-
- Models
-
-
-
-
-
-
- {ALL_MODEL_TYPES.map((type) => {
- return (
-
- );
- })}
-
- {isSelecting && (
-
- All
- {selectedID.length} Selected
-
+
+
+
+
+
+
+ Models
+
+
+ }
- onClick={() => setIsSelecting(false)}
- aria-label="cancel"
- />
+ py={1}
+ mr={8}
+ onClick={customUrlDownload}
+ >
+ Custom URL Install
+
- )}
- {loading && (
-
- )}
-
-
-
-
- {models?.map((model) => {
- return (
-
+
+ {ALL_MODEL_TYPES.map((type) => {
+ return (
+
+ );
+ })}
+
+ {isSelecting && (
+
+ All
+ {selectedID.length} Selected
+ }
+ onClick={() => setIsSelecting(false)}
+ aria-label="cancel"
/>
- );
- })}
-
-
-
-
+
+ )}
+ {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 !== "" && (
-
- )} */}
-
- );
-}