'commit'

dsAiSupport/src/main/java/com/dsideal/aiSupport/Util/Liblib/UltraImg2Img.java

@@ -9,12 +9,16 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 
 public class UltraImg2Img extends LibLibCommon {
     //日志
     private static final Logger log = LoggerFactory.getLogger(UltraImg2Img.class);
     // 图生图API路径
     private static final String IMG_TO_IMG_PATH = "/api/generate/webui/img2img/ultra";
+    // 查询任务状态API路径
+    private static final String QUERY_STATUS_PATH = "/api/generate/webui/status";
 
 
     /**
@@ -113,6 +117,103 @@ public class UltraImg2Img extends LibLibCommon {
         return generateUuid;
     }
 
+    /**
+     * 查询生图任务结果
+     * 
+     * @param generateUuid 生图任务UUID
+     * @return 任务结果信息
+     * @throws IOException 异常信息
+     */
+    public static JSONObject queryTaskResult(String generateUuid) throws IOException {
+        // 创建OkHttpClient
+        OkHttpClient client = createHttpClient();
+        
+        // 构建请求体
+        JSONObject requestBody = new JSONObject();
+        requestBody.put("generateUuid", generateUuid);
+        
+        // 获取API路径
+        String uri = QUERY_STATUS_PATH;
+        
+        // 生成签名信息
+        LibLibCommon.SignatureInfo signInfo = LibLibCommon.sign(uri);
+        
+        // 构建带签名的URL
+        HttpUrl.Builder urlBuilder = HttpUrl.parse(API_BASE_URL + uri).newBuilder()
+                .addQueryParameter("AccessKey", accessKey)
+                .addQueryParameter("Signature", signInfo.getSignature())
+                .addQueryParameter("Timestamp", String.valueOf(signInfo.getTimestamp()))
+                .addQueryParameter("SignatureNonce", signInfo.getSignatureNonce());
+        
+        // 创建请求
+        MediaType mediaType = MediaType.parse("application/json");
+        RequestBody body = RequestBody.create(mediaType, requestBody.toJSONString());
+        Request request = new Request.Builder()
+                .url(urlBuilder.build())
+                .method("POST", body)
+                .addHeader("Content-Type", "application/json")
+                .build();
+        
+        // 执行请求
+        log.info("查询生图任务结果: {}", requestBody.toJSONString());
+        Response response = client.newCall(request).execute();
+        
+        // 处理响应
+        if (!response.isSuccessful()) {
+            String errorMsg = "查询生图任务结果失败，状态码: " + response.code();
+            log.error(errorMsg);
+            throw new IOException(errorMsg);
+        }
+        
+        // 解析响应
+        String responseBody = response.body().string();
+        log.info("查询生图任务结果响应: {}", responseBody);
+        
+        JSONObject responseJson = JSON.parseObject(responseBody);
+        int code = responseJson.getIntValue("code");
+        
+        if (code != 0) {
+            String errorMsg = "查询生图任务结果失败，错误码: " + code + ", 错误信息: " + responseJson.getString("msg");
+            log.error(errorMsg);
+            throw new IOException(errorMsg);
+        }
+        
+        return responseJson.getJSONObject("data");
+    }
+    
+    /**
+     * 获取生成图片的URL列表
+     * 
+     * @param generateUuid 生图任务UUID
+     * @return 图片URL列表
+     * @throws IOException 异常信息
+     */
+    public static List<String> getGeneratedImageUrls(String generateUuid) throws IOException {
+        JSONObject resultData = queryTaskResult(generateUuid);
+        List<String> imageUrls = new ArrayList<>();
+        
+        // 检查生成状态
+        int generateStatus = resultData.getIntValue("generateStatus");
+        if (generateStatus == 5) { // 5表示生成成功
+            if (resultData.containsKey("images")) {
+                for (Object imageObj : resultData.getJSONArray("images")) {
+                    JSONObject imageJson = (JSONObject) imageObj;
+                    String imageUrl = imageJson.getString("imageUrl");
+                    if (imageUrl != null && !imageUrl.isEmpty()) {
+                        // 清理URL，移除可能的反引号和多余空格
+                        imageUrl = imageUrl.trim().replace("`", "");
+                        imageUrls.add(imageUrl);
+                    }
+                }
+            }
+        } else {
+            log.info("生图任务尚未完成，当前状态: {}, 完成百分比: {}%", 
+                    generateStatus, resultData.getIntValue("percentCompleted"));
+        }
+        
+        return imageUrls;
+    }
+
     /**
      * 使用示例
      */
@@ -133,15 +234,68 @@ public class UltraImg2Img extends LibLibCommon {
             double cfgScale = 3.5;
             // 随机种子，-1表示随机
             long seed = -1;
-
+    
             // 提交图生图任务
             String generateUuid = submitImageToImageTask(
                     templateUuid, prompt, sourceImageUrl,
                     width, height, steps, cfgScale, seed, null, 0
             );
+            
             // 输出生成任务UUID
             log.info("图生图任务已提交，任务UUID: {}", generateUuid);
-            log.info("请使用此UUID查询任务进度和结果");
+            
+            // 每5秒查询一次任务结果，直到任务完成或失败
+            boolean isCompleted = false;
+            int maxRetries = 60; // 最多尝试60次，即5分钟
+            int retryCount = 0;
+            
+            while (!isCompleted && retryCount < maxRetries) {
+                try {
+                    // 等待5秒
+                    Thread.sleep(5000);
+                    log.info("第{}次查询任务结果...", retryCount + 1);
+                    
+                    // 查询任务结果
+                    JSONObject resultData = queryTaskResult(generateUuid);
+                    int generateStatus = resultData.getIntValue("generateStatus");
+                    int percentCompleted = resultData.getIntValue("percentCompleted");
+                    
+                    log.info("任务状态: {}, 完成百分比: {}%", generateStatus, percentCompleted);
+                    
+                    // 检查任务是否完成或失败
+                    if (generateStatus == 5) { // 5表示生成成功
+                        isCompleted = true;
+                        log.info("任务已完成！");
+                        
+                        // 获取生成的图片URL
+                        List<String> imageUrls = getGeneratedImageUrls(generateUuid);
+                        if (!imageUrls.isEmpty()) {
+                            log.info("生成的图片URL:");
+                            for (String imageUrl : imageUrls) {
+                                log.info(imageUrl);
+                            }
+                        } else {
+                            log.info("未找到生成的图片");
+                        }
+                    } else if (generateStatus == 4) { // 4表示生成失败
+                        isCompleted = true;
+                        log.error("任务失败: {}", resultData.getString("generateMsg"));
+                    }
+                } catch (IOException e) {
+                    log.error("查询任务结果失败", e);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    log.error("线程被中断", e);
+                    break;
+                }
+                
+                retryCount++;
+            }
+            
+            if (!isCompleted) {
+                log.warn("达到最大重试次数，任务可能仍在处理中");
+            }
+            
         } catch (Exception e) {
             log.error("图生图任务执行失败", e);
         }