diff options
Diffstat (limited to 'packages/backend/src/core/AiService.ts')
| -rw-r--r-- | packages/backend/src/core/AiService.ts | 18 |
1 files changed, 13 insertions, 5 deletions
diff --git a/packages/backend/src/core/AiService.ts b/packages/backend/src/core/AiService.ts index 059e335eff..c0596446dd 100644 --- a/packages/backend/src/core/AiService.ts +++ b/packages/backend/src/core/AiService.ts @@ -4,6 +4,7 @@ import { dirname } from 'node:path'; import { Inject, Injectable } from '@nestjs/common'; import * as nsfw from 'nsfwjs'; import si from 'systeminformation'; +import { Mutex } from 'async-mutex'; import type { Config } from '@/config.js'; import { DI } from '@/di-symbols.js'; import { bindThis } from '@/decorators.js'; @@ -17,6 +18,7 @@ let isSupportedCpu: undefined | boolean = undefined; @Injectable() export class AiService { private model: nsfw.NSFWJS; + private modelLoadMutex: Mutex = new Mutex(); constructor( @Inject(DI.config) @@ -31,16 +33,22 @@ export class AiService { const cpuFlags = await this.getCpuFlags(); isSupportedCpu = REQUIRED_CPU_FLAGS.every(required => cpuFlags.includes(required)); } - + if (!isSupportedCpu) { console.error('This CPU cannot use TensorFlow.'); return null; } - + const tf = await import('@tensorflow/tfjs-node'); - - if (this.model == null) this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 }); - + + if (this.model == null) { + await this.modelLoadMutex.runExclusive(async () => { + if (this.model == null) { + this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 }); + } + }); + } + const buffer = await fs.promises.readFile(path); const image = await tf.node.decodeImage(buffer, 3) as any; try { |