Commit f15ecda727df5ca29febd81f36a7637ddaa3edf5
1 parent
b0b7a3df
【增加】MidjourneyJob 被动拉去
Showing
1 changed file
with
74 additions
and
0 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/job/MidjourneyJob.java
0 → 100644
| 1 | +package cn.iocoder.yudao.module.ai.job; | |
| 2 | + | |
| 3 | +import cn.hutool.core.collection.CollUtil; | |
| 4 | +import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum; | |
| 5 | +import cn.iocoder.yudao.framework.quartz.core.handler.JobHandler; | |
| 6 | +import cn.iocoder.yudao.module.ai.client.MidjourneyProxyClient; | |
| 7 | +import cn.iocoder.yudao.module.ai.client.vo.MidjourneyNotifyReqVO; | |
| 8 | +import cn.iocoder.yudao.module.ai.dal.dataobject.image.AiImageDO; | |
| 9 | +import cn.iocoder.yudao.module.ai.dal.mysql.image.AiImageMapper; | |
| 10 | +import cn.iocoder.yudao.module.ai.enums.image.AiImageStatusEnum; | |
| 11 | +import cn.iocoder.yudao.module.ai.service.image.AiImageService; | |
| 12 | +import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper; | |
| 13 | +import lombok.extern.slf4j.Slf4j; | |
| 14 | +import org.springframework.beans.factory.annotation.Autowired; | |
| 15 | +import org.springframework.stereotype.Component; | |
| 16 | + | |
| 17 | +import java.util.ArrayList; | |
| 18 | +import java.util.List; | |
| 19 | +import java.util.Map; | |
| 20 | +import java.util.stream.Collectors; | |
| 21 | + | |
| 22 | +/** | |
| 23 | + * midjourney job 定时拉去 midjourney 绘制状态 | |
| 24 | + * | |
| 25 | + * @author fansili | |
| 26 | + * @time 2024/6/5 14:55 | |
| 27 | + * @since 1.0 | |
| 28 | + */ | |
| 29 | +@Component | |
| 30 | +@Slf4j | |
| 31 | +public class MidjourneyJob implements JobHandler { | |
| 32 | + | |
| 33 | + @Autowired | |
| 34 | + private MidjourneyProxyClient midjourneyProxyClient; | |
| 35 | + @Autowired | |
| 36 | + private AiImageMapper imageMapper; | |
| 37 | + @Autowired | |
| 38 | + private AiImageService imageService; | |
| 39 | + | |
| 40 | + @Override | |
| 41 | + public String execute(String param) throws Exception { | |
| 42 | + // 1、获取 midjourney 平台，状态在 “进行中” 的 image | |
| 43 | + log.info("Midjourney 同步 - 开始..."); | |
| 44 | + List<AiImageDO> imageList = imageMapper.selectList( | |
| 45 | + new LambdaUpdateWrapper<AiImageDO>() | |
| 46 | + .eq(AiImageDO::getStatus, AiImageStatusEnum.IN_PROGRESS.getStatus()) | |
| 47 | + .eq(AiImageDO::getPlatform, AiPlatformEnum.MIDJOURNEY.getPlatform()) | |
| 48 | + ); | |
| 49 | + log.info("Midjourney 同步 - 任务数量 {}!", imageList.size()); | |
| 50 | + if (CollUtil.isEmpty(imageList)) { | |
| 51 | + return "Midjourney 同步 - 数量为空!"; | |
| 52 | + } | |
| 53 | + // 2、批量拉去 task 信息 | |
| 54 | + List<MidjourneyNotifyReqVO> taskList = midjourneyProxyClient | |
| 55 | + .listByCondition(imageList.stream().map(AiImageDO::getTaskId).collect(Collectors.toSet())); | |
| 56 | + Map<String, MidjourneyNotifyReqVO> taskIdMap = taskList.stream().collect(Collectors.toMap(MidjourneyNotifyReqVO::getId, o -> o)); | |
| 57 | + // 3、更新 image 状态 | |
| 58 | + List<AiImageDO> updateImageList = new ArrayList<>(); | |
| 59 | + for (AiImageDO aiImageDO : imageList) { | |
| 60 | + // 3.1 排除掉空的情况 | |
| 61 | + if (!taskIdMap.containsKey(aiImageDO.getTaskId())) { | |
| 62 | + log.warn("Midjourney 同步 - {} 任务为空!", aiImageDO.getTaskId()); | |
| 63 | + continue; | |
| 64 | + } | |
| 65 | + // 3.2 获取通知对象 | |
| 66 | + MidjourneyNotifyReqVO notifyReqVO = taskIdMap.get(aiImageDO.getTaskId()); | |
| 67 | + // 3.2 构建更新对象 | |
| 68 | + updateImageList.add(imageService.buildUpdateImage(aiImageDO.getId(), notifyReqVO)); | |
| 69 | + } | |
| 70 | + // 4、批了更新 updateImageList | |
| 71 | + imageMapper.updateBatch(updateImageList); | |
| 72 | + return "Midjourney 同步 - ".concat(String.valueOf(updateImageList.size())).concat(" 任务!"); | |
| 73 | + } | |
| 74 | +} | ... | ... |