replace prisma blob saved with datasets

app/api/collections/route.ts

@@ -5,7 +5,7 @@ const prisma = new PrismaClient()
 export async function POST(request: Request) {
   const { ids, page } = await request.json()
 
-  const images = await prisma.image.findMany({
+  const collections = await prisma.collection.findMany({
     orderBy: {
       id: 'desc'
     },
@@ -18,11 +18,11 @@ export async function POST(request: Request) {
     skip: page * 15
   })
 
-  const total = await prisma.image.count()
+  const total = await prisma.collection.count()
 
   return Response.json(
     {
-      images, 
+      collections, 
       pagination: {
         total,
         page: page + 1,

app/api/route.ts

@@ -1,5 +1,7 @@
 import { PrismaClient } from '@prisma/client'
 
+import { UploaderDataset } from './uploader'
+
 const prisma = new PrismaClient()
 
 export async function POST(
@@ -43,11 +45,12 @@ export async function POST(
       if (isNSFW?.error && isNSFW?.estimated_time) {
         setTimeout(() => {
           checkIfIsNSFW(blob)
+          console.log('retry')
         }, isNSFW?.estimated_time * 100);
       } else resolve(isNSFW)
     })
   }
-
+  
   const isNSFW: any = await checkIfIsNSFW(blob)
   if (isNSFW?.error) return Response.json({ status: 500, ok: false, message: isNSFW?.error });
   if (isNSFW?.length) {
@@ -56,17 +59,19 @@ export async function POST(
       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({
+  const name = Date.now() + `-${inputs.replace(/[^a-zA-Z0-9]/g, '-').slice(0, 10).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,
-      blob: bytes,
+      file_name: name,
     },
   })
   
-  return Response.json({ blob: new_blob, status: 200, ok: true, headers });
+  return Response.json({ image: new_image, status: 200, ok: true, headers });
 
 }

app/api/uploader.ts

@@ -0,0 +1,30 @@
+import { uploadFiles } from "@huggingface/hub";
+import type { RepoDesignation, Credentials } from "@huggingface/hub";
+
+export const UploaderDataset = async (blob: Blob, name: string) => {
+  const repo: RepoDesignation = { type: "dataset", name: "enzostvs/stable-diffusion-tpu-generations" };
+  const credentials: Credentials = { accessToken: process.env.NEXT_PUBLIC_APP_HF_TOKEN as string };
+
+  const res: any = await uploadFiles({
+    repo,
+    credentials,
+    files: [
+      {
+        path: `images/${name}.png`,
+        content: blob,
+      },
+    ],
+  });
+
+  if (res?.error) return {
+    status: 500,
+    ok: false,
+    message: res?.error
+  };
+
+  return {
+    status: 200,
+    ok: true,
+  };
+
+}

components/main/collections/collection.tsx

@@ -20,12 +20,6 @@ export const Collection: React.FC<Props> = ({
 }) => {
   const { setPrompt } = useInputGeneration();
 
-  const bufferToBase64 = useMemo(() => {
-    const base64Flag = "data:image/jpeg;base64,";
-    const imageStr = arrayBufferToBase64(collection.blob.data);
-    return base64Flag + imageStr;
-  }, [collection]);
-
   const formatDate = useMemo(() => {
     const date = new Date(collection.createdAt);
     return date.toLocaleDateString();
@@ -66,9 +60,9 @@ export const Collection: React.FC<Props> = ({
         </div>
         <div
           style={{
-            backgroundImage: `url(${bufferToBase64})`,
+            backgroundImage: `url(https://huggingface.co/datasets/enzostvs/stable-diffusion-tpu-generations/resolve/main/images/${collection.file_name}.png)`,
           }}
-          className="rounded-[33px] bg-red-400 bg-cover absolute top-0 left-0 w-full h-full z-[-1] transition-all duration-200 group-hover:scale-110 bg-center"
+          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"
         />
       </motion.div>
     </div>

components/main/hooks/useCollections.ts

@@ -28,8 +28,9 @@ export const useCollections = (category: string) => {
       if (!response.ok) {
         throw new Error(data.message)
       }
+      console.log(data)
       return {
-        images: data?.images,
+        images: data?.collections,
         pagination: data?.pagination,
       };
     },

components/main/hooks/useInputGeneration.ts

@@ -41,16 +41,12 @@ export const useInputGeneration = () => {
       if (!hasMadeFirstGeneration) setFirstGenerationDone()
       client.setQueryData(["collections"], (old: any) => {
         return {
-          pagination: old.pagination,
+          pagination: old?.pagination,
           images: [{
             id,
             loading: true,
-            blob: {
-              type: "image/png",
-              data: new ArrayBuffer(0),
-            },
             prompt
-          }, ...old.images as Image[]]
+          }, ...old?.images as Image[]]
         }
       })
 
@@ -72,7 +68,7 @@ export const useInputGeneration = () => {
         newArray[index] = !data.ok ? {
           ...newArray[index],
           error: data.message
-        } : data?.blob as Image
+        } : data?.image as Image
 
         return {
           ...old,
@@ -82,8 +78,8 @@ export const useInputGeneration = () => {
 
       if (!data.ok) return null
 
-      setGenerationsId(myGenerationsId?.length ? [...myGenerationsId, data?.blob?.id] : [data?.blob?.id])
-      setOpen(data?.blob?.id)
+      setGenerationsId(myGenerationsId?.length ? [...myGenerationsId, data?.image?.id] : [data?.image?.id])
+      setOpen(data?.image?.id)
       return data ?? {}
     }
   )

components/modal/modal.tsx

@@ -34,13 +34,6 @@ export const Modal: React.FC<Props> = ({ id, onClose }) => {
 
   if (!collection) return null;
 
-  const bufferToBase64 = useMemo(() => {
-    if (!collection) return;
-    const base64Flag = "data:image/jpeg;base64,";
-    const imageStr = arrayBufferToBase64(collection.blob.data);
-    return base64Flag + imageStr;
-  }, [collection]);
-
   const formatDate = useMemo(() => {
     if (!collection) return;
     const date = new Date(collection.createdAt);
@@ -64,7 +57,7 @@ export const Modal: React.FC<Props> = ({ id, onClose }) => {
         exit="exit"
       >
         <Image
-          src={bufferToBase64 as string}
+          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}
@@ -73,7 +66,7 @@ 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(${bufferToBase64})`,
+            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">

next.config.js

@@ -4,6 +4,14 @@ const nextConfig = {
     serverActions: true,
     serverComponentsExternalPackages: ['sharp', 'onnxruntime-node'],
   },
+  images: {
+    remotePatterns: [
+      {
+        protocol: "https",
+        hostname: "huggingface.co",
+      },
+    ],
+  },
 }
 
 module.exports = nextConfig

package-lock.json

@@ -8,6 +8,7 @@
       "name": "hf-fast-image-generation",
       "version": "0.1.0",
       "dependencies": {
+        "@huggingface/hub": "^0.11.3",
         "@prisma/client": "^5.5.2",
         "@tanstack/react-query": "^4.32.6",
         "classnames": "^2.3.2",
@@ -162,6 +163,17 @@
         "node": "^12.22.0 || ^14.17.0 || >=16.0.0"
       }
     },
+    "node_modules/@huggingface/hub": {
+      "version": "0.11.3",
+      "resolved": "https://registry.npmjs.org/@huggingface/hub/-/hub-0.11.3.tgz",
+      "integrity": "sha512-GpdW09CQNgmZsSvVDlh7A1Vrr1NSfeEcAdtvRPaF9Cs6q2/6fywK0HInZHgrbPcXuawsnoIFvuePkboTol2Ttg==",
+      "dependencies": {
+        "hash-wasm": "^4.9.0"
+      },
+      "engines": {
+        "node": ">=18"
+      }
+    },
     "node_modules/@humanwhocodes/config-array": {
       "version": "0.11.13",
       "resolved": "https://registry.npmjs.org/@humanwhocodes/config-array/-/config-array-0.11.13.tgz",
@@ -2496,6 +2508,11 @@
         "url": "https://github.com/sponsors/ljharb"
       }
     },
+    "node_modules/hash-wasm": {
+      "version": "4.10.0",
+      "resolved": "https://registry.npmjs.org/hash-wasm/-/hash-wasm-4.10.0.tgz",
+      "integrity": "sha512-a0NjBNWjavvMalm/pPSEJ00MPDjRG8rv9D5BK7dBQTLGwAOVWqnTEUggaYs5szATB5UK5ULeIQr7QJXbczAZYA=="
+    },
     "node_modules/hasown": {
       "version": "2.0.0",
       "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.0.tgz",

package.json

@@ -9,6 +9,7 @@
     "lint": "next lint"
   },
   "dependencies": {
+    "@huggingface/hub": "^0.11.3",
     "@prisma/client": "^5.5.2",
     "@tanstack/react-query": "^4.32.6",
     "classnames": "^2.3.2",

prisma/migrations/20231027162901_/migration.sql

@@ -1,7 +0,0 @@
--- CreateTable
-CREATE TABLE "Image" (
-    "id" INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
-    "prompt" TEXT NOT NULL,
-    "blob" BLOB NOT NULL,
-    "createdAt" DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
-);

prisma/migrations/migration_lock.toml

@@ -1,3 +0,0 @@
-# Please do not edit this file manually
-# It should be added in your version-control system (i.e. Git)
-provider = "sqlite"

prisma/schema.prisma

@@ -7,9 +7,9 @@ datasource db {
   url      = "file://data/dev.db"
 }
 
-model Image {
+model Collection {
   id        Int      @id @default(autoincrement())
   prompt    String
-  blob      Bytes
+  file_name String
   createdAt DateTime @default(now())
 }

utils/type.ts

@@ -9,10 +9,7 @@ export interface Collection {
 
 export interface Image {
   id: string;
-  blob: {
-    type: string;
-    data: ArrayBuffer
-  };
+  file_name: string;
   prompt: string;
   createdAt: string;
   error?: string;