Merge pull request #2667 from SillyTavern/dataroot-cache Move transformers.js model cache under the data root

fc02898a97897df3e1285989d603ea76f4760a37

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

Signed

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	19	declare module 'express-session' {

package-lock.json+1 -1 Viewed

                 "rate-limiter-flexible": "^5.0.0",
                 "response-time": "^2.3.2",
                 "sanitize-filename": "^1.6.3",
                 "sillytavern-transformers": "^2.14.6",
                 "simple-git": "^3.19.1",
                 "tiktoken": "^1.0.15",
                 "vectra": "^0.2.2",

package.json+1 -1 Viewed

         "rate-limiter-flexible": "^5.0.0",
         "response-time": "^2.3.2",
         "sanitize-filename": "^1.6.3",
         "sillytavern-transformers": "^2.14.6",
         "simple-git": "^3.19.1",
         "tiktoken": "^1.0.15",
         "vectra": "^0.2.2",

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_dircacheDir = path.join(processglobal.cwd()DATA_ROOT, 'cache_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: 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 ${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) {
     global.DATA_ROOT = dataRoot;
     console.log('Using data root:', color.green(global.DATA_ROOT));
     console.log();
     await storage.init({
         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(global.DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]);
+    }
     DIRECTORIES_CACHE.set(handle, directories);
     return directories;