Move transformers.js model cache under the data root

c2057da348e0ff9b26feeaf9d03d65db12e87378

Cohee <18619528+Cohee1207@users.noreply.github.com>

index.d.ts+5 -0 Viewed

@@ -9,6 +9,11 @@ declare global {
9	};	9	};
10	}	10	}
11	}	11	}
		12
		13	/**
		14	* The root directory for user data.
		15	*/
		16	var DATA_ROOT: string;
12	}	17	}
13		18
14	declare module 'express-session' {	19	declare module 'express-session' {

src/transformers.mjs+36 -3 Viewed

 import { pipeline, env, RawImage, Pipeline } from 'sillytavern-transformers';
 import { getConfigValue } from './util.js';
 import path from 'path';
+import fs from 'fs';
 configureTransformers();

         configField: 'extras.textToSpeechModel',
         quantized: false,
     },
-}
+};
 /**
  * Gets a RawImage object from a base64-encoded image.

+    }
+}
+async function migrateCacheToDataDir() {
+    const oldCacheDir = path.join(process.cwd(), 'cache');
+    const newCacheDir = path.join(global.DATA_ROOT, '_cache');
+    if (!fs.existsSync(newCacheDir)) {
+        fs.mkdirSync(newCacheDir, { recursive: true });
+    }
+    if (fs.existsSync(oldCacheDir) && fs.statSync(oldCacheDir).isDirectory()) {
+        const files = fs.readdirSync(oldCacheDir);
+        if (files.length === 0) {
+            return;
+        }
+        console.log('Migrating model cache files to data directory. Please wait...');
+        for (const file of files) {
+            try {
+                const oldPath = path.join(oldCacheDir, file);
+                const newPath = path.join(newCacheDir, file);
+                fs.cpSync(oldPath, newPath, { recursive: true, force: true });
+                fs.rmSync(oldPath, { recursive: true, force: true });
+            } catch (error) {
+                console.warn('Failed to migrate cache file. The model will be re-downloaded.', error);
+            }
+        }
+    }
+}
 /**
  * Gets the transformers.js pipeline for a given task.
  * @param {import('sillytavern-transformers').PipelineType} task The task to get the pipeline for

  * @returns {Promise<Pipeline>} Pipeline for the task
  */
 async function getPipeline(task, forceModel = '') {
+    await migrateCacheToDataDir();
     if (tasks[task].pipeline) {
         if (forceModel === '' || tasks[task].currentModel === forceModel) {
             return tasks[task].pipeline;

         await tasks[task].pipeline.dispose();
+    }
-    const cache_dir = path.join(process.cwd(), 'cache');
+    const cacheDir = path.join(global.DATA_ROOT, '_cache');
     const model = forceModel || getModelForTask(task);
     const localOnly = getConfigValue('extras.disableAutoDownload', false);
     console.log('Initializing transformers.js pipeline for task', task, 'with model', model);
-    const instance = await pipeline(task, model, { cache_dir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly });
+    const instance = await pipeline(task, model, { cache_dir: cacheDir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly });
     tasks[task].pipeline = instance;
     tasks[task].currentModel = model;
     return instance;

src/users.js+5 -11 Viewed

 const ANON_CSRF_SECRET = crypto.randomBytes(64).toString('base64');
 /**
- * The root directory for user data.
- * @type {string}
- */
-let DATA_ROOT = './data';
-/**
  * Cache for user directories.
  * @type {Map<string, UserDirectoryList>}
  */

     console.log();
     console.log(color.magenta('Preparing to migrate user data...'));
-    console.log(`All public data will be moved to the ${DATA_ROOT} directory.`);
+    console.log(`All public data will be moved to the ${global.DATA_ROOT} directory.`);
     console.log('This process may take a while depending on the amount of data to move.');
     console.log(`Backups will be placed in the ${PUBLIC_DIRECTORIES.backups} directory.`);
     console.log(`The process will start in ${TIMEOUT} seconds. Press Ctrl+C to cancel.`);

  * @returns {Promise<void>}
  */
 async function initUserStorage(dataRoot) {
-    DATA_ROOT = dataRoot;
+    global.DATA_ROOT = dataRoot;
-    console.log('Using data root:', color.green(DATA_ROOT));
+    console.log('Using data root:', color.green(global.DATA_ROOT));
     console.log();
     await storage.init({
-        dir: path.join(DATA_ROOT, '_storage'),
+        dir: path.join(global.DATA_ROOT, '_storage'),
         ttl: false, // Never expire
     });

     const directories = structuredClone(USER_DIRECTORY_TEMPLATE);
     for (const key in directories) {
-        directories[key] = path.join(DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]);
+        directories[key] = path.join(global.DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]);
+    }
     DIRECTORIES_CACHE.set(handle, directories);
     return directories;