From a32007c8fd85482256cf87566d8406cb6bec050b Mon Sep 17 00:00:00 2001
From: zjt <937178071@qq.com>
Date: Sat, 25 May 2024 22:31:49 +0800
Subject: [PATCH] get image

---
 src/biz/txt23d.ts  |  7 +++-
 src/biz/txt2img.ts | 95 ++++++++++++++++++++++++----------------------
 src/biz/upScale.ts |  7 +++-
 3 files changed, 62 insertions(+), 47 deletions(-)

diff --git a/src/biz/txt23d.ts b/src/biz/txt23d.ts
index 1665541..2cac4e0 100644
--- a/src/biz/txt23d.ts
+++ b/src/biz/txt23d.ts
@@ -20,7 +20,8 @@ const Txt23DHandler: RequestHandler<any, any> = async (ctx) => {
     const { prompt, url } = requestBody;
     const inputNode = selectNodeFromApiJSONbyID(txt23dApiFormatJSON, "33");
     inputNode.inputs.string = prompt;
-    ctx.body = {
+    try {
+           ctx.body = {
         url: ctx.body = baseUrl + '/view?filename=' + await new Promise((resolve, reject) => {
             const taskID = Math.random().toFixed(10);
             const ws = new WebSocket(`${baseWsUrl}/ws?clientId=${taskID}`);
@@ -52,5 +53,9 @@ const Txt23DHandler: RequestHandler<any, any> = async (ctx) => {
             }
         })
     }
+    } catch (error) {
+        console.log(error);
+    }
+ 
 }
 export default Txt23DHandler;
\ No newline at end of file
diff --git a/src/biz/txt2img.ts b/src/biz/txt2img.ts
index 4a5276b..5fd7012 100644
--- a/src/biz/txt2img.ts
+++ b/src/biz/txt2img.ts
@@ -22,58 +22,63 @@ const Txt2ImgHandler: RequestHandler<Txt2ImgRequest, any> = async (ctx, next) =>
     const { prompt, loraDetail, loraModel } = requestBody;
     const inputNode = selectNodeFromApiJSONbyID(txt2imgAPIformatJSON, "33");
     let temp = "";
-    let preViewBlobs:any = [];
+    let preViewBlobs: any = [];
     inputNode.inputs.string = prompt;
-    ctx.body = {
-        data: await new Promise((resolve, reject) => {
-            const taskID = Math.random().toFixed(10);
-            const ws = new WebSocket(`${baseWsUrl}/ws?clientId=${taskID}`);
-            ws.onopen = () => {
-                try {
-                    axios.post("/prompt", {
-                        client_id: taskID,
-                        prompt: txt2imgAPIformatJSON,
-                        // extra_data: txt2imgAPIformatExtraData
-                    });
-                    console.log(`ok`);
-                } catch (error) {
-                    console.log(error);
-                }
+    try {
+        ctx.body = {
+            data: await new Promise((resolve, reject) => {
+                const taskID = Math.random().toFixed(10);
+                const ws = new WebSocket(`${baseWsUrl}/ws?clientId=${taskID}`);
+                ws.onopen = () => {
+                    try {
+                        axios.post("/prompt", {
+                            client_id: taskID,
+                            prompt: txt2imgAPIformatJSON,
+                            // extra_data: txt2imgAPIformatExtraData
+                        });
+                        console.log(`ok`);
+                    } catch (error) {
+                        console.log(error);
+                    }
 
-                ws.onmessage = (event) => {
-                    if (typeof event.data === "string") {
-                        const { type, data } = JSON.parse(event.data);
-                        // if (type === "executed") {
-                        //     if (data.node === '94') {
-                        //         resolve(data.output.images[0].filename)
-                        //     }
-                        // }
-                        if (type === "executed") {
-                            if (data.node === '21') {
-                                temp = data.output.text[0];
-                            }
-                            if (data.node === '94') {
-                                console.log(data.output);
-                                const filePath = './data.json';
+                    ws.onmessage = (event) => {
+                        if (typeof event.data === "string") {
+                            const { type, data } = JSON.parse(event.data);
+                            // if (type === "executed") {
+                            //     if (data.node === '94') {
+                            //         resolve(data.output.images[0].filename)
+                            //     }
+                            // }
+                            if (type === "executed") {
+                                if (data.node === '21') {
+                                    temp = data.output.text[0];
+                                }
+                                if (data.node === '94') {
+                                    console.log(data.output);
+                                    const filePath = './data.json';
 
-                                readJSONFile(filePath)
-                                    .then((data: any) => {
-                                        console.log('读取的JSON数据:', data);
-                                        // 修改对象
-                                        data.push({ prompt: prompt, url: baseUrl + '/view?filename=' + data.output.images[0].filename, like: 0, index: data.length })
-                                        // 保存回JSON文件
-                                        saveJSONFile(filePath, data);
-                                    })
-                                resolve({ prompt: temp, url: baseUrl + '/view?filename=' + data.output.images[0].filename });
+                                    readJSONFile(filePath)
+                                        .then((data: any) => {
+                                            console.log('读取的JSON数据:', data);
+                                            // 修改对象
+                                            data.push({ prompt: prompt, url: baseUrl + '/view?filename=' + data.output.images[0].filename, like: 0, index: data.length })
+                                            // 保存回JSON文件
+                                            saveJSONFile(filePath, data);
+                                        })
+                                    resolve({ prompt: temp, url: baseUrl + '/view?filename=' + data.output.images[0].filename });
+                                }
+                            } else if (type === "executing") {
                             }
-                        }else if (type === "executing") {
+                        } else if (event.data instanceof Buffer) {
+
                         }
-                    } else if (event.data instanceof Buffer) {
-
                     }
                 }
-            }
-        })
+            })
+        }
+    } catch (error) {
+        console.log(error);
     }
+
 }
 export default Txt2ImgHandler;
\ No newline at end of file
diff --git a/src/biz/upScale.ts b/src/biz/upScale.ts
index fbbede0..1a03f15 100644
--- a/src/biz/upScale.ts
+++ b/src/biz/upScale.ts
@@ -21,7 +21,8 @@ const UpscaleHandler: RequestHandler<any ,any> = async (ctx) => {
     const { prompt, url } = requestBody;
     const inputNode = selectNodeFromApiJSONbyID(upScaleAPIFormatJSON, "12");
     inputNode.inputs.image = url;
-    ctx.body = {
+    try {
+           ctx.body = {
         url: baseUrl + '/view?filename=' + await new Promise((resolve, reject) => {
             const taskID = Math.random().toFixed(10);
             const ws = new WebSocket(`${baseWsUrl}/ws?clientId=${taskID}`);
@@ -51,6 +52,10 @@ const UpscaleHandler: RequestHandler<any ,any> = async (ctx) => {
                 }
             }
         })
+    } 
+    }catch (error) {
+        console.log(error);
     }
+
 }
 export default UpscaleHandler;
\ No newline at end of file