Merge pull request #9 from joyhope/master

optional image save, depends on config
novicezk · May 15, 2023 · 1a374cc · 1a374cc
2 parents 21c5a2f + 2f2ca72
commit 1a374cc
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 27 deletions.
diff --git a/src/config.ts b/src/config.ts
@@ -12,5 +12,5 @@ export const config: IConfig = {
   mjProxyEndpoint: process.env.MJ_PROXY_ENDPOINT || "http://localhost:8022/mj",
   blockWords: process.env.BLOCK_WORDS?.split(",") || [],
   httpProxy: process.env.HTTP_PROXY || "",
-  imagesPath: process.env.IMAGE_PATH || "./data",
+  imagesPath: process.env.IMAGE_PATH || "",
 };
diff --git a/src/main.ts b/src/main.ts
@@ -1,10 +1,8 @@
 import { WechatyBuilder } from "wechaty";
-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 { config } from "./config.js";
 
 import express, { Application, Request, Response } from "express";
 
@@ -76,30 +74,15 @@ app.post("/notify", async (req: Request, res: Response): Promise<Response> => {
       if (action == 'UPSCALE') {
         await room.say(`@${userName} \n🎨 图片放大，用时: ${displayMilliseconds(time)}\n✨ ${description}`);
 
-        let image:FileBox;
-
-        if (config.httpProxy == "") {
-          image = FileBox.fromUrl(req.body.imageUrl);
-        } else {
-          const savedFileName = await downloadImage(req.body.imageUrl);
-          image = FileBox.fromFile(savedFileName);
-        }
-
+        const image =  await downloadImage(req.body.imageUrl);
         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`);
 
-        let image:FileBox;
-
-        if (config.httpProxy == "") {
-          image = FileBox.fromUrl(req.body.imageUrl);
-        } else {
-          const savedFileName = await downloadImage(req.body.imageUrl);
-          console.log(`saved`);
-          image = FileBox.fromFile(savedFileName);
-        }
+        const image = await downloadImage(req.body.imageUrl);
         room.say(image);
       }
     }

diff --git a/src/mj-api.ts b/src/mj-api.ts
@@ -1,5 +1,6 @@
 import { config } from "./config.js";
 import axios from 'axios';
+import { FileBox } from 'file-box';
 import type { AxiosInstance, AxiosRequestConfig, AxiosResponse } from 'axios';
 import { HttpsProxyAgent } from "https-proxy-agent"
 import * as fs from 'fs';
@@ -31,8 +32,7 @@ export async function submitTask(params: any): Promise<string> {
     }
 }
 
-
-export async function downloadImage(url: string): Promise<string> {
+export async function downloadImage(url: string): Promise<FileBox> {
     const response: AxiosResponse = await axios({
         method: 'GET',
         url: url,
@@ -42,10 +42,11 @@ export async function downloadImage(url: string): Promise<string> {
     });
 
     const filename = url.split('/')!.pop()!;
+    if (config.imagesPath!="") {
+      fs.writeFileSync(config.imagesPath + '/' + filename, response.data, 'binary');
+    }
+    const fileBuffer = Buffer.from(response.data, 'binary');
 
-    // Write the image data to a file
-    fs.writeFileSync(config.imagesPath + '/' + filename, response.data, 'binary');
+    return  FileBox.fromBuffer(fileBuffer, filename);
 
-    // Return the filename as a string
-    return filename;
 }