Update lib/diffusion.js

lib/diffusion.js

@@ -1,43 +1,54 @@
 import axios from "axios";
 
 async function stablediff(prompt) {
-    const BASE_URL = "https://stabilityai-stable-diffusion-3-medium.hf.space/";
-    const session_hash = Math.random().toString(36).substring(2).slice(1);
-
-    const resp = await axios({
-        method: "POST",
-        url: BASE_URL + "queue/join?__theme=light",
-        data: {
-            data: [
-                prompt,
-                "",
-                0,
-                true,
-                1024,
-                1024,
-                5,
-                28
-            ],
-            event_data: null,
-            fn_index: 1,
-            session_hash,
-            trigger_id: 4
-        }
-    });
-
-    if(resp.data.event_id) {
-        const response = await axios({
-            method: "GET",
-            url: BASE_URL + "queue/data?session_hash=" + session_hash
+    try { 
+        const BASE_URL = "https://stabilityai-stable-diffusion-3-medium.hf.space/";
+        const session_hash = Math.random().toString(36).substring(2).slice(1);
+    
+        const resp = await axios({
+            method: "POST",
+            url: BASE_URL + "queue/join?__theme=light",
+            data: {
+                data: [
+                    prompt,
+                    "",
+                    0,
+                    true,
+                    1024,
+                    1024,
+                    5,
+                    28
+                ],
+                event_data: null,
+                fn_index: 1,
+                session_hash,
+                trigger_id: 4
+            }
         });
-        const splited = response.data.split("\n");
-        const processStartsLine = splited.find(line => line.includes('process_completed'));
-        const processStartsData = JSON.parse(processStartsLine.replace('data: ', ''));
-
-        if(processStartsData?.success) {
-            return processStartsData.output.data;
+    
+        if(resp.data.event_id) {
+            const response = await axios({
+                method: "GET",
+                url: BASE_URL + "queue/data?session_hash=" + session_hash
+            });
+            const splited = response.data.split("\n");
+            const processStartsLine = splited.find(line => line.includes('process_completed'));
+            const processStartsData = JSON.parse(processStartsLine.replace('data: ', ''));
+    
+            const processStartsLine2 = splited.find(line => line.includes('progress'));
+            const processStartsData2 = JSON.parse(processStartsLine.replace('data: ', ''));
+    
+            if(processStartsData?.success) {
+                return processStartsData.output.data;
+            } else if(processStartsData?.success) {
+                return processStartsData2.output.data;
+            }
         }
+    } catch(e) {
+        console.log("error:" + e.message);
+        return e.message;
     }
 }
 
+
 export { stablediff }