File size: 2,250 Bytes
4f2c36e
 
 
6b92c38
 
 
 
2d49c67
 
 
 
 
 
8b95140
2d49c67
6b92c38
 
 
8b95140
6b92c38
2d49c67
6b92c38
 
0891679
 
 
6b92c38
 
 
4f2c36e
848e268
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
53353f2
2d49c67
 
53353f2
8b95140
 
 
 
4f2c36e
 
 
 
 
 
 
 
 
6b92c38
0891679
 
4f2c36e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import { PrismaClient } from '@prisma/client'

const prisma = new PrismaClient()

export async function POST(
  request: Request,
) {
  const global_headers = {
    Authorization: `Bearer ${process.env.NEXT_PUBLIC_APP_HF_TOKEN}`,
    'Content-Type': 'application/json',
    ['x-use-cache']: "0"
  }

  const { inputs, negative_prompt } = await request.json()  
  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
    }),
    headers: global_headers,
  })

  const res = await response.clone().json().catch(() => ({}));
  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) {
        console.log("retrying...")
        setTimeout(() => {
          checkIfIsNSFW(blob)
        }, isNSFW?.estimated_time * 1000);
      } else resolve(isNSFW)
    })
  }

  const isNSFW: any = await checkIfIsNSFW(blob)
  console.log(isNSFW)

  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.8) {
      return Response.json({ status: 401, ok: false, message: "Image is not safe for work." });
    }
  }
  
  const arrayBuffer = await blob.arrayBuffer()
  const bytes = Buffer.from(arrayBuffer)

  const new_blob = await prisma.image.create({
    data: {
      prompt: inputs,
      blob: bytes,
    },
  })
  
  return Response.json({ blob: new_blob, status: 200, ok: true, headers });

}