Merge pull request #53 from kammalshettys/error_check

exception handling in txtToImg service
a16z-infra · Jul 10, 2023 · 44db208 · 44db208
2 parents 312f507 + 2ac222f
commit 44db208
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 22 deletions.
diff --git a/src/app/api/txt2img/route.ts b/src/app/api/txt2img/route.ts
@@ -9,15 +9,19 @@ export async function POST(request: Request) {
   const replicate = new Replicate({
     auth: process.env.REPLICATE_API_TOKEN || "",
   });
+  try{
+    const output = await replicate.run(
+      "stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
+      {
+        input: {
+          prompt,
+        },
+      }
+    );
+    return NextResponse.json(output);
+  }
+  catch(error:any){
+    return NextResponse.json({error:error.message},{status:500});
+  }
 
-  const output = await replicate.run(
-    "stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
-    {
-      input: {
-        prompt,
-      },
-    }
-  );
-
-  return NextResponse.json(output);
 }
diff --git a/src/components/TextToImgModal.tsx b/src/components/TextToImgModal.tsx
@@ -15,21 +15,34 @@ export default function TextToImgModal({
 }) {
   const [imgSrc, setImgSrc] = useState("");
   const [loading, setLoading] = useState(false);
+  const [errorMsg, setErrorMessage] = useState("");
   const onSubmit = async (e: any) => {
     e.preventDefault();
     setLoading(true);
-    const response = await fetch("/api/txt2img", {
-      method: "POST",
-      body: JSON.stringify({
-        prompt: e.target.value,
-      }),
-      headers: {
-        "Content-Type": "application/json",
-      },
-    });
-    const data = await response.json();
-    setImgSrc(data[0]);
-    setLoading(false);
+    setErrorMessage("")
+    try {
+      const response = await fetch("/api/txt2img", {
+        method: "POST",
+        body: JSON.stringify({
+          prompt: e.target.value,
+        }),
+        headers: {
+          "Content-Type": "application/json",
+        },
+      });
+      const data = await response.json();
+      if (response.ok) {
+        setImgSrc(data[0]);
+        setLoading(false);
+      }
+      else {
+        throw Error(data.error);
+      }
+    }
+    catch (error: any) {
+      setErrorMessage(error.message);
+      setLoading(false);
+    }
   };
   return (
     <Transition.Root show={open} as={Fragment}>
@@ -117,6 +130,11 @@ export default function TextToImgModal({
                     </svg>
                   </p>
                 )}
+
+                {errorMsg ?
+                  (<p className="text-sm text-white">
+                    {errorMsg}
+                  </p>) : null}
               </Dialog.Panel>
             </Transition.Child>
           </div>