'commit'

dsLightRag/Liblib/LiblibUtil.py

@@ -45,4 +45,67 @@ class LiblibUtil:
             return None
         except Exception as e:
             print(f"处理异常: {str(e)}")
-            return None
+            return None
+
+    # ===== 从T1_VersionGet.py迁移的方法 =====
+    def get_model_version_info(self, version_uuid):
+        """获取模型版本信息"""
+        uri = "/api/model/version/get"
+        payload = {"versionUuid": version_uuid}
+        
+        model_info = self.post_request(uri, payload)
+        if model_info:
+            return {
+                'modelName': model_info.get('modelName'),
+                'versionName': model_info.get('versionName'),
+                'commercialUse': model_info.get('commercialUse'),
+                'modelUrl': model_info.get('modelUrl')
+            }
+        return None
+
+    def generate_text_to_image(self, template_uuid, generate_params):
+        """调用text2img接口生成图片"""
+        uri = "/api/generate/webui/text2img"
+        payload = {
+            "templateUuid": template_uuid,
+            "generateParams": generate_params
+        }
+        return self.post_request(uri, payload)
+
+    # ===== 从T2.py迁移的方法 =====
+    def get_generation_status(self, generate_uuid):
+        """查询生图任务状态和结果"""
+        uri = "/api/generate/webui/status"
+        payload = {"generateUuid": generate_uuid}
+        return self.post_request(uri, payload)
+
+    def wait_for_generation_completion(self, generate_uuid, check_interval=5, max_wait_time=300):
+        """等待生图任务完成并获取最终结果"""
+        start_time = time.time()
+        
+        while time.time() - start_time < max_wait_time:
+            status_data = self.get_generation_status(generate_uuid)
+            
+            if status_data:
+                status = status_data.get('status')
+                progress = status_data.get('percentCompleted', 0) * 100
+                print(f"生成进度：{progress:.2f}%")
+                
+                if status == 'completed':
+                    print("生图任务已完成")
+                    images = status_data.get('images', [])
+                    if images:
+                        print("生成的最终图片：")
+                        for img in images:
+                            print(f"- {img.get('imageUrl')}")
+                    return status_data
+                elif status == 'failed':
+                    print(f"生图任务失败: {status_data.get('errorMessage', '未知错误')}")
+                    return None
+                else:
+                    print(f"生图任务处理中，当前状态: {status}")
+            
+            time.sleep(check_interval)
+        
+        print(f"生图任务超时（{max_wait_time}秒）")
+        return None

dsLightRag/Liblib/T1_VersionGet.py

@@ -1,39 +1,13 @@
 from .LiblibUtil import LiblibUtil
 
 
-def get_model_version_info(version_uuid):
-    """获取模型版本信息"""
-    liblib_util = LiblibUtil()
-    uri = "/api/model/version/get"
-    payload = {"versionUuid": version_uuid}
-    
-    model_info = liblib_util.post_request(uri, payload)
-    if model_info:
-        return {
-            'modelName': model_info.get('modelName'),
-            'versionName': model_info.get('versionName'),
-            'commercialUse': model_info.get('commercialUse'),
-            'modelUrl': model_info.get('modelUrl')
-        }
-    return None
-
-
-def generate_text_to_image(template_uuid, generate_params):
-    """调用text2img接口生成图片"""
-    liblib_util = LiblibUtil()
-    uri = "/api/generate/webui/text2img"
-    payload = {
-        "templateUuid": template_uuid,
-        "generateParams": generate_params
-    }
-    return liblib_util.post_request(uri, payload)
-
-
 if __name__ == '__main__':
+    liblib = LiblibUtil()
+    
     # 测试获取模型版本信息
     print("===== 测试获取模型版本信息 =====")
     version_uuid = "390e7df48aed45d4bc3f0bcbc89fa44c"
-    model_info = get_model_version_info(version_uuid)
+    model_info = liblib.get_model_version_info(version_uuid)
     
     if model_info:
         print(f"模型名称：{model_info['modelName']}")
@@ -57,7 +31,7 @@ if __name__ == '__main__':
         "imgCount": 1
     }
     
-    image_data = generate_text_to_image(template_uuid, generate_params)
+    image_data = liblib.generate_text_to_image(template_uuid, generate_params)
     if image_data:
         print(f"图片生成成功，任务UUID: {image_data.get('generateUuid')}")
     else:

dsLightRag/Liblib/T2.py

@@ -1,57 +1,20 @@
 from .LiblibUtil import LiblibUtil
-import time
-
-
-def get_generation_status(generate_uuid):
-    """查询生图任务状态和结果"""
-    liblib_util = LiblibUtil()
-    uri = "/api/generate/webui/status"
-    payload = {"generateUuid": generate_uuid}
-    return liblib_util.post_request(uri, payload)
-
-
-def wait_for_generation_completion(generate_uuid, check_interval=5, max_wait_time=300):
-    """等待生图任务完成并获取最终结果"""
-    liblib_util = LiblibUtil()
-    start_time = time.time()
-    
-    while time.time() - start_time < max_wait_time:
-        status_data = get_generation_status(generate_uuid)
-        
-        if status_data:
-            status = status_data.get('status')
-            progress = status_data.get('percentCompleted', 0) * 100
-            print(f"生成进度：{progress:.2f}%")
-            
-            if status == 'completed':
-                print("生图任务已完成")
-                images = status_data.get('images', [])
-                if images:
-                    print("生成的最终图片：")
-                    for img in images:
-                        print(f"- {img.get('imageUrl')}")
-                return status_data
-            elif status == 'failed':
-                print(f"生图任务失败: {status_data.get('errorMessage', '未知错误')}")
-                return None
-            else:
-                print(f"生图任务处理中，当前状态: {status}")
-        
-        time.sleep(check_interval)
-    
-    print(f"生图任务超时（{max_wait_time}秒）")
-    return None
 
 
 if __name__ == '__main__':
+    liblib = LiblibUtil()
     print("===== 测试查询生图结果 =====")
     test_generate_uuid = "061d402be5af492d9949d53719912737"
     
     if test_generate_uuid:
-        status_data = get_generation_status(test_generate_uuid)
+        # 直接查询状态
+        status_data = liblib.get_generation_status(test_generate_uuid)
         if status_data:
             print(f"生图任务状态：{status_data}")
         else:
             print("查询生图任务状态失败")
+        
+        # 等待任务完成（如需测试请取消注释）
+        # liblib.wait_for_generation_completion(test_generate_uuid)
     else:
         print("请先设置有效的生图任务UUID")