setupcomfyuimodels 1.0.3 → 1.2.0

package/package.json

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

package/src/downloadAssets.ts

@@ -9,28 +9,40 @@ import {
   stopProgressBar,
   updateKey,
 } from "./progressbar";
-import { getConfig } from "./getConfig";
+import { Config } 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) => {
-  const config = await getConfig(configLocation);
+export const downloadAssets = async (config: Config, targetFolder: string) => {
+  const filesConfigs = [
+    config.files,
+    ...Object.entries(config.groups)
+      .filter(([groupName]) => config.groupsEnabled[groupName])
+      .map(([_, f]) => f.files),
+  ];
+  const foldersConfigs = [
+    config.folders,
+    ...Object.entries(config.groups)
+      .filter(([groupName]) => config.groupsEnabled[groupName])
+      .map(([_, f]) => f.folders),
+  ];
+
   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)),
+    ...filesConfigs.flatMap((config) =>
+      Object.entries(config.files).map(([targetLocation, url]) => {
+        const fileName = path.basename(targetLocation),
+          dir = path.dirname(targetLocation);
+        return scheduleDownload(targetFolder, dir, fileName, url);
+      }),
+    ),
+    ...foldersConfigs.flatMap((config) =>
+      Object.entries(config).flatMap(([folder, urls]) =>
+        urls.map((url) => scheduleDownload(targetFolder, folder, null, url)),
+      ),
     ),
   ];
   setTotalProgress(0, filePromises.length);
@@ -76,10 +88,25 @@ const results = [] as {
   error?: string;
 }[];
 const scheduleDownload = async (
+  globalTargetLocation: string,
   targetLocation: string,
   targetFileName: string | null,
   url: string,
 ) => {
+  const matchingJob = scheduledDownloads.find(
+    (d) =>
+      d.targetLocation === targetLocation &&
+      d.targetFileName === targetFileName,
+  );
+
+  if (matchingJob) {
+    if (matchingJob.url !== url) {
+      throw new Error(
+        `Conflicting download jobs for ${targetLocation}/${targetFileName || ""}: ${matchingJob.url} and ${url}`,
+      );
+    }
+    return;
+  }
   scheduledDownloads.push({ targetLocation, url, targetFileName });
   if (concurrentDownloads < MAX_CONCURRENT_DOWNLOADS) {
     while (scheduledDownloads.length > 0) {
@@ -88,6 +115,7 @@ const scheduleDownload = async (
         concurrentDownloads++;
         try {
           const skipped = await downloadFile(
+            globalTargetLocation,
             nextDownload.targetLocation,
             nextDownload.targetFileName,
             nextDownload.url,
@@ -111,6 +139,7 @@ const scheduleDownload = async (
 };
 
 const downloadFile = async (
+  globalTargetLocation: string,
   targetLocation: string,
   targetFileName: string | null,
   url: string,
@@ -118,7 +147,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,192 @@
+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 { CIVITAI_TOKEN, HF_TOKEN } = process.env;
+
+const MAX_CONCURRENT_DOWNLOADS = 8;
+let concurrentDownloads = 0;
+
+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(targetFolder, dir, fileName, url);
+    }),
+    ...Object.entries(config.folders).flatMap(([folder, urls]) =>
+      urls.map((url) => scheduleDownload(targetFolder, 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 (
+  globalTargetLocation: string,
+  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(
+            globalTargetLocation,
+            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 (
+  globalTargetLocation: string,
+  targetLocation: string,
+  targetFileName: string | null,
+  url: string,
+) => {
+  const getTargetId = () => path.join(targetLocation, targetFileName || "");
+  try {
+    const targetDir = path.normalize(
+      `${globalTargetLocation}/${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/getConfig.ts

@@ -1,11 +1,24 @@
 import { readFileSync } from "fs";
 import { parse } from "comment-json";
 
+type FilesAndFolders = {
+  files: Record<string, string>;
+  folders: Record<string, string[]>;
+};
+
+export type Config = {
+  groupsEnabled: Record<string, boolean>;
+  groups: Record<string, FilesAndFolders>;
+  replacements?: {
+    file: string;
+    original: string;
+    replacement: string;
+    line: number;
+  }[];
+} & FilesAndFolders;
+
 export const getConfig = async (configLocation: string) => {
   const fileContent = readFileSync(configLocation, "utf-8");
 
-  return parse(fileContent) as unknown as {
-    files: Record<string, string>;
-    folders: Record<string, string[]>;
-  };
+  return parse(fileContent) as unknown as Config;
 };

package/src/getConfig.ts~

@@ -0,0 +1,11 @@
+import { readFileSync } from "fs";
+import { parse } from "comment-json";
+
+export const getConfig = async (configLocation: string) => {
+  const fileContent = readFileSync(configLocation, "utf-8");
+
+  return parse(fileContent) as unknown as {
+    files: Record<string, string>;
+    folders: Record<string, string[]>;
+  };
+};

package/src/index.ts

@@ -1,13 +1,18 @@
 #!/usr/bin/env node
 import { Command } from "@commander-js/extra-typings";
 import { downloadAssets } from "./downloadAssets";
+import { replaceInFiles } from "./replacements";
+import { getConfig } from "./getConfig";
 
 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);
+  .option("-l, --location <path>", "Target download location")
+  .action(async (configLocation: string, options) => {
+    const config = await getConfig(configLocation);
+    await downloadAssets(config, options.location || "./");
+    replaceInFiles(config, options.location || "./");
   });
 program.parse();

package/src/index.ts~

@@ -0,0 +1,14 @@
+#!/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")
+  .option("-l, --location <path>", "Target download location")
+  .action(async (configLocation: string, options) => {
+    await downloadAssets(configLocation, options.location || "./");
+  });
+program.parse();

package/src/replacements.ts

@@ -0,0 +1,37 @@
+import * as fs from "fs";
+import * as path from "path";
+import { Config } from "./getConfig";
+
+export const replaceInFiles = (config: Config, targetFolder: string) => {
+  const replacements = config.replacements || [];
+  for (const replacement of replacements) {
+    const filePath = path.join(targetFolder, replacement.file);
+    if (fs.existsSync(filePath)) {
+      let fileContent = fs.readFileSync(filePath, "utf-8");
+      if (fileContent.includes(replacement.original)) {
+        const lines = fileContent.split("\n");
+        if (
+          lines.length >= replacement.line &&
+          lines[replacement.line - 1].includes(replacement.original)
+        ) {
+          lines[replacement.line - 1] = lines[replacement.line - 1].replace(
+            replacement.original,
+            replacement.replacement,
+          );
+          fileContent = lines.join("\n");
+        } else {
+          console.warn(
+            `Line number mismatch in ${replacement.file}. Expected line ${replacement.line} to contain the original code.`,
+          );
+        }
+
+        fs.writeFileSync(filePath, fileContent, "utf-8");
+        console.log(`Patched ${replacement.file} at line ${replacement.line}`);
+      } else {
+        console.warn(`Original code not found in ${replacement.file}`);
+      }
+    } else {
+      console.warn(`File not found: ${replacement.file}`);
+    }
+  }
+};