From 660781afd9fe96de44e16725ad27f82922f628e8 Mon Sep 17 00:00:00 2001 From: MeiMei <30769358+mei23@users.noreply.github.com> Date: Tue, 12 Jul 2022 10:38:57 +0900 Subject: [PATCH] Fix crash at startup if TensorFlow is not supported (#8984) * Lazy loading tensorflow * CHANGELOG * CHANGELOG * Check CPU flags * . --- CHANGELOG.md | 7 ++++++ .../backend/src/services/detect-sensitive.ts | 24 +++++++++++++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 37e3d52a0..1e35a2e66 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,13 @@ You should also include the user name that made the change. --> +## 12.x.x (unreleased) + +### Improvements + +### Bugfixes +- Server: Fix crash at startup if TensorFlow is not supported @mei23 + ## 12.112.3 (2022/07/09) ### Improvements diff --git a/packages/backend/src/services/detect-sensitive.ts b/packages/backend/src/services/detect-sensitive.ts index 0fa263599..2ade39d52 100644 --- a/packages/backend/src/services/detect-sensitive.ts +++ b/packages/backend/src/services/detect-sensitive.ts @@ -2,19 +2,34 @@ import * as fs from 'node:fs'; import { fileURLToPath } from 'node:url'; import { dirname } from 'node:path'; import * as nsfw from 'nsfwjs'; -import * as tf from '@tensorflow/tfjs-node'; +import si from 'systeminformation'; const _filename = fileURLToPath(import.meta.url); const _dirname = dirname(_filename); +const REQUIRED_CPU_FLAGS = ['avx2', 'fma']; +let isSupportedCpu: undefined | boolean = undefined; + let model: nsfw.NSFWJS; export async function detectSensitive(path: string): Promise { try { + if (isSupportedCpu === undefined) { + const cpuFlags = await 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 (model == null) 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 tf.Tensor3D; + const image = await tf.node.decodeImage(buffer, 3) as any; try { const predictions = await model.classify(image); return predictions; @@ -26,3 +41,8 @@ export async function detectSensitive(path: string): Promise { + const str = await si.cpuFlags(); + return str.split(/\s+/); +}