huangge1199/vue-pro
1
0
Fork 0
mirror of https://gitee.com/huangge1199_admin/vue-pro.git synced 2025-01-18 19:20:05 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

【增加】MidjourneyJob 被动拉去

Browse Source
This commit is contained in:
cherishsince 2024-06-05 15:20:53 +08:00
parent b0b7a3df09
commit f15ecda727
1 changed files with 74 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/job/MidjourneyJob.java Normal file
View File

@ -0,0 +1,74 @@
package cn.iocoder.yudao.module.ai.job;
import cn.hutool.core.collection.CollUtil;
import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
import cn.iocoder.yudao.framework.quartz.core.handler.JobHandler;
import cn.iocoder.yudao.module.ai.client.MidjourneyProxyClient;
import cn.iocoder.yudao.module.ai.client.vo.MidjourneyNotifyReqVO;
import cn.iocoder.yudao.module.ai.dal.dataobject.image.AiImageDO;
import cn.iocoder.yudao.module.ai.dal.mysql.image.AiImageMapper;
import cn.iocoder.yudao.module.ai.enums.image.AiImageStatusEnum;
import cn.iocoder.yudao.module.ai.service.image.AiImageService;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
/**
* midjourney job 定时拉去 midjourney 绘制状态
*
* @author fansili
* @time 2024/6/5 14:55
* @since 1.0
*/
@Component
@Slf4j
public class MidjourneyJob implements JobHandler {
@Autowired
private MidjourneyProxyClient midjourneyProxyClient;
@Autowired
private AiImageMapper imageMapper;
@Autowired
private AiImageService imageService;
@Override
public String execute(String param) throws Exception {
// 1、获取 midjourney 平台，状态在 “进行中” 的 image
log.info("Midjourney 同步 - 开始...");
List<AiImageDO> imageList = imageMapper.selectList(
new LambdaUpdateWrapper<AiImageDO>()
.eq(AiImageDO::getStatus, AiImageStatusEnum.IN_PROGRESS.getStatus())
.eq(AiImageDO::getPlatform, AiPlatformEnum.MIDJOURNEY.getPlatform())
);
log.info("Midjourney 同步 - 任务数量 {}!", imageList.size());
if (CollUtil.isEmpty(imageList)) {
return "Midjourney 同步 - 数量为空!";
}
// 2、批量拉去 task 信息
List<MidjourneyNotifyReqVO> taskList = midjourneyProxyClient
.listByCondition(imageList.stream().map(AiImageDO::getTaskId).collect(Collectors.toSet()));
Map<String, MidjourneyNotifyReqVO> taskIdMap = taskList.stream().collect(Collectors.toMap(MidjourneyNotifyReqVO::getId, o -> o));
// 3、更新 image 状态
List<AiImageDO> updateImageList = new ArrayList<>();
for (AiImageDO aiImageDO : imageList) {
// 3.1 排除掉空的情况
if (!taskIdMap.containsKey(aiImageDO.getTaskId())) {
log.warn("Midjourney 同步 - {} 任务为空!", aiImageDO.getTaskId());
continue;
}
// 3.2 获取通知对象
MidjourneyNotifyReqVO notifyReqVO = taskIdMap.get(aiImageDO.getTaskId());
// 3.2 构建更新对象
updateImageList.add(imageService.buildUpdateImage(aiImageDO.getId(), notifyReqVO));
}
// 4、批了更新 updateImageList
imageMapper.updateBatch(updateImageList);
return "Midjourney 同步 - ".concat(String.valueOf(updateImageList.size())).concat(" 任务!");
}
}