add admin rule

app/api/collections/[id]/route.ts

@@ -7,7 +7,13 @@ const prisma = new PrismaClient()
 export async function DELETE(request: Request, { params }: { params: { id: string } }) {
   const { headers } = request
 
-  if (!headers?.get('Authorization') || headers.get('Authorization') !== `Bearer ${process.env.HF_TOKEN}`) {
+  const staff_flag_id = headers.get('x-staff-flag-id') ?? undefined
+
+  // @ts-ignore
+  const HF_ADMIN = process?.env?.HF_ADMIN.split(',') ?? []
+  const is_admin = staff_flag_id ? HF_ADMIN.includes(staff_flag_id) : false
+
+  if (!is_admin) {
     return Response.json(
       {
         status: 401,

app/api/collections/[id]/visibility/route.ts

@@ -0,0 +1,55 @@
+import { PrismaClient } from '@prisma/client'
+
+const prisma = new PrismaClient()
+
+export async function PUT(request: Request, { params }: { params: { id: string } }) {
+  const { headers } = request
+
+  const staff_flag_id = headers.get('x-staff-flag-id') ?? undefined
+
+  // @ts-ignore
+  const HF_ADMIN = process?.env?.HF_ADMIN.split(',') ?? []
+  const is_admin = staff_flag_id ? HF_ADMIN.includes(staff_flag_id) : false
+
+  if (!is_admin) {
+    return Response.json(
+      {
+        status: 401,
+        ok: false,
+        message: 'Wrong castle fam :/'
+      }
+    )
+  }
+
+  const collection = await prisma.collection.findUnique({
+    where: {
+      id: parseInt(params.id)
+    }
+  })
+  if (!collection) {
+    return Response.json(
+      {
+        status: 404,
+        ok: false
+      }
+    )
+  }
+
+  const new_image = await prisma.collection.update({
+    where: {
+      id: parseInt(params.id)
+    },
+    data: {
+      is_visible: true
+    }
+  })
+
+  return Response.json(
+    {
+      image: new_image,
+      message: `Image ${params.id} has been updated`,
+      status: 200,
+      ok: true
+    }
+  )
+}

app/api/collections/route.ts

@@ -3,10 +3,19 @@ import { PrismaClient } from '@prisma/client'
 const prisma = new PrismaClient()
 
 export async function GET(request: Request) {
+  const { headers } = request
   const { searchParams } = new URL(request.url)
   const userId = searchParams.get('id') ?? undefined
   const page = searchParams.get('page') ? parseInt(searchParams.get('page') as string) : 0
 
+  const staff_flag_id = headers.get('x-staff-flag-id') ?? undefined
+
+  // @ts-ignore
+  const HF_ADMIN = process?.env?.HF_ADMIN.split(',') ?? []
+  const is_admin = staff_flag_id ? HF_ADMIN.includes(staff_flag_id) : false
+
+  console.log('is_admin', HF_ADMIN, is_admin, staff_flag_id)
+
   const collections = await prisma.collection.findMany({
     orderBy: {
       id: 'desc'
@@ -14,6 +23,9 @@ export async function GET(request: Request) {
     where: {
       userId: {
         equals: userId
+      },
+      is_visible: {
+        equals: is_admin ? undefined : true
       }
     },
     take: 15,

app/api/login/route.ts

@@ -2,7 +2,7 @@ export async function GET() {
   const REDIRECT_URI = `https://${process.env.SPACE_HOST}/login/callback`
   return Response.json(
     {
-      redirect: `https://huggingface.co/oauth/authorize?client_id=${process.env.OAUTH_CLIENT_ID}&redirect_uri=${REDIRECT_URI}&scope=openid%20profile&state=STATE&response_type=code`,
+      redirect: `https://huggingface.co/oauth/authorize?client_id=${process.env.OAUTH_CLIENT_ID}&redirect_uri=${REDIRECT_URI}&scope=openid%20profile%20email&state=STATE&response_type=code`,
       status: 200,
       ok: true
     }

app/api/me/route.ts

@@ -13,10 +13,16 @@ export async function GET() {
   })
 
   const res = await request.clone().json().catch(() => ({}));
+  // @ts-ignore
+  const HF_ADMIN = process?.env?.HF_ADMIN.split(',') ?? []
+  const is_admin = res?.sub ? HF_ADMIN.includes(res?.sub) : false
 
   return Response.json(
     {
-      user: res,
+      user: {
+        ...res,
+        is_admin,
+      },
       status: 200,
       ok: true
     }

app/api/route.ts

@@ -41,7 +41,7 @@ export async function POST(
   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 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 });

components/button/index.tsx

@@ -1,15 +1,17 @@
 import classNames from "classnames";
+import { MouseEvent } from "react";
 
 interface Props {
   children: React.ReactNode;
   disabled?: boolean;
-  theme?: "primary" | "secondary" | "white";
-  onClick?: () => void;
+  theme?: "primary" | "secondary" | "white" | "danger" | "success";
+  onClick?: (e: MouseEvent<HTMLButtonElement>) => void;
 }
 export const Button: React.FC<Props> = ({
   children,
   disabled,
   theme = "primary",
+  onClick,
   ...props
 }) => {
   return (
@@ -19,11 +21,22 @@ export const Button: React.FC<Props> = ({
         {
           "bg-primary text-white border-primary": theme === "primary",
           "bg-white text-gray-900 border-white": theme === "white",
+          "bg-red-500 text-white border-red-500": theme === "danger",
+          "bg-green-500 text-white border-green-500": theme === "success",
           "!bg-gray-400 !text-gray-600 !cursor-not-allowed !border-gray-400":
             disabled,
         }
       )}
       {...props}
+      onClick={
+        onClick
+          ? (e) => {
+              e.stopPropagation();
+              e.preventDefault();
+              onClick(e);
+            }
+          : undefined
+      }
     >
       {children}
     </button>

components/main/collections/collection.tsx

@@ -1,9 +1,11 @@
 import { useMemo } from "react";
 import { motion } from "framer-motion";
+import classNames from "classnames";
+import { AiFillEyeInvisible } from "react-icons/ai";
 
-import { Collection as CollectionType, Image } from "@/utils/type";
+import { Image } from "@/utils/type";
 import { useInputGeneration } from "@/components/main/hooks/useInputGeneration";
-import { arrayBufferToBase64 } from "@/utils";
+import { Button } from "@/components/button";
 
 interface Props {
   index: number;
@@ -63,8 +65,18 @@ export const Collection: React.FC<Props> = ({
           style={{
             backgroundImage: `url(https://huggingface.co/datasets/enzostvs/stable-diffusion-tpu-generations/resolve/main/images/${collection.file_name}.png)`,
           }}
-          className="rounded-[33px] bg-gray-950 bg-cover absolute top-0 left-0 w-full h-full z-[-1] transition-all duration-200 group-hover:scale-110 bg-center"
+          className={classNames(
+            "rounded-[33px] bg-gray-950 bg-cover absolute top-0 left-0 w-full h-full z-[-1] transition-all duration-200 group-hover:scale-110 bg-center",
+            {
+              "opacity-40": !collection.is_visible,
+            }
+          )}
         />
+        {!collection.is_visible && (
+          <div className="flex items-center justify-end px-5 py-6">
+            <AiFillEyeInvisible className="text-white text-2xl" />
+          </div>
+        )}
       </motion.div>
     </div>
   );

components/main/hooks/useCollections.ts

@@ -1,14 +1,13 @@
 import { useState } from "react";
 import { useQuery, useQueryClient } from "@tanstack/react-query";
-import { useLocalStorage, useUpdateEffect } from "react-use";
+import { useUpdateEffect } from "react-use";
 import _ from "lodash";
 
 import { useUser } from "@/utils/useUser";
 
 export const useCollections = (category: string) => {
   const [loading, setLoading] = useState(false);
-  const { user } = useUser();
-  const [myGenerationsId] = useLocalStorage<any>('my-own-generations', []);
+  const { user, loading: userLoading } = useUser();
 
   const client = useQueryClient();
   
@@ -24,6 +23,9 @@ export const useCollections = (category: string) => {
       queryParams.append('page', '0');
 
       const response = await fetch(`/api/collections?${queryParams.toString()}`, {
+        headers: {
+          'x-staff-flag-id': user?.sub
+        },
         method: "GET",
       })
 
@@ -38,6 +40,7 @@ export const useCollections = (category: string) => {
       };
     },
     {
+      enabled: !userLoading,
       refetchOnMount: false,
       refetchOnWindowFocus: false,
       refetchOnReconnect: false,      
@@ -51,6 +54,9 @@ export const useCollections = (category: string) => {
       queryParams.append('page', data?.pagination?.page,);
 
       const response = await fetch(`/api/collections?${queryParams.toString()}`, {
+        headers: {
+          'x-staff-flag-id': user?.sub
+        },
         method: "GET",
       })
 

components/modal/modal.tsx

@@ -1,9 +1,12 @@
 import { useMemo } from "react";
 import { motion } from "framer-motion";
 import Image from "next/image";
+import { BsFillTrashFill } from "react-icons/bs";
+import { AiFillCheckCircle } from "react-icons/ai";
 
-import { arrayBufferToBase64 } from "@/utils";
 import { useCollection } from "./useCollection";
+import { Button } from "../button";
+import { useUser } from "@/utils/useUser";
 
 interface Props {
   id: string;
@@ -30,16 +33,17 @@ const dropIn = {
 };
 
 export const Modal: React.FC<Props> = ({ id, onClose }) => {
-  const { collection } = useCollection(id);
-
-  if (!collection) return null;
+  const { collection, updateVisibility, remove } = useCollection(id);
+  const { user } = useUser();
 
   const formatDate = useMemo(() => {
     if (!collection) return;
-    const date = new Date(collection.createdAt);
+    const date = new Date(collection?.createdAt);
     return date.toLocaleDateString();
   }, [collection?.createdAt]);
 
+  console.log(user);
+
   return (
     <motion.div
       onClick={onClose}
@@ -50,14 +54,28 @@ export const Modal: React.FC<Props> = ({ id, onClose }) => {
     >
       <motion.div
         onClick={(e) => e.stopPropagation()}
-        className="max-w-2xl h-auto w-full z-[1] rounded-3xl overflow-hidden flex items-center justify-center flex-col gap-4 bg-white/30 backdrop-blur-sm px-2 pb-2 pt-2"
+        className="max-w-2xl h-auto w-full z-[1] rounded-3xl overflow-hidden relative flex items-center justify-center flex-col gap-4 bg-white/30 backdrop-blur-sm px-2 pb-2 pt-2"
         variants={dropIn}
         initial="hidden"
         animate="visible"
         exit="exit"
       >
+        {user?.is_admin && (
+          <div className="absolute p-2 rounded-full top-4 left-4 flex items-center justify-start gap-2 bg-black/20 backdrop-blur">
+            {!collection?.is_visible && (
+              <Button theme="white" onClick={updateVisibility}>
+                <AiFillCheckCircle />
+                Validate
+              </Button>
+            )}
+            <Button theme="danger" onClick={remove}>
+              <BsFillTrashFill />
+              Delete
+            </Button>
+          </div>
+        )}
         <Image
-          src={`https://huggingface.co/datasets/enzostvs/stable-diffusion-tpu-generations/resolve/main/images/${collection.file_name}.png?raw=true`}
+          src={`https://huggingface.co/datasets/enzostvs/stable-diffusion-tpu-generations/resolve/main/images/${collection?.file_name}.png?raw=true`}
           alt="Generated image"
           className="object-center object-contain w-full h-full rounded-2xl"
           width={1024}
@@ -66,13 +84,13 @@ export const Modal: React.FC<Props> = ({ id, onClose }) => {
         <div
           className="bg-cover bg-center w-full h-full rounded-2xl bg-no-repeat z-[-1] absolute top-0 left-0 opacity-90 blur-xl"
           style={{
-            backgroundImage: `url(https://huggingface.co/datasets/enzostvs/stable-diffusion-tpu-generations/resolve/main/images/${collection.file_name}.png?raw=true)`,
+            backgroundImage: `url(https://huggingface.co/datasets/enzostvs/stable-diffusion-tpu-generations/resolve/main/images/${collection?.file_name}.png?raw=true)`,
           }}
         />
         <div className="text-left w-full px-4 pb-3 pt-2">
           <p className="text-sm font-medium text-white/60 mb-1">{formatDate}</p>
           <p className="text-xl font-semibold text-white lowercase leading-snug">
-            {collection.prompt}
+            {collection?.prompt}
           </p>
         </div>
       </motion.div>

components/modal/useCollection.ts

@@ -1,9 +1,14 @@
 import { useMemo, useState } from "react"
 import { useQuery, useQueryClient } from "@tanstack/react-query"
 
-import { Collection } from "@/utils/type"
+import { Collection, Image } from "@/utils/type"
+import { useUser } from "@/utils/useUser"
+import { set } from "lodash"
 
 export const useCollection = (id?: string) => {
+  const { user } = useUser()
+  const [loading, setLoading] = useState(false)
+
   const { data: open } = useQuery(["modal"], () => {
     return null
   }, {
@@ -18,14 +23,67 @@ export const useCollection = (id?: string) => {
 
   const collection = useMemo(() => {
     const collections = client.getQueryData<Collection>(["collections"])
-    if (!collections?.images) return null
-
+    if (!collections?.images) {
+      setOpen(null)
+      return null
+    }
+ 
     return collections?.images?.find((collection) => collection.id === id)
-  }, [id])
+  }, [id, loading])
+
+  const updateVisibility = async () => {
+    setLoading(true)
+    const response = await fetch(`/api/collections/${collection?.id}/visibility`, {
+      method: "PUT",
+      headers: {
+        'x-staff-flag-id': user?.sub
+      }
+    })
+
+    const data = await response.json()
+    if (data.ok) {
+      client.setQueryData(["collections"], (old: any) => {
+        return {
+          ...old,
+          images: old.images.map((collection: Image) => {
+            if (collection.id === data.image.id) {
+              return data.image
+            }
+            return collection
+          })
+        }
+      })
+    }
+    setLoading(false)
+  }
+
+  const remove = async () => {
+    setLoading(true)
+    const response = await fetch(`/api/collections/${collection?.id}`, {
+      method: "DELETE",
+      headers: {
+        'x-staff-flag-id': user?.sub
+      }
+    })
+
+    const data = await response.json()
+    if (data.ok) {
+      client.setQueryData(["collections"], (old: any) => {
+        return {
+          ...old,
+          images: old.images.filter((col: Image) => col.id !== collection?.id)
+        }
+      })
+      setOpen(null)
+    }
+    setLoading(false)
+  }
 
   return {
     collection,
     open,
-    setOpen
+    setOpen,
+    updateVisibility,
+    remove
   }
 }

prisma/schema.prisma

@@ -4,13 +4,14 @@ generator client {
 
 datasource db {
   provider = "sqlite"
-  url      = "file://data/dev.db"
+  url      = "file:../data/dev.db"
 }
 
 model Collection {
-  id        Int      @id @default(autoincrement())
-  prompt    String
-  file_name String
-  createdAt DateTime @default(now())
-  userId    String?  @default("")
+  id         Int      @id @default(autoincrement())
+  prompt     String
+  file_name  String
+  is_visible Boolean  @default(false)
+  createdAt  DateTime @default(now())
+  userId     String?  @default("")
 }

utils/type.ts

@@ -12,6 +12,7 @@ export interface Image {
   file_name: string;
   prompt: string;
   createdAt: string;
+  is_visible: boolean;
   error?: string;
   loading?: boolean;
 }