feat: ✨ Added NSFW filter

src/app/api/generate/route.ts

@@ -5,7 +5,7 @@ import { nanoid } from "nanoid";
 import { NextRequest, NextResponse } from "next/server";
 import Appwrite from "~/server/appwrite";
 
-import { ALLOWED_FILE_TYPES, MAX_FILE_SIZE } from "~/lib/constants";
+import { ALLOWED_FILE_TYPES, MAX_FILE_SIZE, SITE_URL } from "~/lib/constants";
 import replicate from "~/server/replicate";
 
 const ratelimit = new Ratelimit({
@@ -59,11 +59,6 @@ export async function POST(req: NextRequest) {
     access: "public",
   });
 
-  const SITE_URL =
-    process.env.NODE_ENV === "development"
-      ? "https://2d39-2a01-cb1c-8104-d500-a98d-c719-3dd3-d02c.ngrok-free.app"
-      : "https://micro-scale.vercel.app";
-
   try {
     const savedPrediction = await Appwrite.database.createDocument(
       process.env.NEXT_PUBLIC_APPWRITE_DATABASE_ID as string,
@@ -76,15 +71,15 @@ export async function POST(req: NextRequest) {
       }
     );
 
-    const webhook = new URL(`${SITE_URL}/api/webhook/save-image`);
+    const webhook = new URL(`${SITE_URL}/api/webhook/safety`);
     // setting an id to the webhook so we can identify the prediction later
     webhook.searchParams.set("id", savedPrediction.$id);
+    webhook.searchParams.set("blobURL", blob.url);
 
     const prediction = await replicate.predictions.create({
       version:
-        "32fdb2231d00a10d33754cc2ba794a2dfec94216579770785849ce6f149dbc69",
+        "88c3624a13d60bb5ecd0cb215e49e39d2a2135c211bcb94fc801d3def46803c4",
       input: {
-        scale: 4,
         image: blob.url,
       },
       webhook: webhook.toString(),

src/app/api/webhook/safety/route.ts

@@ -0,0 +1,55 @@
+import { NextRequest, NextResponse } from "next/server";
+import { SITE_URL } from "~/lib/constants";
+import appwrite from "~/server/appwrite";
+import replicate from "~/server/replicate";
+
+export async function POST(req: NextRequest) {
+  const body = (await req.json()) as { output: { nsfw_detected: boolean } };
+
+  const predictionId = req.nextUrl.searchParams.get("id");
+  const blobURL = req.nextUrl.searchParams.get("blobURL");
+
+  if (!predictionId) {
+    return NextResponse.json(
+      { message: "An error has occured." },
+      { status: 500 }
+    );
+  }
+
+  const isNSFW = body.output.nsfw_detected;
+
+  if (isNSFW) {
+    await appwrite.database.updateDocument(
+      process.env.NEXT_PUBLIC_APPWRITE_DATABASE_ID as string,
+      process.env.NEXT_PUBLIC_APPWRITE_PREDICTIONS_ID as string,
+      predictionId as string,
+      {
+        is_generated: false,
+        nsfw: true,
+      }
+    );
+
+    return NextResponse.json({ message: "NSFW" }, { status: 422 });
+  }
+
+  try {
+    const webhook = new URL(`${SITE_URL}/api/webhook/save-image`);
+    webhook.searchParams.set("id", predictionId);
+
+    await replicate.predictions.create({
+      version:
+        "32fdb2231d00a10d33754cc2ba794a2dfec94216579770785849ce6f149dbc69",
+      input: {
+        scale: 4,
+        image: blobURL,
+      },
+      webhook: webhook.toString(),
+      webhook_events_filter: ["completed"],
+    });
+
+    return NextResponse.json({ message: "OK" });
+  } catch (error) {
+    console.log(error);
+    return NextResponse.json({ message: "ERROR" }, { status: 500 });
+  }
+}

src/app/api/webhook/save-image/route.ts

@@ -22,12 +22,18 @@ type Prediction = {
     get: string;
   };
 };
+
 export async function POST(req: NextRequest) {
   const body = (await req.json()) as Prediction;
 
-  const params = new URL(req.url).searchParams;
+  const predictionId = req.nextUrl.searchParams.get("id");
 
-  const predictionId = params.get("id");
+  if (!predictionId) {
+    return NextResponse.json(
+      { message: "An error has occured." },
+      { status: 500 }
+    );
+  }
 
   try {
     await appwrite.database.updateDocument(

src/domain/generate/index.tsx

@@ -107,10 +107,28 @@ const Generate = () => {
         is_generated: boolean;
         output: string;
         input: string;
+        nsfw: boolean;
       };
       time: number;
     };
 
+    if (prediction.nsfw) {
+      reset();
+      setImages(initialState);
+      form.setError("file", {
+        type: "nsfw",
+        message: "NSFW images are not allowed.",
+      });
+      setPredictionId(null);
+      setIsPolling(false);
+
+      if (inputFileRef.current) {
+        inputFileRef.current.value = "";
+      }
+
+      return;
+    }
+
     if (prediction.is_generated) {
       setImages({
         original: prediction.input,
@@ -119,7 +137,7 @@ const Generate = () => {
       setPredictionId(null);
       setIsPolling(false);
     }
-  }, [pollingQuery.data, pollingQuery.data?.time]);
+  }, [pollingQuery.data, pollingQuery.data?.time, pollingQuery.data?.nsfw]);
 
   /**
    * ? Handle first visit

src/lib/constants.ts

@@ -1,3 +1,8 @@
 export const MAX_FILE_SIZE = 1024 * 1024 * 4; // 4MB
 
 export const ALLOWED_FILE_TYPES = ["image/png", "image/jpeg"];
+
+export const SITE_URL =
+  process.env.NODE_ENV === "development"
+    ? "https://fc90-2a01-cb1c-8104-d500-79ed-cd48-e2b4-7640.ngrok-free.app"
+    : "https://micro-scale.vercel.app";