aidlab/aida_back
7
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

TASK:batch toProductImage;chatStream;

Browse Source
This commit is contained in:
shahaibo
2025-06-05 17:06:11 +08:00
parent 38ca810068
commit 249351bf52
4 changed files with 97 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/main/java/com/ai/da/common/RabbitMQ/GenerateConsumer.java
View File

@@ -387,11 +387,11 @@ public class GenerateConsumer {
} }
} else { } else {
// 修改redis中的数据状态为exception // 修改redis中的数据状态为exception
String key = toProductImageResultKey + ":" + generateResult.get("task_id"); String key = toProductImageResultKey + ":" + generateResult.get("tasks_id");
redisUtil.addToString(key, new Gson().toJson(new GenerateResultVO(generateResult.getString("task_id"), null, null, "Fail")), CommonConstant.GENERATE_RESULT_EXPIRE_TIME); redisUtil.addToString(key, new Gson().toJson(new GenerateResultVO(generateResult.getString("tasks_id"), null, null, "Fail")), CommonConstant.GENERATE_RESULT_EXPIRE_TIME);
// 将异常信息存到exception中 // 将异常信息存到exception中
HashMap<String, String> exceptionInfo = new HashMap<>(); HashMap<String, String> exceptionInfo = new HashMap<>();
exceptionInfo.put(generateResult.getString("task_id"), generateResult.getString("data")); exceptionInfo.put(generateResult.getString("tasks_id"), generateResult.getString("data"));
// 存redis // 存redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo); redisUtil.addToMap(exceptionMapKey, exceptionInfo);
} }
@@ -400,7 +400,7 @@ public class GenerateConsumer {
try { try {
channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false); channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
// 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除 // 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除
redisUtil.removeFromZSet(consumptionOrderKey, generateResult.getString("task_id")); redisUtil.removeFromZSet(consumptionOrderKey, generateResult.getString("tasks_id"));
} catch (IOException exception) { } catch (IOException exception) {
log.error("手动确认，取消返回队列，不再重新消费"); log.error("手动确认，取消返回队列，不再重新消费");
} }
@@ -408,13 +408,13 @@ public class GenerateConsumer {
String exceptionMessage = JSONObject.toJSONString(generateResult) + String exceptionMessage = JSONObject.toJSONString(generateResult) +
" Exception message ： " + e.getMessage(); " Exception message ： " + e.getMessage();
HashMap<String, String> exceptionInfo = new HashMap<>(); HashMap<String, String> exceptionInfo = new HashMap<>();
exceptionInfo.put(String.valueOf(generateResult.get("task_id")), exceptionMessage); exceptionInfo.put(String.valueOf(generateResult.get("tasks_id")), exceptionMessage);
// 存redis // 存redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo); redisUtil.addToMap(exceptionMapKey, exceptionInfo);
} }
long end = System.currentTimeMillis(); long end = System.currentTimeMillis();
log.info("tasks_id : {}, end , message : {}, 执行时长： {} 毫秒", generateResult.get("task_id"), generateResult.get("message"), (end - start)); log.info("tasks_id : {}, end , message : {}, 执行时长： {} 毫秒", generateResult.get("tasks_id"), generateResult.get("message"), (end - start));
log.info("============ProcessToProductImageBatchResult End listening=========="); log.info("============ProcessToProductImageBatchResult End listening==========");
} }
@@ -426,23 +426,23 @@ public class GenerateConsumer {
log.info("relightBatch response : {}", generateResult); log.info("relightBatch response : {}", generateResult);
try { try {
log.info("task_id : {} start ", generateResult.get("task_id")); log.info("task_id : {} start ", generateResult.get("tasks_id"));
if (!StringUtils.isEmpty(generateResult.getString("progress"))) { if (!StringUtils.isEmpty(generateResult.getString("progress"))) {
String progress = generateResult.getString("progress"); String progress = generateResult.getString("progress");
JSONArray result = generateResult.getJSONArray("result"); JSONObject result = generateResult.getJSONObject("result_data");
String url = null; String url = null;
if (!StringUtils.isEmpty(result)) { if (!StringUtils.isEmpty(result)) {
url = result.getString(0); url = result.getString("image_url");
String taskId = generateResult.getString("task_id"); String taskId = generateResult.getString("tasks_id");
userLikeGroupService.relightBatch(taskId, url, progress); userLikeGroupService.relightBatch(taskId, url, progress);
} }
} else { } else {
// 修改redis中的数据状态为exception // 修改redis中的数据状态为exception
String key = relightResultKey + ":" + generateResult.get("task_id"); String key = relightResultKey + ":" + generateResult.get("tasks_id");
redisUtil.addToString(key, new Gson().toJson(new GenerateResultVO(generateResult.getString("task_id"), null, null, "Fail")), CommonConstant.GENERATE_RESULT_EXPIRE_TIME); redisUtil.addToString(key, new Gson().toJson(new GenerateResultVO(generateResult.getString("tasks_id"), null, null, "Fail")), CommonConstant.GENERATE_RESULT_EXPIRE_TIME);
// 将异常信息存到exception中 // 将异常信息存到exception中
HashMap<String, String> exceptionInfo = new HashMap<>(); HashMap<String, String> exceptionInfo = new HashMap<>();
exceptionInfo.put(generateResult.getString("task_id"), generateResult.getString("data")); exceptionInfo.put(generateResult.getString("tasks_id"), generateResult.getString("data"));
// 存redis // 存redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo); redisUtil.addToMap(exceptionMapKey, exceptionInfo);
} }
@@ -451,7 +451,7 @@ public class GenerateConsumer {
try { try {
channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false); channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
// 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除 // 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除
redisUtil.removeFromZSet(consumptionOrderKey, generateResult.getString("task_id")); redisUtil.removeFromZSet(consumptionOrderKey, generateResult.getString("tasks_id"));
} catch (IOException exception) { } catch (IOException exception) {
log.error("手动确认，取消返回队列，不再重新消费"); log.error("手动确认，取消返回队列，不再重新消费");
} }
@@ -459,13 +459,13 @@ public class GenerateConsumer {
String exceptionMessage = JSONObject.toJSONString(generateResult) + String exceptionMessage = JSONObject.toJSONString(generateResult) +
" Exception message ： " + e.getMessage(); " Exception message ： " + e.getMessage();
HashMap<String, String> exceptionInfo = new HashMap<>(); HashMap<String, String> exceptionInfo = new HashMap<>();
exceptionInfo.put(String.valueOf(generateResult.get("task_id")), exceptionMessage); exceptionInfo.put(String.valueOf(generateResult.get("tasks_id")), exceptionMessage);
// 存redis // 存redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo); redisUtil.addToMap(exceptionMapKey, exceptionInfo);
} }
long end = System.currentTimeMillis(); long end = System.currentTimeMillis();
log.info("task_id : {}, end , message : {}, 执行时长： {} 毫秒", generateResult.get("task_id"), generateResult.get("message"), (end - start)); log.info("task_id : {}, end , message : {}, 执行时长： {} 毫秒", generateResult.get("tasks_id"), generateResult.get("message"), (end - start));
log.info("============ProcessRelightBatchResult End listening=========="); log.info("============ProcessRelightBatchResult End listening==========");
} }

