From 8727c4de92bcd04dc47dc82c44f81f8a5cf864f4 Mon Sep 17 00:00:00 2001
From: cherishsince <cherishsince@aliyun.com>
Date: Wed, 5 Jun 2024 10:44:16 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90=E4=BC=98=E5=8C=96=E3=80=91Midjourney?=
 =?UTF-8?q?=20=E6=8F=90=E4=BA=A4=E4=BB=BB=E5=8A=A1=20=E4=BF=9D=E5=AD=98=20?=
 =?UTF-8?q?ImageOptions?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../vo/AiImageMidjourneyImagineReqVO.java     |  4 ++--
 .../ai/dal/dataobject/image/AiImageDO.java    |  2 +-
 .../ai/service/image/AiImageServiceImpl.java  | 19 ++++++++++++++-----
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/image/vo/AiImageMidjourneyImagineReqVO.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/image/vo/AiImageMidjourneyImagineReqVO.java
index e83dd2c87..64c519d57 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/image/vo/AiImageMidjourneyImagineReqVO.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/image/vo/AiImageMidjourneyImagineReqVO.java
@@ -27,10 +27,10 @@ public class AiImageMidjourneyImagineReqVO {
     private String model;
 
     @Schema(description = "图片宽度 --ar 设置")
-    private String width;
+    private Integer width;
 
     @Schema(description = "图片高度 --ar 设置")
-    private String height;
+    private Integer height;
 
     @Schema(description = "版本号 --v 设置")
     private String version;
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/image/AiImageDO.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/image/AiImageDO.java
index 378b09940..1fc498996 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/image/AiImageDO.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/image/AiImageDO.java
@@ -92,7 +92,7 @@ public class AiImageDO extends BaseDO {
      * 2. {@link StabilityAiImageOptions}
      */
     @TableField(typeHandler = JacksonTypeHandler.class)
-    private Map<String, String> options;
+    private Map<String, Object> options;
 
     /**
      * 绘画 response
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
index b99c951e9..2a0dfbd1f 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
@@ -1,5 +1,6 @@
 package cn.iocoder.yudao.module.ai.service.image;
 
+import cn.hutool.core.bean.BeanUtil;
 import cn.hutool.core.codec.Base64;
 import cn.hutool.core.exceptions.ExceptionUtil;
 import cn.hutool.core.map.MapUtil;
@@ -138,9 +139,9 @@ public class AiImageServiceImpl implements AiImageService {
         aiImageDO.setUserId(loginUserId);
         aiImageDO.setPrompt(req.getPrompt());
         aiImageDO.setPlatform(AiPlatformEnum.MIDJOURNEY.getPlatform());
-        aiImageDO.setModel(null);
-        aiImageDO.setWidth(null);
-        aiImageDO.setHeight(null);
+        aiImageDO.setModel(req.getModel());
+        aiImageDO.setWidth(req.getWidth());
+        aiImageDO.setHeight(req.getHeight());
         aiImageDO.setStatus(AiImageStatusEnum.IN_PROGRESS.getStatus());
         // 2、保存 image
         imageMapper.insert(aiImageDO);
@@ -156,15 +157,23 @@ public class AiImageServiceImpl implements AiImageService {
         if (!MidjourneySubmitCodeEnum.SUCCESS_CODES.contains(submitRespVO.getCode())) {
             throw exception(AI_IMAGE_MIDJOURNEY_SUBMIT_FAIL, submitRespVO.getDescription());
         }
+        // 7、构建 imageOptions 参数
+        MidjourneyImageOptions imageOptions = new MidjourneyImageOptions()
+                .setWidth(req.getWidth())
+                .setHeight(req.getHeight())
+                .setModel(req.getModel())
+                .setVersion(req.getVersion())
+                .setState(imagineReqVO.getState());
+        // 8、更新 taskId 和参数
         imageMapper.updateById(new AiImageDO()
                 .setId(aiImageDO.getId())
                 .setTaskId(submitRespVO.getResult())
+                .setOptions(BeanUtil.beanToMap(imageOptions))
         );
         return aiImageDO.getId();
     }
 
 
-
     @Override
     public void deleteImageMy(Long id, Long userId) {
         // 1. 校验是否存在
@@ -269,7 +278,7 @@ public class AiImageServiceImpl implements AiImageService {
      * @param model
      * @return
      */
-    private String buildParams(String width, String height, String version, MidjourneyModelEnum model) {
+    private String buildParams(Integer width, Integer height, String version, MidjourneyModelEnum model) {
         StringBuilder params = new StringBuilder();
         //  --ar 来设置尺寸
         params.append(String.format(" --ar %s:%s ", width, height));