Commit e97408b3ac1d2c05c96855e622b362db6e33aca3

Authored by cherishsince
1 parent 63a8cc24

【优化】dall 绘画,改为异步。

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
@@ -35,6 +35,9 @@ import org.springframework.transaction.annotation.Transactional; @@ -35,6 +35,9 @@ import org.springframework.transaction.annotation.Transactional;
35 35
36 import java.util.Collections; 36 import java.util.Collections;
37 import java.util.List; 37 import java.util.List;
  38 +import java.util.concurrent.LinkedBlockingQueue;
  39 +import java.util.concurrent.ThreadPoolExecutor;
  40 +import java.util.concurrent.TimeUnit;
38 41
39 /** 42 /**
40 * ai 作图 43 * ai 作图
@@ -53,6 +56,8 @@ public class AiImageServiceImpl implements AiImageService { @@ -53,6 +56,8 @@ public class AiImageServiceImpl implements AiImageService {
53 private final OpenAiImageClient openAiImageClient; 56 private final OpenAiImageClient openAiImageClient;
54 private final MidjourneyWebSocketStarter midjourneyWebSocketStarter; 57 private final MidjourneyWebSocketStarter midjourneyWebSocketStarter;
55 private final MidjourneyInteractionsApi midjourneyInteractionsApi; 58 private final MidjourneyInteractionsApi midjourneyInteractionsApi;
  59 + private static ThreadPoolExecutor EXECUTOR = new ThreadPoolExecutor(
  60 + 3, 5, 1, TimeUnit.HOURS, new LinkedBlockingQueue<>(32));
56 61
57 @PostConstruct 62 @PostConstruct
58 public void startMidjourney() { 63 public void startMidjourney() {
@@ -89,34 +94,48 @@ public class AiImageServiceImpl implements AiImageService { @@ -89,34 +94,48 @@ public class AiImageServiceImpl implements AiImageService {
89 94
90 @Override 95 @Override
91 public AiImageDallRespVO dallDrawing(AiImageDallReqVO req) { 96 public AiImageDallRespVO dallDrawing(AiImageDallReqVO req) {
92 - // 获取 model  
93 - OpenAiImageModelEnum openAiImageModelEnum = OpenAiImageModelEnum.valueOfModel(req.getModel());  
94 - OpenAiImageStyleEnum openAiImageStyleEnum = OpenAiImageStyleEnum.valueOfStyle(req.getStyle());  
95 - try {  
96 - // 转换openai 参数  
97 - OpenAiImageOptions openAiImageOptions = new OpenAiImageOptions();  
98 - openAiImageOptions.setModel(openAiImageModelEnum.getModel());  
99 - openAiImageOptions.setStyle(openAiImageStyleEnum.getStyle());  
100 - openAiImageOptions.setSize(req.getSize());  
101 - ImageResponse imageResponse = openAiImageClient.call(new ImagePrompt(req.getPrompt(), openAiImageOptions));  
102 - // 发送  
103 - ImageGeneration imageGeneration = imageResponse.getResult();  
104 - // 图片保存到服务器  
105 - String filePath = fileApi.createFile(HttpUtil.downloadBytes(imageGeneration.getOutput().getUrl()));  
106 - // 保存数据库  
107 - AiImageDO aiImageDO = doSave(req.getPrompt(), req.getSize(), req.getModel(),  
108 - filePath, imageGeneration.getOutput().getUrl(), AiImageStatusEnum.COMPLETE, null,  
109 - null, null, null);  
110 - // 转换 AiImageDallDrawingRespVO  
111 - return AiImageConvert.INSTANCE.convertAiImageDallDrawingRespVO(aiImageDO);  
112 - } catch (AiException aiException) {  
113 - // 保存数据库  
114 - AiImageDO aiImageDO = doSave(req.getPrompt(), req.getSize(), req.getModel(),  
115 - null, null, AiImageStatusEnum.FAIL, aiException.getMessage(),  
116 - null, null, null);  
117 - // 发送错误信息  
118 - return AiImageConvert.INSTANCE.convertAiImageDallDrawingRespVO(aiImageDO);  
119 - } 97 + // 保存数据库
  98 + AiImageDO aiImageDO = doSave(req.getPrompt(), req.getSize(), req.getModel(),
  99 + null, null, AiImageStatusEnum.IN_PROGRESS, null,
  100 + null, null, null);
  101 + // 异步执行
  102 + EXECUTOR.execute(() -> {
  103 + try {
  104 +
  105 + // 获取 model
  106 + OpenAiImageModelEnum openAiImageModelEnum = OpenAiImageModelEnum.valueOfModel(req.getModel());
  107 + OpenAiImageStyleEnum openAiImageStyleEnum = OpenAiImageStyleEnum.valueOfStyle(req.getStyle());
  108 +
  109 + // 转换openai 参数
  110 + OpenAiImageOptions openAiImageOptions = new OpenAiImageOptions();
  111 + openAiImageOptions.setModel(openAiImageModelEnum.getModel());
  112 + openAiImageOptions.setStyle(openAiImageStyleEnum.getStyle());
  113 + openAiImageOptions.setSize(req.getSize());
  114 + ImageResponse imageResponse = openAiImageClient.call(new ImagePrompt(req.getPrompt(), openAiImageOptions));
  115 + // 发送
  116 + ImageGeneration imageGeneration = imageResponse.getResult();
  117 + // 图片保存到服务器
  118 + String filePath = fileApi.createFile(HttpUtil.downloadBytes(imageGeneration.getOutput().getUrl()));
  119 + // 更新数据库
  120 + aiImageMapper.updateById(
  121 + new AiImageDO()
  122 + .setId(aiImageDO.getId())
  123 + .setStatus(AiImageStatusEnum.COMPLETE.getStatus())
  124 + .setPicUrl(filePath)
  125 + .setOriginalPicUrl(imageGeneration.getOutput().getUrl())
  126 + );
  127 + } catch (AiException aiException) {
  128 + // 更新错误信息
  129 + aiImageMapper.updateById(
  130 + new AiImageDO()
  131 + .setId(aiImageDO.getId())
  132 + .setStatus(AiImageStatusEnum.FAIL.getStatus())
  133 + .setErrorMessage(aiException.getMessage())
  134 + );
  135 + }
  136 + });
  137 + // 转换 AiImageDallDrawingRespVO
  138 + return AiImageConvert.INSTANCE.convertAiImageDallDrawingRespVO(aiImageDO);
120 } 139 }
121 140
122 @Override 141 @Override