setupcomfyuimodels 1.0.3 → 1.1.0

package/package.json

@@ -3,7 +3,7 @@
   "bin": {
     "setupcomfyuimodels": "dist/index.js"
   },
-  "version": "1.0.3",
+  "version": "1.1.0",
   "description": "",
   "main": "src/index.ts",
   "scripts": {

package/src/downloadAssets.ts

@@ -12,25 +12,24 @@ import {
 import { getConfig } from "./getConfig";
 import chalk from "chalk";
 
-const { WORKSPACE, CIVITAI_TOKEN, HF_TOKEN } = process.env;
-
-if (!WORKSPACE) {
-  throw new Error("WORKSPACE environment variable is not set.");
-}
+const { CIVITAI_TOKEN, HF_TOKEN } = process.env;
 
 const MAX_CONCURRENT_DOWNLOADS = 8;
 let concurrentDownloads = 0;
 
-export const downloadAssets = async (configLocation: string) => {
+export const downloadAssets = async (
+  configLocation: string,
+  targetFolder: string,
+) => {
   const config = await getConfig(configLocation);
   const filePromises = [
     ...Object.entries(config.files).map(([targetLocation, url]) => {
       const fileName = path.basename(targetLocation),
         dir = path.dirname(targetLocation);
-      return scheduleDownload(dir, fileName, url);
+      return scheduleDownload(targetFolder, dir, fileName, url);
     }),
     ...Object.entries(config.folders).flatMap(([folder, urls]) =>
-      urls.map((url) => scheduleDownload(folder, null, url)),
+      urls.map((url) => scheduleDownload(targetFolder, folder, null, url)),
     ),
   ];
   setTotalProgress(0, filePromises.length);
@@ -76,6 +75,7 @@ const results = [] as {
   error?: string;
 }[];
 const scheduleDownload = async (
+  globalTargetLocation: string,
   targetLocation: string,
   targetFileName: string | null,
   url: string,
@@ -88,6 +88,7 @@ const scheduleDownload = async (
         concurrentDownloads++;
         try {
           const skipped = await downloadFile(
+            globalTargetLocation,
             nextDownload.targetLocation,
             nextDownload.targetFileName,
             nextDownload.url,
@@ -111,6 +112,7 @@ const scheduleDownload = async (
 };
 
 const downloadFile = async (
+  globalTargetLocation: string,
   targetLocation: string,
   targetFileName: string | null,
   url: string,
@@ -118,7 +120,7 @@ const downloadFile = async (
   const getTargetId = () => path.join(targetLocation, targetFileName || "");
   try {
     const targetDir = path.normalize(
-      `${WORKSPACE}/ComfyUI/models/${targetLocation}`,
+      `${globalTargetLocation}/${targetLocation}`,
     );
 
     fs.mkdirSync(targetDir, { recursive: true });

package/src/downloadAssets.ts~

@@ -0,0 +1,190 @@
+import * as fs from "fs";
+import path from "path";
+import { Readable } from "stream";
+import {
+  incrementTotalProgress,
+  log,
+  removeKey,
+  setTotalProgress,
+  stopProgressBar,
+  updateKey,
+} from "./progressbar";
+import { getConfig } from "./getConfig";
+import chalk from "chalk";
+
+const { WORKSPACE, CIVITAI_TOKEN, HF_TOKEN } = process.env;
+
+if (!WORKSPACE) {
+  throw new Error("WORKSPACE environment variable is not set.");
+}
+
+const MAX_CONCURRENT_DOWNLOADS = 8;
+let concurrentDownloads = 0;
+
+export const downloadAssets = async (configLocation: string) => {
+  const config = await getConfig(configLocation);
+  const filePromises = [
+    ...Object.entries(config.files).map(([targetLocation, url]) => {
+      const fileName = path.basename(targetLocation),
+        dir = path.dirname(targetLocation);
+      return scheduleDownload(dir, fileName, url);
+    }),
+    ...Object.entries(config.folders).flatMap(([folder, urls]) =>
+      urls.map((url) => scheduleDownload(folder, null, url)),
+    ),
+  ];
+  setTotalProgress(0, filePromises.length);
+  await Promise.all(filePromises);
+  stopProgressBar();
+  console.log(
+    chalk.bgWhite("Skipped downloads:\n"),
+    results
+      .filter((r) => r.skipped)
+      .map((r) => `${r.targetLocation}/${r.targetFileName || ""} from ${r.url}`)
+      .join("\n"),
+  );
+  console.log(
+    chalk.bgWhite("Successful downloads:\n"),
+    results
+      .filter((r) => r.downloaded)
+      .map((r) => `${r.targetLocation}/${r.targetFileName || ""} from ${r.url}`)
+      .join("\n"),
+  );
+  console.log(
+    chalk.bgWhite("Download errors:\n"),
+    results
+      .filter((r) => r.error)
+      .map(
+        (r) =>
+          `${r.targetLocation}/${r.targetFileName || ""} from ${r.url} had error: ${r.error}`,
+      )
+      .join("\n"),
+  );
+};
+
+const scheduledDownloads = [] as {
+  targetLocation: string;
+  targetFileName: string | null;
+  url: string;
+}[];
+const results = [] as {
+  targetLocation: string;
+  targetFileName: string | null;
+  url: string;
+  skipped?: boolean;
+  downloaded?: boolean;
+  error?: string;
+}[];
+const scheduleDownload = async (
+  targetLocation: string,
+  targetFileName: string | null,
+  url: string,
+) => {
+  scheduledDownloads.push({ targetLocation, url, targetFileName });
+  if (concurrentDownloads < MAX_CONCURRENT_DOWNLOADS) {
+    while (scheduledDownloads.length > 0) {
+      const nextDownload = scheduledDownloads.shift();
+      if (nextDownload) {
+        concurrentDownloads++;
+        try {
+          const skipped = await downloadFile(
+            nextDownload.targetLocation,
+            nextDownload.targetFileName,
+            nextDownload.url,
+          );
+          results.push({
+            ...nextDownload,
+            skipped,
+            downloaded: !skipped,
+          });
+        } catch (error) {
+          results.push({
+            ...nextDownload,
+            error: (error as Error).message,
+          });
+        }
+        incrementTotalProgress();
+        concurrentDownloads--;
+      }
+    }
+  }
+};
+
+const downloadFile = async (
+  targetLocation: string,
+  targetFileName: string | null,
+  url: string,
+) => {
+  const getTargetId = () => path.join(targetLocation, targetFileName || "");
+  try {
+    const targetDir = path.normalize(
+      `${WORKSPACE}/ComfyUI/models/${targetLocation}`,
+    );
+
+    fs.mkdirSync(targetDir, { recursive: true });
+    if (targetFileName && fs.existsSync(path.join(targetDir, targetFileName))) {
+      log(`File already exists, skipping: ${getTargetId()}`);
+      return true;
+    }
+    const isCivitAi = url.match(/https:\/\/civitai.com/);
+    const isHugging = url.match(/https:\/\/huggingface.co/);
+
+    const headers: Record<string, string> = {};
+    if (isCivitAi && CIVITAI_TOKEN) {
+      headers["Authorization"] = `Bearer ${CIVITAI_TOKEN}`;
+    }
+    if (isHugging && HF_TOKEN) {
+      headers["Authorization"] = `Bearer ${HF_TOKEN}`;
+    }
+
+    const response = await fetch(url, { headers });
+    if (!response.ok || !response.body) {
+      log(`Error for url: ${url}`);
+      throw new Error(
+        `Failed to download ${url}: ${response.status} ${response.statusText}`,
+      );
+    }
+    const total = Number(response.headers.get("content-length"));
+    const contentDisposition = response.headers.get("content-disposition");
+    if (!targetFileName) {
+      const fileNameMatch = contentDisposition?.match(
+        /filename="?(.+?)"?($|;)/,
+      );
+      if (fileNameMatch) {
+        targetFileName = fileNameMatch[1];
+      } else {
+        log(`Error for url: ${url}`);
+        throw new Error(
+          `Failed to extract filename from content-disposition for ${url}. Provide filename.`,
+        );
+      }
+    }
+    const filePath = path.join(targetDir, targetFileName);
+    if (fs.existsSync(filePath)) {
+      log(`File already exists, skipping: ${getTargetId()}`);
+      response.body?.cancel();
+      return true;
+    }
+    log(`Downloading ${getTargetId()} from ${url}`);
+
+    updateKey(getTargetId(), 0, total);
+    const readableNodeStream = Readable.fromWeb(response.body);
+    const fileStream = fs.createWriteStream(filePath + ".part");
+
+    let current = 0;
+    await new Promise<void>((resolve, reject) => {
+      readableNodeStream.on("data", (chunk: Buffer) => {
+        current += chunk.length;
+        updateKey(getTargetId(), current, total);
+      });
+      readableNodeStream.pipe(fileStream);
+      readableNodeStream.on("error", reject);
+      fileStream.on("finish", resolve);
+    });
+    fs.renameSync(filePath + ".part", filePath);
+    log(`Done downloading: ${url}`);
+    return false;
+  } finally {
+    removeKey(getTargetId());
+  }
+};

package/src/index.ts

@@ -7,7 +7,8 @@ program.name("Setup ComfyUi Models");
 program
   .command("download")
   .argument("<config>", "JSON config file location")
-  .action(async (configLocation: string) => {
-    await downloadAssets(configLocation);
+  .option("-l, --location <path>", "Target download location")
+  .action(async (configLocation: string, options) => {
+    await downloadAssets(configLocation, options.location || "./");
   });
 program.parse();

package/src/index.ts~

@@ -0,0 +1,13 @@
+#!/usr/bin/env node
+import { Command } from "@commander-js/extra-typings";
+import { downloadAssets } from "./downloadAssets";
+
+const program = new Command();
+program.name("Setup ComfyUi Models");
+program
+  .command("download")
+  .argument("<config>", "JSON config file location")
+  .action(async (configLocation: string) => {
+    await downloadAssets(configLocation);
+  });
+program.parse();