feat tensorflow WASM

Dockerfile

@@ -13,7 +13,7 @@ ENV NSFW_VALUE 0.7
 ENV OCR_SENSITIVE false
 ENV OCR_API_KEY ""
 ENV SENSITIVE_WORDS "https://raw.githubusercontent.com/Tohrusky/chinese-sensitive-go/main/dict/boss.txt"
-ENV TENSORFLOW_BACKEND cpu
+ENV TENSORFLOW_BACKEND wasm
 
 EXPOSE 3008
 

packages/nsfw/config.ts

@@ -12,4 +12,4 @@ export const SENSITIVE_WORDS =
   process.env.SENSITIVE_WORDS ||
   'https://raw.githubusercontent.com/Tohrusky/chinese-sensitive-go/main/dict/boss.txt' // sensitive words dictionary, load online
 
-export const TENSORFLOW_BACKEND = process.env.TENSORFLOW_BACKEND || 'cpu' // tensorflow backend
+export const TENSORFLOW_BACKEND = process.env.TENSORFLOW_BACKEND || 'wasm' // tensorflow backend

packages/nsfw/index.ts

@@ -5,7 +5,8 @@ import { loadModel } from './src/inference/nsfw'
 import { handleNSFW } from './src/services/nsfw'
 import { inject } from '@vercel/analytics'
 import { initMintFilter } from './src/inference/sensitive'
-import { OCR_SENSITIVE } from './config'
+import { OCR_SENSITIVE, TENSORFLOW_BACKEND } from './config'
+import * as tf from '@tensorflow/tfjs'
 
 inject()
 
@@ -23,16 +24,21 @@ app.use(bodyParser())
 app.use(router.routes())
 app.use(router.allowedMethods())
 
-loadModel().then(async () => {
-  console.log('model loaded successfully')
+tf.setBackend(TENSORFLOW_BACKEND).then(async () => {
+  await tf.ready()
+  console.log('Tensorflow backend: ', tf.getBackend())
 
-  // 预加载敏感词库
-  if (OCR_SENSITIVE) {
-    await initMintFilter()
-  }
+  loadModel().then(async () => {
+    console.log('model loaded successfully')
 
-  app.listen(3008, () => {
-    console.log('3008 is listening')
+    // 预加载敏感词库
+    if (OCR_SENSITIVE) {
+      await initMintFilter()
+    }
+
+    app.listen(3008, () => {
+      console.log('3008 is listening')
+    })
   })
 })
 

packages/nsfw/package.json

@@ -21,6 +21,7 @@
     "tesseract.js": "^4.1.2",
     "@napi-rs/canvas": "^0.1.44",
     "@tensorflow/tfjs": "^4.10.0",
+    "@tensorflow/tfjs-backend-wasm": "4.11.0",
     "koa": "^2.14.2",
     "koa-body": "^6.0.1",
     "koa-router": "^12.0.0",

packages/nsfw/src/lib/nsfwjs/index.ts

@@ -1,8 +1,6 @@
 import * as tf from '@tensorflow/tfjs'
+import '@tensorflow/tfjs-backend-wasm'
 import { NSFW_CLASSES } from './nsfw_classes'
-import { TENSORFLOW_BACKEND } from '../../../config'
-
-tf.setBackend(TENSORFLOW_BACKEND).then(() => console.log('Tensorflow backend: ', tf.getBackend()))
 
 export type frameResult = {
   index: number