【增加】aiImage 增加 dall-e-2、dall-e-3

cherishsince
1 parent 844ae285
Showing 5 changed files with 85 additions and 4 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/AiImageController.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/mapper/AiImageMapper.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/impl/AiImageServiceImpl.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/vo/AiImageDallDrawingReq.java
@@ -34,6 +34,6 @@ public class AiImageController {
     public SseEmitter dallDrawing(@Validated @ModelAttribute AiImageDallDrawingReq req) {
         Utf8SseEmitter sseEmitter = new Utf8SseEmitter();
         aiImageService.dallDrawing(req, sseEmitter);
-        return null;
+        return sseEmitter;
     }
 }
@@ -17,15 +17,24 @@ import lombok.experimental.Accessors;
  */
 @Data
 @Accessors(chain = true)
-@TableName("ai_drawing_image")
-public class AiDrawingImageDO extends BaseDO {
+@TableName("ai_image")
+public class AiImageDO extends BaseDO {
     @TableId(type = IdType.AUTO)
     private Long id;
+    @Schema(description = "用户id")
+    private Long userId;
+
     @Schema(description = "提示词")
     private String prompt;
+    @Schema(description = "模型")
+    private String modal;
+
+    @Schema(description = "生成图像的尺寸大小。对于dall-e-2模型，尺寸可为256x256, 512x512, 或 1024x1024。对于dall-e-3模型，尺寸可为1024x1024, 1792x1024, 或 1024x1792。")
+    private String size;
+
     @Schema(description = "绘画状态：提交、排队、绘画中、绘画完成、绘画失败")
     private String drawingStatus;
+package cn.iocoder.yudao.module.ai.mapper;
+
+import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
+import cn.iocoder.yudao.module.ai.dal.dataobject.AiImageDO;
+import org.apache.ibatis.annotations.Mapper;
+import org.springframework.stereotype.Repository;
+
+/**
+ * ai image
+ *
+ * @author fansili
+ * @time 2024/4/28 14:01
+ * @since 1.0
+ */
+@Repository
+@Mapper
+public interface AiImageMapper extends BaseMapperX<AiImageDO> {
+
+
+}
 package cn.iocoder.yudao.module.ai.service.impl;
+import cn.iocoder.yudao.framework.ai.image.Image;
+import cn.iocoder.yudao.framework.ai.image.ImageGeneration;
+import cn.iocoder.yudao.framework.ai.image.ImagePrompt;
+import cn.iocoder.yudao.framework.ai.image.ImageResponse;
+import cn.iocoder.yudao.framework.ai.imageopenai.OpenAiImageClient;
+import cn.iocoder.yudao.framework.ai.imageopenai.OpenAiImageModelEnum;
+import cn.iocoder.yudao.framework.ai.imageopenai.OpenAiImageOptions;
+import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
 import cn.iocoder.yudao.module.ai.controller.Utf8SseEmitter;
+import cn.iocoder.yudao.module.ai.dal.dataobject.AiImageDO;
+import cn.iocoder.yudao.module.ai.mapper.AiImageMapper;
 import cn.iocoder.yudao.module.ai.service.AiImageService;
 import cn.iocoder.yudao.module.ai.vo.AiImageDallDrawingReq;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.http.MediaType;
 import org.springframework.stereotype.Service;
+import java.io.IOException;
+
 /**
  * ai 作图
  *
@@ -19,9 +32,39 @@ import org.springframework.stereotype.Service;
 @Slf4j
 public class AiImageServiceImpl implements AiImageService {
+    private final AiImageMapper aiImageMapper;
+    private final OpenAiImageClient openAiImageClient;
     @Override
     public void dallDrawing(AiImageDallDrawingReq req, Utf8SseEmitter sseEmitter) {
-
+        // 获取 model
+        OpenAiImageModelEnum openAiImageModelEnum = OpenAiImageModelEnum.valueOfModel(req.getModal());
+        //
+        OpenAiImageOptions openAiImageOptions = new OpenAiImageOptions();
+        openAiImageOptions.setModel(openAiImageModelEnum);
+        ImageResponse imageResponse = openAiImageClient.call(new ImagePrompt(req.getPrompt(), openAiImageOptions));
+        // 发送
+        ImageGeneration imageGeneration = imageResponse.getResult();
+        try {
+            sseEmitter.send(imageGeneration, MediaType.APPLICATION_JSON);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        } finally {
+            // 发送 complete
+            sseEmitter.complete();
+        }
+        Long loginUserId = SecurityFrameworkUtils.getLoginUserId();
+        //
+        // 保存数据库
+        Image output = imageGeneration.getOutput();
+        String imageUrl = output.getUrl();
+        AiImageDO aiImageDO = new AiImageDO();
+        aiImageDO.setId(null);
+        aiImageDO.setPrompt(req.getPrompt());
+        aiImageDO.setSize(req.getSize());
+        aiImageDO.setModal(req.getModal());
+        aiImageDO.setUserId(loginUserId);
+        aiImageDO.setDrawingImageUrl(imageUrl);
+        aiImageMapper.insert(aiImageDO);
     }
 }
@@ -19,4 +19,13 @@ public class AiImageDallDrawingReq {
     @Schema(description = "提示词")
     @NotNull(message = "提示词不能为空!")
     private String prompt;
+
+    @Schema(description = "模型")
+    @NotNull(message = "模型不能为空")
+    private String modal;
+
+    @Schema(description = "生成图像的尺寸大小。对于dall-e-2模型，尺寸可为256x256, 512x512, 或 1024x1024。对于dall-e-3模型，尺寸可为1024x1024, 1792x1024, 或 1024x1792。")
+    @NotNull(message = "size不能为空!")
+    private String size;
+
 }