Merge pull request #114 from Calcium-Ion/midjourney-proxy-plus

feat: support midjourney-proxy-plus

controller/channel-billing.go

@@ -214,8 +214,8 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
 		return 0, errors.New("尚未实现")
 	case common.ChannelTypeCustom:
 		baseURL = channel.GetBaseURL()
-	case common.ChannelTypeOpenAISB:
-		return updateChannelOpenAISBBalance(channel)
+	//case common.ChannelTypeOpenAISB:
+	//	return updateChannelOpenAISBBalance(channel)
 	case common.ChannelTypeAIProxy:
 		return updateChannelAIProxyBalance(channel)
 	case common.ChannelTypeAPI2GPT:

controller/midjourney.go

@@ -10,145 +10,14 @@ import (
 	"log"
 	"net/http"
 	"one-api/common"
+	"one-api/dto"
 	"one-api/model"
-	relay2 "one-api/relay"
 	"one-api/service"
 	"strconv"
 	"strings"
 	"time"
 )
 
-/*func UpdateMidjourneyTask() {
-	//revocer
-	//imageModel := "midjourney"
-	ctx := context.TODO()
-	imageModel := "midjourney"
-	defer func() {
-		if err := recover(); err != nil {
-			log.Printf("UpdateMidjourneyTask panic: %v", err)
-		}
-	}()
-	for {
-		time.Sleep(time.Duration(15) * time.Second)
-		tasks := model.GetAllUnFinishTasks()
-		if len(tasks) != 0 {
-			common.LogInfo(ctx, fmt.Sprintf("检测到未完成的任务数有: %v", len(tasks)))
-			for _, task := range tasks {
-				common.LogInfo(ctx, fmt.Sprintf("未完成的任务信息: %v", task))
-				midjourneyChannel, err := model.GetChannelById(task.ChannelId, true)
-				if err != nil {
-					common.LogError(ctx, fmt.Sprintf("UpdateMidjourneyTask: %v", err))
-					task.FailReason = fmt.Sprintf("获取渠道信息失败，请联系管理员，渠道ID：%d", task.ChannelId)
-					task.Status = "FAILURE"
-					task.Progress = "100%"
-					err := task.Update()
-					if err != nil {
-						common.LogInfo(ctx, fmt.Sprintf("UpdateMidjourneyTask error: %v", err))
-						continue
-					}
-					continue
-				}
-				requestUrl := fmt.Sprintf("%s/mj/task/%s/fetch", *midjourneyChannel.BaseURL, task.MjId)
-				common.LogInfo(ctx, fmt.Sprintf("requestUrl: %s", requestUrl))
-
-				req, err := http.NewRequest("GET", requestUrl, bytes.NewBuffer([]byte("")))
-				if err != nil {
-					common.LogInfo(ctx, fmt.Sprintf("Get Task error: %v", err))
-					continue
-				}
-
-				// 设置超时时间
-				timeout := time.Second * 5
-				ctx, cancel := context.WithTimeout(context.Background(), timeout)
-
-				// 使用带有超时的 context 创建新的请求
-				req = req.WithContext(ctx)
-
-				req.Header.Set("Content-Type", "application/json")
-				//req.Header.Set("ApiKey", "Bearer midjourney-proxy")
-				req.Header.Set("mj-api-secret", midjourneyChannel.Key)
-				resp, err := httpClient.Do(req)
-				if err != nil {
-					log.Printf("UpdateMidjourneyTask error: %v", err)
-					continue
-				}
-				responseBody, err := io.ReadAll(resp.Body)
-				resp.Body.Close()
-				log.Printf("responseBody: %s", string(responseBody))
-				var responseItem Midjourney
-				// err = json.NewDecoder(resp.Body).Decode(&responseItem)
-				err = json.Unmarshal(responseBody, &responseItem)
-				if err != nil {
-					if strings.Contains(err.Error(), "cannot unmarshal number into Go struct field Midjourney.status of type string") {
-						var responseWithoutStatus MidjourneyWithoutStatus
-						var responseStatus MidjourneyStatus
-						err1 := json.Unmarshal(responseBody, &responseWithoutStatus)
-						err2 := json.Unmarshal(responseBody, &responseStatus)
-						if err1 == nil && err2 == nil {
-							jsonData, err3 := json.Marshal(responseWithoutStatus)
-							if err3 != nil {
-								log.Printf("UpdateMidjourneyTask error1: %v", err3)
-								continue
-							}
-							err4 := json.Unmarshal(jsonData, &responseStatus)
-							if err4 != nil {
-								log.Printf("UpdateMidjourneyTask error2: %v", err4)
-								continue
-							}
-							responseItem.Status = strconv.Itoa(responseStatus.Status)
-						} else {
-							log.Printf("UpdateMidjourneyTask error3: %v", err)
-							continue
-						}
-					} else {
-						log.Printf("UpdateMidjourneyTask error4: %v", err)
-						continue
-					}
-				}
-				task.Code = 1
-				task.Progress = responseItem.Progress
-				task.PromptEn = responseItem.PromptEn
-				task.State = responseItem.State
-				task.SubmitTime = responseItem.SubmitTime
-				task.StartTime = responseItem.StartTime
-				task.FinishTime = responseItem.FinishTime
-				task.ImageUrl = responseItem.ImageUrl
-				task.Status = responseItem.Status
-				task.FailReason = responseItem.FailReason
-				if task.Progress != "100%" && responseItem.FailReason != "" {
-					common.LogWarn(task.MjId + " 构建失败，" + task.FailReason)
-					task.Progress = "100%"
-					err = model.CacheUpdateUserQuota(task.UserId)
-					if err != nil {
-						log.Println("error update user quota cache: " + err.Error())
-					} else {
-						modelRatio := common.GetModelRatio(imageModel)
-						groupRatio := common.GetGroupRatio("default")
-						ratio := modelRatio * groupRatio
-						quota := int(ratio * 1 * 1000)
-						if quota != 0 {
-							err := model.IncreaseUserQuota(task.UserId, quota)
-							if err != nil {
-								log.Println("fail to increase user quota")
-							}
-							logContent := fmt.Sprintf("构图失败 %s，补偿 %s", task.MjId, common.LogQuota(quota))
-							model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
-						}
-					}
-				}
-
-				err = task.Update()
-				if err != nil {
-					log.Printf("UpdateMidjourneyTask error5: %v", err)
-				}
-				log.Printf("UpdateMidjourneyTask success: %v", task)
-				cancel()
-			}
-		}
-	}
-}
-*/
-
 func UpdateMidjourneyTaskBulk() {
 	//imageModel := "midjourney"
 	ctx := context.TODO()
@@ -228,12 +97,16 @@ func UpdateMidjourneyTaskBulk() {
 				common.LogError(ctx, fmt.Sprintf("Get Task Do req error: %v", err))
 				continue
 			}
+			if resp.StatusCode != http.StatusOK {
+				common.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
+				continue
+			}
 			responseBody, err := io.ReadAll(resp.Body)
 			if err != nil {
 				common.LogError(ctx, fmt.Sprintf("Get Task parse body error: %v", err))
 				continue
 			}
-			var responseItems []relay2.Midjourney
+			var responseItems []dto.MidjourneyDto
 			err = json.Unmarshal(responseBody, &responseItems)
 			if err != nil {
 				common.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody)))
