sendnrw/misskey
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(backend): use mutex for nsfw model loading (#11109)

Browse Source 
Co-authored-by: tamaina <tamaina@hotmail.co.jp>
This commit is contained in:
popkirby
2023-07-08 07:27:26 +09:00
committed by GitHub
parent d84796588c
commit 8daca59ca6
4 changed files with 20 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
packages/backend/src/core/AiService.ts
View File

@@ -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)
@@ -39,7 +41,13 @@ export class AiService {
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;