relay-mj.go 23 KB

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