Merge pull request #129 from arslan2012/feat/queue

update queue to show queue in the modal and added custom url download

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'])
+    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)
+    

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 (
+        <AlertDialog
+            isOpen={isOpen}
+            leastDestructiveRef={cancelRef}
+            onClose={onClose}
+        >
+            <AlertDialogOverlay>
+                <AlertDialogContent>
+                    <AlertDialogHeader fontSize='lg' fontWeight='bold'>
+                        Choose Folder
+                    </AlertDialogHeader>
+
+                    <AlertDialogBody>
+                        <Select placeholder='Select option' value={folderPath} onChange={e => setFolderPath(e.target.value)}>
+                            {allFolderPaths.map((folderPath) => (
+                                <option value={folderPath}>{folderPath}</option>
+                            ))}
+                        </Select>
+                    </AlertDialogBody>
+
+                    <AlertDialogFooter>
+                        <Button ref={cancelRef} onClick={onClose}>
+                            Cancel
+                        </Button>
+                        <Button onClick={() => selectFolder(folderPath)} ml={3}>
+                            Confirm
+                        </Button>
+                    </AlertDialogFooter>
+                </AlertDialogContent>
+            </AlertDialogOverlay>
+        </AlertDialog>
+    );
+}

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<string[]>([]);
   const [searchQuery, setSearchQuery] = useState("");
+  const [queue, setQueue] = useState<Queue[]>([]);
+  const { isOpen, onOpen, onClose } = useDisclosure();
+  const file = useRef<CivitiModelFileVersion>();
   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 (
-    <Modal isOpen={true} onClose={onclose} blockScrollOnMount={true}>
-      <ModalOverlay />
-      <ModalContent width={"90%"} maxWidth={"90vw"} height={"90vh"}>
-        <ModalHeader>
-          <HStack gap={2} mb={2} alignItems={"center"}>
-            <Heading size={"md"} mr={2}>
-              Models
-            </Heading>
-            <InstallModelSearchBar setSearchQuery={setSearchQuery} />
-          </HStack>
-          <InstallProgress />
-          <HStack gap={2} mb={2} wrap={"wrap"}>
-            <Button
-              size={"sm"}
-              py={1}
-              onClick={() => {
-                setModelType(undefined);
-              }}
-              isActive={modelType == null}
-            >
-              All
-            </Button>
-            {ALL_MODEL_TYPES.map((type) => {
-              return (
-                <Button
-                  size={"sm"}
-                  py={1}
-                  isActive={modelType === type}
-                  onClick={() => {
-                    setModelType(type);
-                  }}
-                >
-                  {type}
-                </Button>
-              );
-            })}
-          </HStack>
-          {isSelecting && (
-            <HStack gap={3}>
-              <Checkbox isChecked={isAllSelected}>All</Checkbox>
-              <Text fontSize={16}>{selectedID.length} Selected</Text>
-              <IconButton
+    <>
+      <Modal isOpen={true} onClose={onclose} blockScrollOnMount={true}>
+        <ModalOverlay />
+        <ModalContent width={"90%"} maxWidth={"90vw"} height={"90vh"}>
+          <ModalHeader>
+            <HStack gap={2} mb={2} alignItems={"center"}>
+              <Heading size={"md"} mr={2}>
+                Models
+              </Heading>
+              <InstallModelSearchBar setSearchQuery={setSearchQuery} />
+              <Button
                 size={"sm"}
-                icon={<IconX size={19} />}
-                onClick={() => setIsSelecting(false)}
-                aria-label="cancel"
-              />
+                py={1}
+                mr={8}
+                onClick={customUrlDownload}
+              >
+                Custom URL Install
+              </Button>
             </HStack>
-          )}
-          {loading && (
-            <Spinner
-              thickness="4px"
-              emptyColor="gray.200"
-              color="pink.500"
-              size="lg"
-            />
-          )}
-        </ModalHeader>
-        <ModalCloseButton />
-        <ModalBody overflowY={"auto"}>
-          <HStack wrap={"wrap"}>
-            {models?.map((model) => {
-              return (
-                <ModelCard
-                  model={model}
-                  key={model.id}
-                  onClickInstallModel={onClickInstallModel}
-                  installing={installing}
+            <HStack gap={2} mb={2} wrap={"wrap"}>
+              <Button
+                size={"sm"}
+                py={1}
+                onClick={() => {
+                  setModelType(undefined);
+                }}
+                isActive={modelType == null}
+              >
+                All
+              </Button>
+              {ALL_MODEL_TYPES.map((type) => {
+                return (
+                  <Button
+                    size={"sm"}
+                    py={1}
+                    isActive={modelType === type}
+                    onClick={() => {
+                      setModelType(type);
+                    }}
+                  >
+                    {type}
+                  </Button>
+                );
+              })}
+            </HStack>
+            {isSelecting && (
+              <HStack gap={3}>
+                <Checkbox isChecked={isAllSelected}>All</Checkbox>
+                <Text fontSize={16}>{selectedID.length} Selected</Text>
+                <IconButton
+                  size={"sm"}
+                  icon={<IconX size={19} />}
+                  onClick={() => setIsSelecting(false)}
+                  aria-label="cancel"
                 />
-              );
-            })}
-          </HStack>
-        </ModalBody>
-      </ModalContent>
-    </Modal>
+              </HStack>
+            )}
+            {loading && (
+              <Spinner
+                thickness="4px"
+                emptyColor="gray.200"
+                color="pink.500"
+                size="lg"
+              />
+            )}
+          </ModalHeader>
+          <ModalCloseButton />
+          <ModalBody overflowY={"auto"}>
+            <HStack wrap={"wrap"}>
+              {models?.map((model) => {
+                return (
+                  <ModelCard
+                    model={model}
+                    key={model.id}
+                    onClickInstallModel={onClickInstallModel}
+                    installing={installing}
+                  />
+                );
+              })}
+            </HStack>
+            <Stack spacing={5} pos="absolute" bottom="0" left="0" width="50%" zIndex={80} backgroundColor="white" paddingX={5}>
+              {queue.map(({ save_path, progress }) => (
+                <HStack>
+                  <Text fontSize={16} width="40%">{save_path.replace(/^.*[\\/]/, '')}</Text>
+                  <Progress isIndeterminate={!progress} hasStripe width="60%" value={progress} />
+                </HStack>
+              ))}
+            </Stack>
+          </ModalBody>
+        </ModalContent>
+      </Modal>
+      <ChooseFolder isOpen={isOpen} onClose={onClose} selectFolder={downloadModels} />
+    </>
   );
 }
+

ui/src/model-manager/install-models/InstallProgress.tsx

@@ -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 (
-    <HStack>
-      <Text fontWeight={"500"} fontSize={18} py={2}>
-        {progress}
-      </Text>
-      {/* {progress !== "" && (
-        <Button onClick={() => setCancelInstall(true)}>Cancel</Button>
-      )} */}
-    </HStack>
-  );
-}