fetch generated png file support though http proxy

itsharex · May 10, 2023 · ec57625 · ec57625
1 parent 75437e3
commit ec57625
Showing 6 changed files with 108 additions and 22 deletions.
diff --git a/.env.example b/.env.example
@@ -1,2 +1,4 @@
 MJ_PROXY_ENDPOINT="http://localhost:8080/mj"
 BLOCK_WORDS="测试,测试一"
+HTTP_PROXY="http://localhost:1080"
+IMAGE_PATH="."
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -16,6 +16,7 @@
     "dotenv": "^16.0.3",
     "express": "^4.18.2",
     "file-box": "^1.4.15",
+    "https-proxy-agent": "^6.1.0",
     "qrcode": "^1.5.1",
     "wechaty": "^1.20.2",
     "wechaty-puppet-wechat": "^1.18.4"
@@ -38,4 +39,4 @@
     "delay": 500
   },
   "type": "module"
-}
+}
diff --git a/src/config.ts b/src/config.ts
@@ -3,10 +3,14 @@ dotenv.config();
 
 export interface IConfig {
   mjProxyEndpoint: string;
-  blockWords: string[]
+  blockWords: string[];
+  httpProxy: string;
+  imagesPath: String;
 }
 
 export const config: IConfig = {
   mjProxyEndpoint: process.env.MJ_PROXY_ENDPOINT || "http://localhost:8022/mj",
-  blockWords: process.env.BLOCK_WORDS?.split(",") || []
+  blockWords: process.env.BLOCK_WORDS?.split(",") || [],
+  httpProxy: process.env.HTTP_PROXY || "",
+  imagesPath: process.env.IMAGE_PATH || ".",
 };
diff --git a/src/main.ts b/src/main.ts
@@ -3,6 +3,7 @@ import { FileBox } from 'file-box';
 import QRCode from "qrcode";
 import { Bot } from "./bot.js";
 import { displayMilliseconds } from "./utils.js";
+import { downloadImage } from "./mj-api.js";
 
 import express, { Application, Request, Response } from "express";
 
@@ -73,13 +74,25 @@ app.post("/notify", async (req: Request, res: Response): Promise<Response> => {
       const time = req.body.finishTime - req.body.submitTime;
       if (action == 'UPSCALE') {
         await room.say(`@${userName} \n🎨 图片放大，用时: ${displayMilliseconds(time)}\n✨ ${description}`);
+
+        /*
         const image = FileBox.fromUrl(req.body.imageUrl);
+        */
+        const savedFileName = await downloadImage(req.body.imageUrl);
+        const image = FileBox.fromFile(savedFileName);
         room.say(image);
       } else {
         const taskId = req.body.id;
         const prompt = req.body.prompt;
         await room.say(`@${userName} \n🎨 ${action == 'IMAGINE' ? '绘图' : '变换'}成功，用时 ${displayMilliseconds(time)}\n✨ Prompt: ${prompt}\n📨 任务ID: ${taskId}\n🪄 放大 U1～U4 ，变换 V1～V4\n✏️ 使用[/up 任务ID 操作]\n/up ${taskId} U1`);
+
+        /*
         const image = FileBox.fromUrl(req.body.imageUrl);
+        */
+        console.log(`${req.body.imageUrl}`);
+        const savedFileName = await downloadImage(req.body.imageUrl);
+        console.log(`saved`);
+        const image = FileBox.fromFile(savedFileName);
         room.say(image);
       }
     }

diff --git a/src/mj-api.ts b/src/mj-api.ts
@@ -1,6 +1,12 @@
+import { config } from "./config.js";
+import axios from 'axios';
+import type { AxiosInstance, AxiosRequestConfig, AxiosResponse } from 'axios';
+import { HttpsProxyAgent } from "https-proxy-agent"
+import * as fs from 'fs';
+
 import { Request } from "./request.js";
 
-const request = new Request({});
+const request =  new Request({})
 
 export async function submitTask(params: any): Promise<string> {
     let url = "/trigger/submit";
@@ -24,3 +30,22 @@ export async function submitTask(params: any): Promise<string> {
         return "系统异常，请稍后再试";
     }
 }
+
+
+export async function downloadImage(url: string): Promise<string> {
+  const response: AxiosResponse = await axios({
+    method: 'GET',
+    url: url,
+    responseType: 'arraybuffer',
+    httpsAgent: config.httpProxy!="" ?  new HttpsProxyAgent(config.httpProxy) : undefined,
+    timeout: 10000,
+  });
+
+  const filename = url.split('/')!.pop()!;
+
+  // Write the image data to a file
+  fs.writeFileSync(config.imagesPath+'/'+filename, response.data, 'binary');
+
+  // Return the filename as a string
+  return filename;
+}