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 37e3d52a0e..1e35a2e66e 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 0fa263599b..2ade39d524 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<nsfw.predictionType[] | null> {
 	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<nsfw.predictionType
 		return null;
 	}
 }
+
+async function getCpuFlags(): Promise<string[]> {
+	const str = await si.cpuFlags();
+	return str.split(/\s+/);
+}