From 44af65b1581740509ea93a473c63e00f1e3b26b2 Mon Sep 17 00:00:00 2001
From: xswl <yanghao@ohdat.io>
Date: Mon, 17 Jul 2023 00:19:53 +0800
Subject: [PATCH] add ImageProxy

---
 .env.example             |  1 +
 src/discord.message.ts   | 15 ++++++++++-----
 src/discord.ws.ts        | 24 ++++++++++++++++++++----
 src/interfaces/config.ts |  3 +++
 4 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/.env.example b/.env.example
index 511aa3e..7eb1c93 100644
--- a/.env.example
+++ b/.env.example
@@ -3,3 +3,4 @@ SERVER_ID="Server id here"
 CHANNEL_ID="Channel in which commands are sent"
 SESSION_ID="session id here"
 HUGGINGFACE_TOKEN="verify human https://huggingface.co/docs/hub/security-tokens" 
+IMAGE_PROXY="image proxy url here"
\ No newline at end of file
diff --git a/src/discord.message.ts b/src/discord.message.ts
index e031537..188edfb 100644
--- a/src/discord.message.ts
+++ b/src/discord.message.ts
@@ -78,15 +78,20 @@ export class MidjourneyMessage {
           this.log("no attachment");
           break;
         }
-        const imageUrl = item.attachments[0].url;
-        //waiting
+        let uri = item.attachments[0].url;
+        if (this.config.ImageProxy !== "") {
+          uri = uri.replace(
+            "https://cdn.discordapp.com/",
+            this.config.ImageProxy
+          );
+        } //waiting
         if (
           item.attachments[0].filename.startsWith("grid") ||
           item.components.length === 0
         ) {
           this.log(`content`, item.content);
           const progress = this.content2progress(item.content);
-          loading?.(imageUrl, progress);
+          loading?.(uri, progress);
           break;
         }
         //finished
@@ -94,10 +99,10 @@ export class MidjourneyMessage {
         const msg: MJMessage = {
           content,
           id: item.id,
-          uri: imageUrl,
+          uri: uri,
           proxy_url: item.attachments[0].proxy_url,
           flags: item.flags,
-          hash: this.UriToHash(imageUrl),
+          hash: this.UriToHash(uri),
           progress: "done",
           options: formatOptions(item.components),
         };
diff --git a/src/discord.ws.ts b/src/discord.ws.ts
index 19f90da..fda4ab1 100644
--- a/src/discord.ws.ts
+++ b/src/discord.ws.ts
@@ -202,12 +202,18 @@ export class WsMessage {
           this.emit("settings", message);
           return;
         case "describe":
-          // console.log("describe", "meseesage", message);
+          let uri = embeds?.[0]?.image?.url;
+          if (this.config.ImageProxy !== "") {
+            uri = uri.replace(
+              "https://cdn.discordapp.com/",
+              this.config.ImageProxy
+            );
+          }
           const describe: MJDescribe = {
             id: id,
             flags: message.flags,
             descriptions: embeds?.[0]?.description.split("\n\n"),
-            uri: embeds?.[0]?.image?.url,
+            uri: uri,
             proxy_url: embeds?.[0]?.image?.proxy_url,
             options: formatOptions(components),
           };
@@ -378,13 +384,18 @@ export class WsMessage {
 
   private done(message: any) {
     const { content, id, attachments, components, flags } = message;
+    let uri = attachments[0].url;
+    if (this.config.ImageProxy !== "") {
+      uri = uri.replace("https://cdn.discordapp.com/", this.config.ImageProxy);
+    }
+
     const MJmsg: MJMessage = {
       id,
       flags,
       content,
       hash: uriToHash(attachments[0].url),
       progress: "done",
-      uri: attachments[0].url,
+      uri: uri,
       proxy_url: attachments[0].proxy_url,
       options: formatOptions(components),
     };
@@ -405,8 +416,13 @@ export class WsMessage {
     if (!attachments || attachments.length === 0) {
       return;
     }
+
+    let uri = attachments[0].url;
+    if (this.config.ImageProxy !== "") {
+      uri = uri.replace("https://cdn.discordapp.com/", this.config.ImageProxy);
+    }
     const MJmsg: MJMessage = {
-      uri: attachments[0].url,
+      uri: uri,
       proxy_url: attachments[0].proxy_url,
       content: content,
       flags: flags,
diff --git a/src/interfaces/config.ts b/src/interfaces/config.ts
index 232b4df..36acdb8 100644
--- a/src/interfaces/config.ts
+++ b/src/interfaces/config.ts
@@ -21,6 +21,7 @@ export interface MJConfig {
   fetch: FetchFn;
   ApiInterval: number;
   WebSocket: WebSocketCl;
+  ImageProxy: string;
 }
 export interface MJConfigParam {
   SalaiToken: string; //DISCORD_TOKEN
@@ -35,6 +36,7 @@ export interface MJConfigParam {
   HuggingFaceToken?: string; //HuggingFaceToken for verify human
   SessionId?: string;
   DiscordBaseUrl?: string;
+  ImageProxy?: string;
   WsBaseUrl?: string;
   fetch?: FetchFn; //Node.js<18 need node.fetch Or proxy
   WebSocket?: WebSocketCl; //isomorphic-ws Or proxy
@@ -50,6 +52,7 @@ export const DefaultMJConfig: MJConfig = {
   Limit: 50,
   Ws: true,
   MaxWait: 200,
+  ImageProxy: "",
   DiscordBaseUrl: "https://discord.com",
   WsBaseUrl: "wss://gateway.discord.gg?v=9&encoding=json&compress=gzip-stream",
   fetch: fetch,