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 | +} |