@@ -245,10 +118,16 @@ func UpdateMidjourneyTaskBulk() {
 
 			for _, responseItem := range responseItems {
 				task := taskM[responseItem.MjId]
+
+				useTime := (time.Now().UnixNano() / int64(time.Millisecond)) - task.SubmitTime
+				// 如果时间超过一小时，且进度不是100%，则认为任务失败
+				if useTime > 3600000 && task.Progress != "100%" {
+					responseItem.FailReason = "上游任务超时（超过1小时）"
+					responseItem.Status = "FAILURE"
+				}
 				if !checkMjTaskNeedUpdate(task, responseItem) {
 					continue
 				}
-
 				task.Code = 1
 				task.Progress = responseItem.Progress
 				task.PromptEn = responseItem.PromptEn
@@ -259,6 +138,15 @@ func UpdateMidjourneyTaskBulk() {
 				task.ImageUrl = responseItem.ImageUrl
 				task.Status = responseItem.Status
 				task.FailReason = responseItem.FailReason
+				if responseItem.Properties != nil {
+					propertiesStr, _ := json.Marshal(responseItem.Properties)
+					task.Properties = string(propertiesStr)
+				}
+				if responseItem.Buttons != nil {
+					buttonStr, _ := json.Marshal(responseItem.Buttons)
+					task.Buttons = string(buttonStr)
+				}
+
 				if task.Progress != "100%" && responseItem.FailReason != "" {
 					common.LogInfo(ctx, task.MjId+" 构建失败，"+task.FailReason)
 					task.Progress = "100%"
@@ -286,7 +174,7 @@ func UpdateMidjourneyTaskBulk() {
 	}
 }
 
-func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask relay2.Midjourney) bool {
+func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask dto.MidjourneyDto) bool {
 	if oldTask.Code != 1 {
 		return true
 	}

controller/model.go

@@ -4,12 +4,13 @@ import (
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"net/http"
+	"one-api/constant"
 	"one-api/dto"
 	"one-api/model"
 	"one-api/relay"
 	"one-api/relay/channel/ai360"
 	"one-api/relay/channel/moonshot"
-	"one-api/relay/constant"
+	relayconstant "one-api/relay/constant"
 )
 
 // https://platform.openai.com/docs/api-reference/models/list
@@ -59,8 +60,8 @@ func init() {
 		IsBlocking:         false,
 	})
 	// https://platform.openai.com/docs/models/model-endpoint-compatibility
-	for i := 0; i < constant.APITypeDummy; i++ {
-		if i == constant.APITypeAIProxyLibrary {
+	for i := 0; i < relayconstant.APITypeDummy; i++ {
+		if i == relayconstant.APITypeAIProxyLibrary {
 			continue
 		}
 		adaptor := relay.GetAdaptor(i)
@@ -100,6 +101,17 @@ func init() {
 			Parent:     nil,
 		})
 	}
+	for modelName, _ := range constant.MidjourneyModel2Action {
+		openAIModels = append(openAIModels, OpenAIModels{
+			Id:         modelName,
+			Object:     "model",
+			Created:    1626777600,
+			OwnedBy:    "midjourney",
+			Permission: permission,
+			Root:       modelName,
+			Parent:     nil,
+		})
+	}
 	openAIModelsMap = make(map[string]OpenAIModels)
 	for _, model := range openAIModels {
 		openAIModelsMap[model.Id] = model

controller/relay.go

@@ -12,7 +12,6 @@ import (
 	relayconstant "one-api/relay/constant"
 	"one-api/service"
 	"strconv"
-	"strings"
 )
 
 func Relay(c *gin.Context) {
@@ -61,60 +60,35 @@ func Relay(c *gin.Context) {
 }
 
 func RelayMidjourney(c *gin.Context) {
-	relayMode := relayconstant.RelayModeUnknown
-	if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/imagine") {
-		relayMode = relayconstant.RelayModeMidjourneyImagine
-	} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/blend") {
-		relayMode = relayconstant.RelayModeMidjourneyBlend
-	} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/describe") {
-		relayMode = relayconstant.RelayModeMidjourneyDescribe
-	} else if strings.HasPrefix(c.Request.URL.Path, "/mj/notify") {
-		relayMode = relayconstant.RelayModeMidjourneyNotify
-	} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/change") {
-		relayMode = relayconstant.RelayModeMidjourneyChange
-	} else if strings.HasPrefix(c.Request.URL.Path, "/mj/submit/simple-change") {
-		relayMode = relayconstant.RelayModeMidjourneyChange
-	} else if strings.HasSuffix(c.Request.URL.Path, "/fetch") {
-		relayMode = relayconstant.RelayModeMidjourneyTaskFetch
-	} else if strings.HasSuffix(c.Request.URL.Path, "/list-by-condition") {
-		relayMode = relayconstant.RelayModeMidjourneyTaskFetchByCondition
-	}
-
+	relayMode := c.GetInt("relay_mode")
 	var err *dto.MidjourneyResponse
 	switch relayMode {
 	case relayconstant.RelayModeMidjourneyNotify:
 		err = relay.RelayMidjourneyNotify(c)
 	case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
 		err = relay.RelayMidjourneyTask(c, relayMode)
+	case relayconstant.RelayModeMidjourneyTaskImageSeed:
+		err = relay.RelayMidjourneyTaskImageSeed(c)
+	case relayconstant.RelayModeSwapFace:
+		err = relay.RelaySwapFace(c)
 	default:
 		err = relay.RelayMidjourneySubmit(c, relayMode)
 	}
 	//err = relayMidjourneySubmit(c, relayMode)
 	log.Println(err)
 	if err != nil {
-		retryTimesStr := c.Query("retry")
-		retryTimes, _ := strconv.Atoi(retryTimesStr)
-		if retryTimesStr == "" {
-			retryTimes = common.RetryTimes
-		}
-		if retryTimes > 0 {
-			c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s?retry=%d", c.Request.URL.Path, retryTimes-1))
-		} else {
-			if err.Code == 30 {
-				err.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
-			}
-			c.JSON(429, gin.H{
-				"error": fmt.Sprintf("%s %s", err.Description, err.Result),
-				"type":  "upstream_error",
-			})
+		statusCode := http.StatusBadRequest
+		if err.Code == 30 {
+			err.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
+			statusCode = http.StatusTooManyRequests
 		}
+		c.JSON(statusCode, gin.H{
+			"description": fmt.Sprintf("%s %s", err.Description, err.Result),
+			"type":        "upstream_error",
+			"code":        err.Code,
+		})
 		channelId := c.GetInt("channel_id")
 		common.SysError(fmt.Sprintf("relay error (channel #%d): %s", channelId, fmt.Sprintf("%s %s", err.Description, err.Result)))
-		//if shouldDisableChannel(&err.Error) {
-		//	channelId := c.GetInt("channel_id")
-		//	channelName := c.GetString("channel_name")
-		//	disableChannel(channelId, channelName, err.Result)
-		//};''''''''''''''''''''''''''''''''
 	}
 }