setupcomfyuimodels 1.1.0 → 1.2.0

package/package.json

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

package/src/downloadAssets.ts

@@ -9,7 +9,7 @@ import {
   stopProgressBar,
   updateKey,
 } from "./progressbar";
-import { getConfig } from "./getConfig";
+import { Config } from "./getConfig";
 import chalk from "chalk";
 
 const { CIVITAI_TOKEN, HF_TOKEN } = process.env;
@@ -17,19 +17,32 @@ 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);
+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(targetFolder, dir, fileName, url);
-    }),
-    ...Object.entries(config.folders).flatMap(([folder, urls]) =>
-      urls.map((url) => scheduleDownload(targetFolder, 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);
@@ -80,6 +93,20 @@ const scheduleDownload = async (
   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) {

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/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,6 +1,8 @@
 #!/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");
@@ -9,6 +11,8 @@ program
   .argument("<config>", "JSON config file location")
   .option("-l, --location <path>", "Target download location")
   .action(async (configLocation: string, options) => {
-    await downloadAssets(configLocation, options.location || "./");
+    const config = await getConfig(configLocation);
+    await downloadAssets(config, options.location || "./");
+    replaceInFiles(config, options.location || "./");
   });
 program.parse();

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/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}`);
+    }
+  }
+};