File size: 1,647 Bytes
4f2c36e
 
8b95140
 
 
4f2c36e
6b92c38
 
 
 
8b95140
6b92c38
 
 
 
8b95140
6b92c38
 
3b6777d
6b92c38
5240c42
6b92c38
 
 
0891679
 
 
 
6b92c38
 
 
4f2c36e
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
import { PrismaClient } from '@prisma/client'

import PipelineSingleton from './pipeline';
import { RawImage } from '@xenova/transformers';

const prisma = new PrismaClient()

export async function POST(
  request: Request,
) {
  const { inputs, negative_prompt } = await request.json()  
  const response = await fetch('https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0', {
    method: 'POST',
    body: JSON.stringify({
      inputs,
      negative_prompt
    }),
    headers: {
      Authorization: `Bearer ${process.env.NEXT_PUBLIC_APP_HF_TOKEN}`,
      'Content-Type': 'application/json',
      ['x-use-cache']: "0"
    },
  })

  // avoid json parse error
  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/*");

  // check if image is not safe for work
  const file_url = await RawImage.fromBlob(blob)
  const classifier: any = await PipelineSingleton.getInstance();
  if (classifier) {
    const results = await classifier(file_url);
    if (results?.length && results[0]?.label === "nsfw") {
      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 });

}