Browse Source

feat: 可设置是否转发上游mj图片地址

CaIon 1 year ago
parent
commit
1162683b4d

+ 1 - 1
constant/midjourney.go

@@ -1,8 +1,8 @@
 package constant
 
 var MjNotifyEnabled = false
-
 var MjModeClearEnabled = false
+var MjForwardUrlEnabled = true
 
 const (
 	MjErrorUnknown = 5

+ 8 - 2
controller/midjourney.go

@@ -10,11 +10,11 @@ import (
 	"log"
 	"net/http"
 	"one-api/common"
+	"one-api/constant"
 	"one-api/dto"
 	"one-api/model"
 	"one-api/service"
 	"strconv"
-	"strings"
 	"time"
 )
 
@@ -233,6 +233,12 @@ func GetAllMidjourney(c *gin.Context) {
 	if logs == nil {
 		logs = make([]*model.Midjourney, 0)
 	}
+	if constant.MjForwardUrlEnabled {
+		for i, midjourney := range logs {
+			midjourney.ImageUrl = common.ServerAddress + "/mj/image/" + midjourney.MjId
+			logs[i] = midjourney
+		}
+	}
 	c.JSON(200, gin.H{
 		"success": true,
 		"message": "",
@@ -259,7 +265,7 @@ func GetUserMidjourney(c *gin.Context) {
 	if logs == nil {
 		logs = make([]*model.Midjourney, 0)
 	}
-	if !strings.Contains(common.ServerAddress, "localhost") {
+	if constant.MjForwardUrlEnabled {
 		for i, midjourney := range logs {
 			midjourney.ImageUrl = common.ServerAddress + "/mj/image/" + midjourney.MjId
 			logs[i] = midjourney

+ 3 - 0
model/option.go

@@ -93,6 +93,7 @@ func InitOptionMap() {
 	common.OptionMap["DefaultCollapseSidebar"] = strconv.FormatBool(common.DefaultCollapseSidebar)
 	common.OptionMap["MjNotifyEnabled"] = strconv.FormatBool(constant.MjNotifyEnabled)
 	common.OptionMap["MjModeClearEnabled"] = strconv.FormatBool(constant.MjModeClearEnabled)
+	common.OptionMap["MjForwardUrlEnabled"] = strconv.FormatBool(constant.MjForwardUrlEnabled)
 	common.OptionMap["CheckSensitiveEnabled"] = strconv.FormatBool(constant.CheckSensitiveEnabled)
 	common.OptionMap["CheckSensitiveOnPromptEnabled"] = strconv.FormatBool(constant.CheckSensitiveOnPromptEnabled)
 	//common.OptionMap["CheckSensitiveOnCompletionEnabled"] = strconv.FormatBool(constant.CheckSensitiveOnCompletionEnabled)
@@ -198,6 +199,8 @@ func updateOptionMap(key string, value string) (err error) {
 			constant.MjNotifyEnabled = boolValue
 		case "MjModeClearEnabled":
 			constant.MjModeClearEnabled = boolValue
+		case "MjForwardUrlEnabled":
+			constant.MjForwardUrlEnabled = boolValue
 		case "CheckSensitiveEnabled":
 			constant.CheckSensitiveEnabled = boolValue
 		case "CheckSensitiveOnPromptEnabled":

+ 3 - 1
relay/relay-mj.go

@@ -110,11 +110,13 @@ func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjo
 	midjourneyTask.StartTime = originTask.StartTime
 	midjourneyTask.FinishTime = originTask.FinishTime
 	midjourneyTask.ImageUrl = ""
-	if originTask.ImageUrl != "" {
+	if originTask.ImageUrl != "" && constant.MjForwardUrlEnabled {
 		midjourneyTask.ImageUrl = common.ServerAddress + "/mj/image/" + originTask.MjId
 		if originTask.Status != "SUCCESS" {
 			midjourneyTask.ImageUrl += "?rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
 		}
+	} else {
+		midjourneyTask.ImageUrl = originTask.ImageUrl
 	}
 	midjourneyTask.Status = originTask.Status
 	midjourneyTask.FailReason = originTask.FailReason

+ 7 - 0
web/src/components/OperationSetting.js

@@ -39,6 +39,7 @@ const OperationSetting = () => {
     SensitiveWords: '',
     MjNotifyEnabled: '',
     MjModeClearEnabled: '',
+    MjForwardUrlEnabled: '',
     DrawingEnabled: '',
     DataExportEnabled: '',
     DataExportDefaultTime: 'hour',
@@ -322,6 +323,12 @@ const OperationSetting = () => {
               name='MjNotifyEnabled'
               onChange={handleInputChange}
             />
+            <Form.Checkbox
+              checked={inputs.MjForwardUrlEnabled === 'true'}
+              label='开启之后将上游地址替换为服务器地址'
+              name='MjForwardUrlEnabled'
+              onChange={handleInputChange}
+            />
             <Form.Checkbox
               checked={inputs.MjModeClearEnabled === 'true'}
               label='开启之后会清除用户提示词中的--fast、--relax以及--turbo参数'