From 5df09d5e2af74a099ca65fcfab9e9f0900a28905 Mon Sep 17 00:00:00 2001
From: lloydzhou <lloydzhou@qq.com>
Date: Mon, 15 Jul 2024 20:26:03 +0800
Subject: [PATCH] move code to utils/file

---
 app/components/sd.tsx     |  3 +--
 app/components/ui-lib.tsx | 26 --------------------------
 app/store/sd.ts           |  4 ++--
 app/utils/file.tsx        | 39 +++++++++++++++++++++++++++++++++++++--
 4 files changed, 40 insertions(+), 32 deletions(-)

diff --git a/app/components/sd.tsx b/app/components/sd.tsx
index 40e0c9abd..5d520bd1a 100644
--- a/app/components/sd.tsx
+++ b/app/components/sd.tsx
@@ -29,10 +29,9 @@ import {
   showConfirm,
   showImageModal,
   showModal,
-  IndexDBImage,
 } from "@/app/components/ui-lib";
 import { func } from "prop-types";
-import { useFileDB } from "@/app/utils/file";
+import { useFileDB, IndexDBImage } from "@/app/utils/file";
 
 function getSdTaskStatus(item: any) {
   let s: string;
diff --git a/app/components/ui-lib.tsx b/app/components/ui-lib.tsx
index a859b9695..ca13a6215 100644
--- a/app/components/ui-lib.tsx
+++ b/app/components/ui-lib.tsx
@@ -511,29 +511,3 @@ export function Selector<T>(props: {
     </div>
   );
 }
-
-export function IndexDBImage({ src, alt, onClick, db, className }) {
-  const [data, setData] = useState(src);
-  const imgId = useMemo(
-    () => src.replace("indexeddb://", "").split("@").pop(),
-    [src],
-  );
-  useEffect(() => {
-    db.getByID(imgId)
-      .then(({ data }) => {
-        setData(`data:image/png;base64,${data}`);
-      })
-      .catch((e) => {
-        setData(src);
-      });
-  }, [src, imgId]);
-
-  return (
-    <img
-      className={className}
-      src={data}
-      alt={alt}
-      onClick={(e) => onClick(data, e)}
-    />
-  );
-}
diff --git a/app/store/sd.ts b/app/store/sd.ts
index 77b20d001..70d44cac3 100644
--- a/app/store/sd.ts
+++ b/app/store/sd.ts
@@ -3,6 +3,7 @@ import { showToast } from "@/app/components/ui-lib";
 import { getHeaders } from "@/app/client/api";
 import { createPersistStore } from "@/app/utils/store";
 import { nanoid } from "nanoid";
+import { saveFileData } from "@/app/utils/file";
 
 export const useSdStore = createPersistStore<
   {
@@ -68,11 +69,10 @@ export const useSdStore = createPersistStore<
               return;
             }
             if (resData.finish_reason === "SUCCESS") {
-              db.add({ id: data.id, data: resData.image });
               this.updateDraw({
                 ...data,
                 status: "success",
-                img_data: `indexeddb://${StoreKey.File}@${data.id}`,
+                img_data: saveFileData(db, data.id, resData.image),
               });
             } else {
               this.updateDraw({
diff --git a/app/utils/file.tsx b/app/utils/file.tsx
index 0e77840e4..63039359f 100644
--- a/app/utils/file.tsx
+++ b/app/utils/file.tsx
@@ -1,4 +1,4 @@
-"use client";
+import { useState, useMemo, useEffect } from "react";
 import { initDB } from "react-indexed-db-hook";
 import { StoreKey } from "@/app/constant";
 import { useIndexedDB } from "react-indexed-db-hook";
@@ -23,9 +23,44 @@ export const FileDbConfig = {
 };
 
 export function FileDbInit() {
-  initDB(FileDbConfig);
+  if (typeof window !== "undefined") {
+    initDB(FileDbConfig);
+  }
 }
 
 export function useFileDB() {
   return useIndexedDB(StoreKey.File);
 }
+
+export function saveFileData(db, fileId, data) {
+  // save file content and return url start with `indexeddb://`
+  db.add({ id: fileId, data });
+  return `indexeddb://${StoreKey.File}@${fileId}`;
+}
+
+export async function getFileData(db, fileId, contentType = "image/png") {
+  const { data } = await db.getByID(fileId);
+  return `data:${contentType};base64,${data}`;
+}
+
+export function IndexDBImage({ src, alt, onClick, db, className }) {
+  const [data, setData] = useState(src);
+  const imgId = useMemo(
+    () => src.replace("indexeddb://", "").split("@").pop(),
+    [src],
+  );
+  useEffect(() => {
+    getFileData(db, imgId)
+      .then((data) => setData(data))
+      .catch((e) => setData(src));
+  }, [src, imgId]);
+
+  return (
+    <img
+      className={className}
+      src={data}
+      alt={alt}
+      onClick={(e) => onClick(data, e)}
+    />
+  );
+}