add prompt checker

app/api/route.ts

@@ -1,6 +1,9 @@
 import { PrismaClient } from '@prisma/client'
 
+import list_styles from "@/assets/list_styles.json"
 import { UploaderDataset } from './uploader'
+import { isTextNSFW } from '@/utils/checker/prompt'
+import { isImageNSFW } from '@/utils/checker/image'
 
 const prisma = new PrismaClient()
 
@@ -13,12 +16,17 @@ export async function POST(
     ['x-use-cache']: "0"
   }
 
-  const { inputs, negative_prompt } = await request.json()  
+  const { inputs, style } = await request.json()
+  const findStyle = list_styles.find((item) => item.name === style)
+
+  const textIsNSFW = await isTextNSFW(inputs, global_headers)
+  if (textIsNSFW) return Response.json({ status: 401, ok: false, message: "Prompt is not safe for work." });
+
   const response = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/models/stabilityai/stable-diffusion-xl-base-1.0`, {
     method: 'POST',
     body: JSON.stringify({
-      inputs,
-      negative_prompt
+      inputs: findStyle?.prompt.replace("{prompt}", inputs) ?? inputs,
+      negative_prompt: findStyle?.negative_prompt ?? "",
     }),
     headers: global_headers,
   })
@@ -27,37 +35,9 @@ export async function POST(
   if (res?.error) return Response.json({ status: response.status, ok: false, message: res.error });
 
   const blob = await response.blob()
-  const headers = new Headers();
-  headers.set("Content-Type", "image/*");
-
-  const checkIfIsNSFW = (blob: Blob) => {
-    return new Promise(async (resolve, reject) => {
-      const response_nsfw = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/models/DamarJati/NSFW-Filterization-DecentScan`, {
-        method: 'POST',
-        headers: {
-          ...global_headers,
-          ...headers,
-        },
-        body: blob,
-      })
-      const isNSFW = await response_nsfw.clone().json().catch(() => ({}));
-
-      if (isNSFW?.error && isNSFW?.estimated_time) {
-        setTimeout(() => {
-          checkIfIsNSFW(blob)
-        }, isNSFW?.estimated_time * 100);
-      } else resolve(isNSFW)
-    })
-  }
   
-  const isNSFW: any = await checkIfIsNSFW(blob)
-  if (isNSFW?.error) return Response.json({ status: 500, ok: false, message: isNSFW?.error });
-  if (isNSFW?.length) {
-    const scoreNotSafe = isNSFW?.find((n: { label: string }) => n.label === "no_safe");
-    if (scoreNotSafe?.score > 0.85) {
-      return Response.json({ status: 401, ok: false, message: "Image is not safe for work." });
-    }
-  }
+  const imageIsNSFW = await isImageNSFW(blob, global_headers)
+  if (imageIsNSFW) return Response.json({ status: 401, ok: false, message: "Image is not safe for work." });
 
   const name = Date.now() + `-${inputs.replace(/[^a-zA-Z0-9]/g, '-').slice(0, 10).toLowerCase()}`
   const { ok, message } = await UploaderDataset(blob, name)
@@ -71,6 +51,6 @@ export async function POST(
     },
   })
   
-  return Response.json({ image: new_image, status: 200, ok: true, headers });
+  return Response.json({ image: new_image, status: 200, ok: true });
 
 }

components/main/hooks/useInputGeneration.ts

@@ -50,13 +50,11 @@ export const useInputGeneration = () => {
         }
       })
 
-      const findStyle = list_styles.find((item) => item.name === style)
-
       const response = await fetch("/api", {
         method: "POST",
         body: JSON.stringify({
-          inputs: findStyle?.prompt.replace("{prompt}", prompt) ?? prompt,
-          negative_prompt: findStyle?.negative_prompt ?? "",
+          inputs: prompt,
+          style,
         }),
       })
       const data = await response.json()

utils/checker/image.ts

@@ -0,0 +1,28 @@
+export const isImageNSFW = async (blob: Blob, global_headers: any) => {
+  return new Promise(async (resolve, reject) => {
+    const headers = new Headers();
+    headers.set("Content-Type", "image/*");
+    const request = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/models/DamarJati/NSFW-Filterization-DecentScan`, {
+      method: 'POST',
+      headers: {
+        ...global_headers,
+        ...headers,
+      },
+      body: blob,
+    })
+    const res = await request.clone().json().catch(() => ({}));
+
+    if (res?.error && res?.estimated_time) {
+      setTimeout(() => {
+        isImageNSFW(blob, global_headers)
+      }, res?.estimated_time * 100);
+    } else {
+      if (res?.error) return Response.json({ status: 500, ok: false, message: res?.error });
+      if (res?.length) {
+        const isNSFW = res?.find((n: { label: string }) => n.label === "no_safe")?.score > 0.85 ?? false;
+        resolve(isNSFW)
+      }
+      resolve(true)
+    }
+  })
+}

utils/checker/prompt.ts

@@ -0,0 +1,14 @@
+export const isTextNSFW = async (prompt: string, headers: any) => {
+  return new Promise(async (resolve, reject) => {
+    const request = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/models/michellejieli/NSFW_text_classifier`, {
+      method: 'POST',
+      body: JSON.stringify({
+        inputs: prompt,
+      }),
+      headers: headers,
+    })
+    const res = await request.clone().json().catch(() => ({}));
+    const isNSFW = res?.[0]?.find((item: { label: string, score: number }) => item?.label === "NSFW")?.score > 0.92 ?? false
+    resolve(isNSFW)
+  })
+}