enzostvs's picture
enzostvs HF Staff
rate limiting
538551c
import { NextRequest } from 'next/server'
import { PrismaClient } from '@prisma/client'
import list_styles from "@/assets/list_styles.json"
import { UploaderDataset } from '../../utils/uploader'
import { isTextNSFW } from '@/utils/checker/prompt'
const prisma = new PrismaClient()
export async function POST(
request: NextRequest,
) {
const global_headers = {
Authorization: `Bearer ${process.env.HF_TOKEN}`,
'Content-Type': 'application/json',
['x-use-cache']: "0"
}
const { inputs, style, userId } = await request.json()
const { headers } = request
const findStyle = list_styles.find((item) => item.name === style)
const ip_address = headers.get("x-forwarded-for") ?? request.ip
// if (!headers.get('Authorization')) {
// const count = await prisma.collection.count({
// where: {
// ip_address,
// createdAt: {
// gte: new Date(Date.now() - 24 * 60 * 60 * 1000)
// }
// }
// })
// if (count > 5) return Response.json({ status: 429, ok: false, message: "You have reached the limit of 5 images per day." });
// }
const textIsNSFW = await isTextNSFW(inputs, global_headers)
if (textIsNSFW) return Response.json({ status: 401, ok: false, message: "Prompt doesn’t work, try another prompt" });
const response = await fetch(`${process.env.API_SDXL_URL}`, {
method: 'POST',
body: JSON.stringify({
prompt: findStyle?.prompt.replace("{prompt}", inputs) ?? inputs,
negative_prompt: findStyle?.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 base64Image = res.images[0];
const blob = await fetch(`data:image/png;base64,${base64Image}`).then((r) => r.blob());
// 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, '-').toLowerCase()}`
const { ok, message } = await UploaderDataset(blob, name)
if (!ok) return Response.json({ status: 500, ok: false, message });
const new_image = await prisma.collection.create({
data: {
prompt: inputs,
file_name: name,
userId: userId ?? "",
},
})
return Response.json({ image: new_image, status: 200, ok: true });
}