2
src/main/java/com/ai/da/model/dto/BatchParamDTO.java
View File

@@ -11,4 +11,6 @@ public class BatchParamDTO {
private BigDecimal image_strength; private BigDecimal image_strength;
private String image_url; private String image_url;
private String product_type; private String product_type;
private String direction;
} }

11
src/main/java/com/ai/da/python/PythonService.java
View File

@@ -4314,7 +4314,7 @@ public class PythonService {
throw new BusinessException("toProductImage.interface.exception"); throw new BusinessException("toProductImage.interface.exception");
} }
public Boolean relightBatch(String url, String taskId, String prompt, String direction, String relightType) { public Boolean relightBatch(String taskIdBatch, List<BatchParamDTO> paramList, String userId) {
// todo 限流校验 // todo 限流校验
// AccessLimitUtils.validate("design",5); // AccessLimitUtils.validate("design",5);
OkHttpClient client = new OkHttpClient().newBuilder() OkHttpClient client = new OkHttpClient().newBuilder()
@@ -4326,12 +4326,9 @@ public class PythonService {
MediaType mediaType = MediaType.parse("application/json"); MediaType mediaType = MediaType.parse("application/json");
//关闭FastJson的引用检测 防止出现$ref 现象 //关闭FastJson的引用检测 防止出现$ref 现象
Map<String, Object> map = new HashMap<>(); Map<String, Object> map = new HashMap<>();
map.put("tasks_id", taskId); map.put("batch_tasks_id", taskIdBatch);
map.put("image_url", url); map.put("user_id", userId);
map.put("prompt", prompt); map.put("batch_data_list", paramList);
map.put("direction", direction);
map.put("product_type", relightType);
map.put("batch_size", 1);
log.info("relightImage请求python 参数:####{}", map); log.info("relightImage请求python 参数:####{}", map);
String param = JSON.toJSONString(map, SerializerFeature.WriteNullStringAsEmpty); String param = JSON.toJSONString(map, SerializerFeature.WriteNullStringAsEmpty);
log.info(param); log.info(param);

47
src/main/java/com/ai/da/service/impl/DesignServiceImpl.java
View File

@@ -2039,9 +2039,9 @@ public class DesignServiceImpl extends ServiceImpl<DesignMapper, Design> impleme
}else { }else {
s = "Snow moutain, snowy day, natural light"; s = "Snow moutain, snowy day, natural light";
} }
Map<ToProductImageVO, Integer> toProductImageVOIntegerMap = allocateElements(toProductImageDTO.getToProductImageVOList(), cloudTaskDTO.getNums());
for (ToProductImageVO toProductImageVO : toProductImageDTO.getToProductImageVOList()) { for (ToProductImageVO toProductImageVO : toProductImageDTO.getToProductImageVOList()) {
String taskId = UUID.randomUUID() + "-" + i + "-" + userHolder.getId(); String taskId;
i ++;
if (toProductImageVO.getElementType().equals("ToProductImage")) { if (toProductImageVO.getElementType().equals("ToProductImage")) {
ToProductImageResult toProductImageResult1 = toProductImageResultMapper.selectById(toProductImageVO.getElementId()); ToProductImageResult toProductImageResult1 = toProductImageResultMapper.selectById(toProductImageVO.getElementId());
String relightType = "overall"; String relightType = "overall";
@@ -2052,10 +2052,18 @@ public class DesignServiceImpl extends ServiceImpl<DesignMapper, Design> impleme
if (design.getSingleOverall().equals("single")) { if (design.getSingleOverall().equals("single")) {
relightType = "single"; relightType = "single";
} }
} List<BatchParamDTO> paramList = new ArrayList<>();
List<String> promptList = pythonService.getPrompt(s, toProductImageVOIntegerMap.get(toProductImageVO));
for (int i1 = 0; i1 < toProductImageVOIntegerMap.get(toProductImageVO); i1++) {
BatchParamDTO batchParamDTO = new BatchParamDTO();
taskId = UUID.randomUUID() + "-" + i + "-" + userHolder.getId();
batchParamDTO.setTasks_id(taskId);
batchParamDTO.setPrompt(promptList.get(i1));
batchParamDTO.setImage_url(toProductImageResult1.getUrl());
batchParamDTO.setProduct_type(relightType);
batchParamDTO.setDirection(toProductImageDTO.getDirection());
paramList.add(batchParamDTO);
// 走模型
pythonService.relightBatch(toProductImageResult1.getUrl(), taskId, s, toProductImageDTO.getDirection(), relightType);
ToProductImageResult toProductImageResult = new ToProductImageResult(); ToProductImageResult toProductImageResult = new ToProductImageResult();
toProductImageResult.setElementId(toProductImageResult1.getId()); toProductImageResult.setElementId(toProductImageResult1.getId());
toProductImageResult.setElementType("ToProductImage"); toProductImageResult.setElementType("ToProductImage");
@@ -2074,10 +2082,30 @@ public class DesignServiceImpl extends ServiceImpl<DesignMapper, Design> impleme
toProductImageResult.setDirection(toProductImageDTO.getDirection()); toProductImageResult.setDirection(toProductImageDTO.getDirection());
toProductImageResultMapper.insert(toProductImageResult); toProductImageResultMapper.insert(toProductImageResult);
result.add(toProductImageResult); result.add(toProductImageResult);
// 添加需要扣除的积分到预扣除区
creditsService.addRecordToCreditsDeduction(userHolder.getId(), taskId, CreditsEventsEnum.RELIGHT);
i ++;
}
// 走模型
pythonService.relightBatch(batchTaskId, paramList, userHolder.getId().toString());
}
}else { }else {
ToProductElement toProductElement = toProductElementMapper.selectById(toProductImageVO.getElementId()); ToProductElement toProductElement = toProductElementMapper.selectById(toProductImageVO.getElementId());
// 走模型 // 走模型
pythonService.relightBatch(toProductElement.getUrl(), taskId, s, toProductImageDTO.getDirection(), "overall"); List<BatchParamDTO> paramList = new ArrayList<>();
List<String> promptList = pythonService.getPrompt(s, toProductImageVOIntegerMap.get(toProductImageVO));
for (int i1 = 0; i1 < toProductImageVOIntegerMap.get(toProductImageVO); i1++) {
BatchParamDTO batchParamDTO = new BatchParamDTO();
taskId = UUID.randomUUID() + "-" + i + "-" + userHolder.getId();
batchParamDTO.setTasks_id(taskId);
batchParamDTO.setPrompt(promptList.get(i1));
batchParamDTO.setImage_url(toProductElement.getUrl());
batchParamDTO.setProduct_type("overall");
batchParamDTO.setDirection(toProductImageDTO.getDirection());
paramList.add(batchParamDTO);
ToProductImageResult toProductImageResult = new ToProductImageResult(); ToProductImageResult toProductImageResult = new ToProductImageResult();
toProductImageResult.setElementId(toProductElement.getId()); toProductImageResult.setElementId(toProductElement.getId());
toProductImageResult.setElementType("ProductElement"); toProductImageResult.setElementType("ProductElement");
@@ -2096,9 +2124,14 @@ public class DesignServiceImpl extends ServiceImpl<DesignMapper, Design> impleme
toProductImageResult.setDirection(toProductImageDTO.getDirection()); toProductImageResult.setDirection(toProductImageDTO.getDirection());
toProductImageResultMapper.insert(toProductImageResult); toProductImageResultMapper.insert(toProductImageResult);
result.add(toProductImageResult); result.add(toProductImageResult);
}
// 添加需要扣除的积分到预扣除区 // 添加需要扣除的积分到预扣除区
creditsService.addRecordToCreditsDeduction(userHolder.getId(), taskId, CreditsEventsEnum.RELIGHT); creditsService.addRecordToCreditsDeduction(userHolder.getId(), taskId, CreditsEventsEnum.RELIGHT);
i ++;
}
// 走模型
pythonService.relightBatch(batchTaskId, paramList, userHolder.getId().toString());
}
} }
CloudTask cloudTask = CopyUtil.copyObject(cloudTaskDTO, CloudTask.class); CloudTask cloudTask = CopyUtil.copyObject(cloudTaskDTO, CloudTask.class);
cloudTask.setProjectId(projectId); cloudTask.setProjectId(projectId);