relay-mj.go 22 KB

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