mjproxy_handler.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679
  1. package relay
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "log"
  8. "net/http"
  9. "strconv"
  10. "strings"
  11. "time"
  12. "github.com/QuantumNous/new-api/common"
  13. "github.com/QuantumNous/new-api/constant"
  14. "github.com/QuantumNous/new-api/dto"
  15. "github.com/QuantumNous/new-api/model"
  16. relaycommon "github.com/QuantumNous/new-api/relay/common"
  17. relayconstant "github.com/QuantumNous/new-api/relay/constant"
  18. "github.com/QuantumNous/new-api/relay/helper"
  19. "github.com/QuantumNous/new-api/service"
  20. "github.com/QuantumNous/new-api/setting"
  21. "github.com/QuantumNous/new-api/setting/system_setting"
  22. "github.com/gin-gonic/gin"
  23. )
  24. func RelayMidjourneyImage(c *gin.Context) {
  25. taskId := c.Param("id")
  26. midjourneyTask := model.GetByOnlyMJId(taskId)
  27. if midjourneyTask == nil {
  28. c.JSON(400, gin.H{
  29. "error": "midjourney_task_not_found",
  30. })
  31. return
  32. }
  33. var httpClient *http.Client
  34. if channel, err := model.CacheGetChannel(midjourneyTask.ChannelId); err == nil {
  35. proxy := channel.GetSetting().Proxy
  36. if proxy != "" {
  37. if httpClient, err = service.NewProxyHttpClient(proxy); err != nil {
  38. c.JSON(400, gin.H{
  39. "error": "proxy_url_invalid",
  40. })
  41. return
  42. }
  43. }
  44. }
  45. if httpClient == nil {
  46. httpClient = service.GetHttpClient()
  47. }
  48. fetchSetting := system_setting.GetFetchSetting()
  49. if err := common.ValidateURLWithFetchSetting(midjourneyTask.ImageUrl, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil {
  50. c.JSON(http.StatusForbidden, gin.H{
  51. "error": fmt.Sprintf("request blocked: %v", err),
  52. })
  53. return
  54. }
  55. resp, err := httpClient.Get(midjourneyTask.ImageUrl)
  56. if err != nil {
  57. c.JSON(http.StatusInternalServerError, gin.H{
  58. "error": "http_get_image_failed",
  59. })
  60. return
  61. }
  62. defer resp.Body.Close()
  63. if resp.StatusCode != http.StatusOK {
  64. responseBody, _ := io.ReadAll(resp.Body)
  65. c.JSON(resp.StatusCode, gin.H{
  66. "error": string(responseBody),
  67. })
  68. return
  69. }
  70. // 从Content-Type头获取MIME类型
  71. contentType := resp.Header.Get("Content-Type")
  72. if contentType == "" {
  73. // 如果无法确定内容类型,则默认为jpeg
  74. contentType = "image/jpeg"
  75. }
  76. // 设置响应的内容类型
  77. c.Writer.Header().Set("Content-Type", contentType)
  78. // 将图片流式传输到响应体
  79. _, err = io.Copy(c.Writer, resp.Body)
  80. if err != nil {
  81. log.Println("Failed to stream image:", err)
  82. }
  83. return
  84. }
  85. func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse {
  86. var midjRequest dto.MidjourneyDto
  87. err := common.UnmarshalBodyReusable(c, &midjRequest)
  88. if err != nil {
  89. return &dto.MidjourneyResponse{
  90. Code: 4,
  91. Description: "bind_request_body_failed",
  92. Properties: nil,
  93. Result: "",
  94. }
  95. }
  96. midjourneyTask := model.GetByOnlyMJId(midjRequest.MjId)
  97. if midjourneyTask == nil {
  98. return &dto.MidjourneyResponse{
  99. Code: 4,
  100. Description: "midjourney_task_not_found",
  101. Properties: nil,
  102. Result: "",
  103. }
  104. }
  105. midjourneyTask.Progress = midjRequest.Progress
  106. midjourneyTask.PromptEn = midjRequest.PromptEn
  107. midjourneyTask.State = midjRequest.State
  108. midjourneyTask.SubmitTime = midjRequest.SubmitTime
  109. midjourneyTask.StartTime = midjRequest.StartTime
  110. midjourneyTask.FinishTime = midjRequest.FinishTime
  111. midjourneyTask.ImageUrl = midjRequest.ImageUrl
  112. midjourneyTask.VideoUrl = midjRequest.VideoUrl
  113. videoUrlsStr, _ := json.Marshal(midjRequest.VideoUrls)
  114. midjourneyTask.VideoUrls = string(videoUrlsStr)
  115. midjourneyTask.Status = midjRequest.Status
  116. midjourneyTask.FailReason = midjRequest.FailReason
  117. err = midjourneyTask.Update()
  118. if err != nil {
  119. return &dto.MidjourneyResponse{
  120. Code: 4,
  121. Description: "update_midjourney_task_failed",
  122. }
  123. }
  124. return nil
  125. }
  126. func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjourneyTask dto.MidjourneyDto) {
  127. midjourneyTask.MjId = originTask.MjId
  128. midjourneyTask.Progress = originTask.Progress
  129. midjourneyTask.PromptEn = originTask.PromptEn
  130. midjourneyTask.State = originTask.State
  131. midjourneyTask.SubmitTime = originTask.SubmitTime
  132. midjourneyTask.StartTime = originTask.StartTime
  133. midjourneyTask.FinishTime = originTask.FinishTime
  134. midjourneyTask.ImageUrl = ""
  135. if originTask.ImageUrl != "" && setting.MjForwardUrlEnabled {
  136. midjourneyTask.ImageUrl = system_setting.ServerAddress + "/mj/image/" + originTask.MjId
  137. if originTask.Status != "SUCCESS" {
  138. midjourneyTask.ImageUrl += "?rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
  139. }
  140. } else {
  141. midjourneyTask.ImageUrl = originTask.ImageUrl
  142. }
  143. if originTask.VideoUrl != "" {
  144. midjourneyTask.VideoUrl = originTask.VideoUrl
  145. }
  146. midjourneyTask.Status = originTask.Status
  147. midjourneyTask.FailReason = originTask.FailReason
  148. midjourneyTask.Action = originTask.Action
  149. midjourneyTask.Description = originTask.Description
  150. midjourneyTask.Prompt = originTask.Prompt
  151. if originTask.Buttons != "" {
  152. var buttons []dto.ActionButton
  153. err := json.Unmarshal([]byte(originTask.Buttons), &buttons)
  154. if err == nil {
  155. midjourneyTask.Buttons = buttons
  156. }
  157. }
  158. if originTask.VideoUrls != "" {
  159. var videoUrls []dto.ImgUrls
  160. err := json.Unmarshal([]byte(originTask.VideoUrls), &videoUrls)
  161. if err == nil {
  162. midjourneyTask.VideoUrls = videoUrls
  163. }
  164. }
  165. if originTask.Properties != "" {
  166. var properties dto.Properties
  167. err := json.Unmarshal([]byte(originTask.Properties), &properties)
  168. if err == nil {
  169. midjourneyTask.Properties = &properties
  170. }
  171. }
  172. return
  173. }
  174. func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyResponse {
  175. var swapFaceRequest dto.SwapFaceRequest
  176. err := common.UnmarshalBodyReusable(c, &swapFaceRequest)
  177. if err != nil {
  178. return service.MidjourneyErrorWrapper(constant.MjRequestError, "bind_request_body_failed")
  179. }
  180. info.InitChannelMeta(c)
  181. if swapFaceRequest.SourceBase64 == "" || swapFaceRequest.TargetBase64 == "" {
  182. return service.MidjourneyErrorWrapper(constant.MjRequestError, "sour_base64_and_target_base64_is_required")
  183. }
  184. modelName := service.CovertMjpActionToModelName(constant.MjActionSwapFace)
  185. priceData, err := helper.ModelPriceHelperPerCall(c, info)
  186. if err != nil {
  187. return &dto.MidjourneyResponse{
  188. Code: 4,
  189. Description: err.Error(),
  190. }
  191. }
  192. userQuota, err := model.GetUserQuota(info.UserId, false)
  193. if err != nil {
  194. return &dto.MidjourneyResponse{
  195. Code: 4,
  196. Description: err.Error(),
  197. }
  198. }
  199. if userQuota-priceData.Quota < 0 {
  200. return &dto.MidjourneyResponse{
  201. Code: 4,
  202. Description: "quota_not_enough",
  203. }
  204. }
  205. requestURL := getMjRequestPath(c.Request.URL.String())
  206. baseURL := c.GetString("base_url")
  207. fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
  208. mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
  209. if err != nil {
  210. return &mjResp.Response
  211. }
  212. defer func() {
  213. if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 {
  214. err := service.PostConsumeQuota(info, priceData.Quota, 0, true)
  215. if err != nil {
  216. common.SysLog("error consuming token remain quota: " + err.Error())
  217. }
  218. tokenName := c.GetString("token_name")
  219. logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace)
  220. other := service.GenerateMjOtherInfo(info, priceData)
  221. model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{
  222. ChannelId: info.ChannelId,
  223. ModelName: modelName,
  224. TokenName: tokenName,
  225. Quota: priceData.Quota,
  226. Content: logContent,
  227. TokenId: info.TokenId,
  228. Group: info.UsingGroup,
  229. Other: other,
  230. })
  231. model.UpdateUserUsedQuotaAndRequestCount(info.UserId, priceData.Quota)
  232. model.UpdateChannelUsedQuota(info.ChannelId, priceData.Quota)
  233. }
  234. }()
  235. midjResponse := &mjResp.Response
  236. midjourneyTask := &model.Midjourney{
  237. UserId: info.UserId,
  238. Code: midjResponse.Code,
  239. Action: constant.MjActionSwapFace,
  240. MjId: midjResponse.Result,
  241. Prompt: "InsightFace",
  242. PromptEn: "",
  243. Description: midjResponse.Description,
  244. State: "",
  245. SubmitTime: info.StartTime.UnixNano() / int64(time.Millisecond),
  246. StartTime: time.Now().UnixNano() / int64(time.Millisecond),
  247. FinishTime: 0,
  248. ImageUrl: "",
  249. Status: "",
  250. Progress: "0%",
  251. FailReason: "",
  252. ChannelId: c.GetInt("channel_id"),
  253. Quota: priceData.Quota,
  254. }
  255. err = midjourneyTask.Insert()
  256. if err != nil {
  257. return service.MidjourneyErrorWrapper(constant.MjRequestError, "insert_midjourney_task_failed")
  258. }
  259. c.Writer.WriteHeader(mjResp.StatusCode)
  260. respBody, err := json.Marshal(midjResponse)
  261. if err != nil {
  262. return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
  263. }
  264. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  265. if err != nil {
  266. return service.MidjourneyErrorWrapper(constant.MjRequestError, "copy_response_body_failed")
  267. }
  268. return nil
  269. }
  270. func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse {
  271. taskId := c.Param("id")
  272. userId := c.GetInt("id")
  273. originTask := model.GetByMJId(userId, taskId)
  274. if originTask == nil {
  275. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_no_found")
  276. }
  277. channel, err := model.GetChannelById(originTask.ChannelId, true)
  278. if err != nil {
  279. return service.MidjourneyErrorWrapper(constant.MjRequestError, "get_channel_info_failed")
  280. }
  281. if channel.Status != common.ChannelStatusEnabled {
  282. return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用")
  283. }
  284. c.Set("channel_id", originTask.ChannelId)
  285. c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
  286. requestURL := getMjRequestPath(c.Request.URL.String())
  287. fullRequestURL := fmt.Sprintf("%s%s", channel.GetBaseURL(), requestURL)
  288. midjResponseWithStatus, _, err := service.DoMidjourneyHttpRequest(c, time.Second*30, fullRequestURL)
  289. if err != nil {
  290. return &midjResponseWithStatus.Response
  291. }
  292. midjResponse := &midjResponseWithStatus.Response
  293. c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
  294. respBody, err := json.Marshal(midjResponse)
  295. if err != nil {
  296. return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
  297. }
  298. service.IOCopyBytesGracefully(c, nil, respBody)
  299. return nil
  300. }
  301. func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse {
  302. userId := c.GetInt("id")
  303. var err error
  304. var respBody []byte
  305. switch relayMode {
  306. case relayconstant.RelayModeMidjourneyTaskFetch:
  307. taskId := c.Param("id")
  308. originTask := model.GetByMJId(userId, taskId)
  309. if originTask == nil {
  310. return &dto.MidjourneyResponse{
  311. Code: 4,
  312. Description: "task_no_found",
  313. }
  314. }
  315. midjourneyTask := coverMidjourneyTaskDto(c, originTask)
  316. respBody, err = json.Marshal(midjourneyTask)
  317. if err != nil {
  318. return &dto.MidjourneyResponse{
  319. Code: 4,
  320. Description: "unmarshal_response_body_failed",
  321. }
  322. }
  323. case relayconstant.RelayModeMidjourneyTaskFetchByCondition:
  324. var condition = struct {
  325. IDs []string `json:"ids"`
  326. }{}
  327. err = c.BindJSON(&condition)
  328. if err != nil {
  329. return &dto.MidjourneyResponse{
  330. Code: 4,
  331. Description: "do_request_failed",
  332. }
  333. }
  334. var tasks []dto.MidjourneyDto
  335. if len(condition.IDs) != 0 {
  336. originTasks := model.GetByMJIds(userId, condition.IDs)
  337. for _, originTask := range originTasks {
  338. midjourneyTask := coverMidjourneyTaskDto(c, originTask)
  339. tasks = append(tasks, midjourneyTask)
  340. }
  341. }
  342. if tasks == nil {
  343. tasks = make([]dto.MidjourneyDto, 0)
  344. }
  345. respBody, err = json.Marshal(tasks)
  346. if err != nil {
  347. return &dto.MidjourneyResponse{
  348. Code: 4,
  349. Description: "unmarshal_response_body_failed",
  350. }
  351. }
  352. }
  353. c.Writer.Header().Set("Content-Type", "application/json")
  354. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  355. if err != nil {
  356. return &dto.MidjourneyResponse{
  357. Code: 4,
  358. Description: "copy_response_body_failed",
  359. }
  360. }
  361. return nil
  362. }
  363. func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dto.MidjourneyResponse {
  364. consumeQuota := true
  365. var midjRequest dto.MidjourneyRequest
  366. err := common.UnmarshalBodyReusable(c, &midjRequest)
  367. if err != nil {
  368. return service.MidjourneyErrorWrapper(constant.MjRequestError, "bind_request_body_failed")
  369. }
  370. relayInfo.InitChannelMeta(c)
  371. if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyAction { // midjourney plus,需要从customId中获取任务信息
  372. mjErr := service.CoverPlusActionToNormalAction(&midjRequest)
  373. if mjErr != nil {
  374. return mjErr
  375. }
  376. relayInfo.RelayMode = relayconstant.RelayModeMidjourneyChange
  377. }
  378. if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyVideo {
  379. midjRequest.Action = constant.MjActionVideo
  380. }
  381. if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyImagine { //绘画任务,此类任务可重复
  382. if midjRequest.Prompt == "" {
  383. return service.MidjourneyErrorWrapper(constant.MjRequestError, "prompt_is_required")
  384. }
  385. midjRequest.Action = constant.MjActionImagine
  386. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyDescribe { //按图生文任务,此类任务可重复
  387. midjRequest.Action = constant.MjActionDescribe
  388. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyEdits { //编辑任务,此类任务可重复
  389. midjRequest.Action = constant.MjActionEdits
  390. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyShorten { //缩短任务,此类任务可重复,plus only
  391. midjRequest.Action = constant.MjActionShorten
  392. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyBlend { //绘画任务,此类任务可重复
  393. midjRequest.Action = constant.MjActionBlend
  394. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyUpload { //绘画任务,此类任务可重复
  395. midjRequest.Action = constant.MjActionUpload
  396. } else if midjRequest.TaskId != "" { //放大、变换任务,此类任务,如果重复且已有结果,远端api会直接返回最终结果
  397. mjId := ""
  398. if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyChange {
  399. if midjRequest.TaskId == "" {
  400. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_id_is_required")
  401. } else if midjRequest.Action == "" {
  402. return service.MidjourneyErrorWrapper(constant.MjRequestError, "action_is_required")
  403. } else if midjRequest.Index == 0 {
  404. return service.MidjourneyErrorWrapper(constant.MjRequestError, "index_is_required")
  405. }
  406. //action = midjRequest.Action
  407. mjId = midjRequest.TaskId
  408. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneySimpleChange {
  409. if midjRequest.Content == "" {
  410. return service.MidjourneyErrorWrapper(constant.MjRequestError, "content_is_required")
  411. }
  412. params := service.ConvertSimpleChangeParams(midjRequest.Content)
  413. if params == nil {
  414. return service.MidjourneyErrorWrapper(constant.MjRequestError, "content_parse_failed")
  415. }
  416. mjId = params.TaskId
  417. midjRequest.Action = params.Action
  418. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyModal {
  419. //if midjRequest.MaskBase64 == "" {
  420. // return service.MidjourneyErrorWrapper(constant.MjRequestError, "mask_base64_is_required")
  421. //}
  422. mjId = midjRequest.TaskId
  423. midjRequest.Action = constant.MjActionModal
  424. } else if relayInfo.RelayMode == relayconstant.RelayModeMidjourneyVideo {
  425. midjRequest.Action = constant.MjActionVideo
  426. if midjRequest.TaskId == "" {
  427. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_id_is_required")
  428. } else if midjRequest.Action == "" {
  429. return service.MidjourneyErrorWrapper(constant.MjRequestError, "action_is_required")
  430. }
  431. mjId = midjRequest.TaskId
  432. }
  433. originTask := model.GetByMJId(relayInfo.UserId, mjId)
  434. if originTask == nil {
  435. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_not_found")
  436. } else { //原任务的Status=SUCCESS,则可以做放大UPSCALE、变换VARIATION等动作,此时必须使用原来的请求地址才能正确处理
  437. if setting.MjActionCheckSuccessEnabled {
  438. if originTask.Status != "SUCCESS" && relayInfo.RelayMode != relayconstant.RelayModeMidjourneyModal {
  439. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_status_not_success")
  440. }
  441. }
  442. channel, err := model.GetChannelById(originTask.ChannelId, true)
  443. if err != nil {
  444. return service.MidjourneyErrorWrapper(constant.MjRequestError, "get_channel_info_failed")
  445. }
  446. if channel.Status != common.ChannelStatusEnabled {
  447. return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用")
  448. }
  449. c.Set("base_url", channel.GetBaseURL())
  450. c.Set("channel_id", originTask.ChannelId)
  451. c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
  452. log.Printf("检测到此操作为放大、变换、重绘,获取原channel信息: %s,%s", strconv.Itoa(originTask.ChannelId), channel.GetBaseURL())
  453. }
  454. midjRequest.Prompt = originTask.Prompt
  455. //if channelType == common.ChannelTypeMidjourneyPlus {
  456. // // plus
  457. //} else {
  458. // // 普通版渠道
  459. //
  460. //}
  461. }
  462. if midjRequest.Action == constant.MjActionInPaint || midjRequest.Action == constant.MjActionCustomZoom {
  463. consumeQuota = false
  464. }
  465. //baseURL := common.ChannelBaseURLs[channelType]
  466. requestURL := getMjRequestPath(c.Request.URL.String())
  467. baseURL := c.GetString("base_url")
  468. //midjRequest.NotifyHook = "http://127.0.0.1:3000/mj/notify"
  469. fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
  470. modelName := service.CovertMjpActionToModelName(midjRequest.Action)
  471. priceData, err := helper.ModelPriceHelperPerCall(c, relayInfo)
  472. if err != nil {
  473. return &dto.MidjourneyResponse{
  474. Code: 4,
  475. Description: err.Error(),
  476. }
  477. }
  478. userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
  479. if err != nil {
  480. return &dto.MidjourneyResponse{
  481. Code: 4,
  482. Description: err.Error(),
  483. }
  484. }
  485. if consumeQuota && userQuota-priceData.Quota < 0 {
  486. return &dto.MidjourneyResponse{
  487. Code: 4,
  488. Description: "quota_not_enough",
  489. }
  490. }
  491. midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
  492. if err != nil {
  493. return &midjResponseWithStatus.Response
  494. }
  495. midjResponse := &midjResponseWithStatus.Response
  496. defer func() {
  497. if consumeQuota && midjResponseWithStatus.StatusCode == 200 {
  498. err := service.PostConsumeQuota(relayInfo, priceData.Quota, 0, true)
  499. if err != nil {
  500. common.SysLog("error consuming token remain quota: " + err.Error())
  501. }
  502. tokenName := c.GetString("token_name")
  503. logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s,ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result)
  504. other := service.GenerateMjOtherInfo(relayInfo, priceData)
  505. model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{
  506. ChannelId: relayInfo.ChannelId,
  507. ModelName: modelName,
  508. TokenName: tokenName,
  509. Quota: priceData.Quota,
  510. Content: logContent,
  511. TokenId: relayInfo.TokenId,
  512. Group: relayInfo.UsingGroup,
  513. Other: other,
  514. })
  515. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, priceData.Quota)
  516. model.UpdateChannelUsedQuota(relayInfo.ChannelId, priceData.Quota)
  517. }
  518. }()
  519. // 文档:https://github.com/novicezk/midjourney-proxy/blob/main/docs/api.md
  520. //1-提交成功
  521. // 21-任务已存在(处理中或者有结果了) {"code":21,"description":"任务已存在","result":"0741798445574458","properties":{"status":"SUCCESS","imageUrl":"https://xxxx"}}
  522. // 22-排队中 {"code":22,"description":"排队中,前面还有1个任务","result":"0741798445574458","properties":{"numberOfQueues":1,"discordInstanceId":"1118138338562560102"}}
  523. // 23-队列已满,请稍后再试 {"code":23,"description":"队列已满,请稍后尝试","result":"14001929738841620","properties":{"discordInstanceId":"1118138338562560102"}}
  524. // 24-prompt包含敏感词 {"code":24,"description":"可能包含敏感词","properties":{"promptEn":"nude body","bannedWord":"nude"}}
  525. // other: 提交错误,description为错误描述
  526. midjourneyTask := &model.Midjourney{
  527. UserId: relayInfo.UserId,
  528. Code: midjResponse.Code,
  529. Action: midjRequest.Action,
  530. MjId: midjResponse.Result,
  531. Prompt: midjRequest.Prompt,
  532. PromptEn: "",
  533. Description: midjResponse.Description,
  534. State: "",
  535. SubmitTime: time.Now().UnixNano() / int64(time.Millisecond),
  536. StartTime: 0,
  537. FinishTime: 0,
  538. ImageUrl: "",
  539. Status: "",
  540. Progress: "0%",
  541. FailReason: "",
  542. ChannelId: c.GetInt("channel_id"),
  543. Quota: priceData.Quota,
  544. }
  545. if midjResponse.Code == 3 {
  546. //无实例账号自动禁用渠道(No available account instance)
  547. channel, err := model.GetChannelById(midjourneyTask.ChannelId, true)
  548. if err != nil {
  549. common.SysLog("get_channel_null: " + err.Error())
  550. }
  551. if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
  552. model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance")
  553. }
  554. }
  555. if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 {
  556. //非1-提交成功,21-任务已存在和22-排队中,则记录错误原因
  557. midjourneyTask.FailReason = midjResponse.Description
  558. consumeQuota = false
  559. }
  560. if midjResponse.Code == 21 { //21-任务已存在(处理中或者有结果了)
  561. // 将 properties 转换为一个 map
  562. properties, ok := midjResponse.Properties.(map[string]interface{})
  563. if ok {
  564. imageUrl, ok1 := properties["imageUrl"].(string)
  565. status, ok2 := properties["status"].(string)
  566. if ok1 && ok2 {
  567. midjourneyTask.ImageUrl = imageUrl
  568. midjourneyTask.Status = status
  569. if status == "SUCCESS" {
  570. midjourneyTask.Progress = "100%"
  571. midjourneyTask.StartTime = time.Now().UnixNano() / int64(time.Millisecond)
  572. midjourneyTask.FinishTime = time.Now().UnixNano() / int64(time.Millisecond)
  573. midjResponse.Code = 1
  574. }
  575. }
  576. }
  577. //修改返回值
  578. if midjRequest.Action != constant.MjActionInPaint && midjRequest.Action != constant.MjActionCustomZoom {
  579. newBody := strings.Replace(string(responseBody), `"code":21`, `"code":1`, -1)
  580. responseBody = []byte(newBody)
  581. }
  582. }
  583. if midjResponse.Code == 1 && midjRequest.Action == "UPLOAD" {
  584. midjourneyTask.Progress = "100%"
  585. midjourneyTask.Status = "SUCCESS"
  586. }
  587. err = midjourneyTask.Insert()
  588. if err != nil {
  589. return &dto.MidjourneyResponse{
  590. Code: 4,
  591. Description: "insert_midjourney_task_failed",
  592. }
  593. }
  594. if midjResponse.Code == 22 { //22-排队中,说明任务已存在
  595. //修改返回值
  596. newBody := strings.Replace(string(responseBody), `"code":22`, `"code":1`, -1)
  597. responseBody = []byte(newBody)
  598. }
  599. //resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
  600. bodyReader := io.NopCloser(bytes.NewBuffer(responseBody))
  601. //for k, v := range resp.Header {
  602. // c.Writer.Header().Set(k, v[0])
  603. //}
  604. c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
  605. _, err = io.Copy(c.Writer, bodyReader)
  606. if err != nil {
  607. return &dto.MidjourneyResponse{
  608. Code: 4,
  609. Description: "copy_response_body_failed",
  610. }
  611. }
  612. err = bodyReader.Close()
  613. if err != nil {
  614. return &dto.MidjourneyResponse{
  615. Code: 4,
  616. Description: "close_response_body_failed",
  617. }
  618. }
  619. return nil
  620. }
  621. type taskChangeParams struct {
  622. ID string
  623. Action string
  624. Index int
  625. }
  626. func getMjRequestPath(path string) string {
  627. requestURL := path
  628. if strings.Contains(requestURL, "/mj-") {
  629. urls := strings.Split(requestURL, "/mj/")
  630. if len(urls) < 2 {
  631. return requestURL
  632. }
  633. requestURL = "/mj/" + urls[1]
  634. }
  635. return requestURL
  636. }