import os import sys import time import requests import json from typing import Dict, Any sys.path.append(os.path.join(os.path.dirname(__file__), '../tools/local/liblibai_controlnet')) from liblibai_client import LibLibAIClient # 常量配置 (根据文档) TEMPLATE_UUID = "e10adc3949ba59abbe56e057f20f883e" # SD1.5 & SDXL 通用自定义参数模板 # 请确保这是一个有效的 SDXL 模型，这样才能匹配底下的 SDXL Canny 模型 # 这里用的是代码里原本的 Checkpoint ID，请根据你们自己的 Liblib 模型库调整！ DEFAULT_CHECKPOINT_ID = "0ea388c7eb854be3ba3c6f65aac6bfd3" class LibLibTestRunner: def __init__(self): self.client = LibLibAIClient() self.models = self._load_models_from_json() # 动态匹配基础算法 XL 的各种控制网模型 self.sdxl_canny = self._get_model_uuid("线稿类", "Canny（硬边缘）", xl_only=True) or "b6806516962f4e1599a93ac4483c3d23" self.sdxl_softedge = self._get_model_uuid("线稿类", "SoftEdge（软边缘）", xl_only=True) or "dda1a0c480bfab9833d9d9a1e4a71fff" self.sdxl_lineart = self._get_model_uuid("线稿类", "Lineart（线稿）", xl_only=True) or "a0f01da42bf48b0ba02c86b6c26b5699" self.sdxl_openpose = self._get_model_uuid("姿态类", "OpenPose（姿态）", xl_only=True) or "2fe4f992a81c5ccbdf8e9851c8c96ff2" self._verify_models() def _load_models_from_json(self): json_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/liblibai_controlnet_models.json")) if os.path.exists(json_path): with open(json_path, "r", encoding="utf-8") as f: return json.load(f) return {} def _get_model_uuid(self, category, subtype, xl_only=True): if category in self.models and subtype in self.models[category]: for model in self.models[category][subtype]: if xl_only and model["base_algorithm"] != "基础算法 XL": continue return model["uuid"] return None def _verify_models(self): print("="*50) print("校验底模信息 (使用 api/model/version/get)") print("="*50) models_to_verify = { "CheckPoint / 底模": DEFAULT_CHECKPOINT_ID, } for name, uuid in models_to_verify.items(): info = self.client.get_model_version_info(uuid) status = f"{info.get('model_name', 'Unknown')} (Base: {info.get('baseAlgo', 'Unknown')})" if info else "Failed to verify" print(f"- {name}: [{uuid}] -> {status}") print("\n") def _submit_task(self, payload: dict) -> str: url = self.client.generate_auth_url("/api/generate/webui/text2img") print(f"Submitting payload...") resp = requests.post(url, json=payload, timeout=10) data = resp.json() if data.get("code") != 0: raise Exception(f"Submit task failed: {data.get('msg')} (code: {data.get('code')})") return data["data"]["generateUuid"] def _wait_and_print_result(self, task_id: str): print(f"Task submitted successfully! Task ID: {task_id}") print("Waiting for result...") timeout = 300 start_time = time.time() while time.time() - start_time < timeout: task_data = self.client.query_task_status(task_id) status = task_data.get("generateStatus") if status == 5: # Success images = [img["imageUrl"] for img in task_data.get("images", [])] print(f"\n[SUCCESSS] Generated images: {images}") # --- 新增: 自动下载图片 --- output_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "output")) os.makedirs(output_dir, exist_ok=True) for i, img_url in enumerate(images): try: img_resp = requests.get(img_url) img_resp.raise_for_status() timestamp = int(time.time()) filename = f"workflow_result_{timestamp}_{i}.png" filepath = os.path.join(output_dir, filename) with open(filepath, "wb") as f: f.write(img_resp.content) print(f"[{i}] 已自动保存到本地: {filepath}") except Exception as e: print(f"图片自动下载失败: {e}") return elif status in [6, 7]: # Failed or Cancelled print(f"\n[FAILED] Task failed. Audit status: {task_data.get('auditStatus')}") return print(".", end="", flush=True) time.sleep(5) print("\n[TIMEOUT] Wait for task timed out.") # 1. 纯文生图测试 def test_text2img(self): print("\n" + "="*50) print("TEST 1: 纯文生图 (Text to Image)") print("="*50) payload = { "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "prompt": "1girl, cute, beautiful landscape, masterpiece", "negativePrompt": "lowres, bad anatomy, text, error", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1 } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 2. 纯图生图测试 def test_img2img(self): print("\n" + "="*50) print("TEST 2: 纯图生图 (Image to Image)") print("="*50) test_img_url = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png" payload = { "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "mode": 0, # 0 代表图生图 "sourceImage": test_img_url, "denoisingStrength": 0.75, # 去噪强度 "prompt": "white line art, cat", "negativePrompt": "lowres, bad anatomy, error", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1 } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 3. 文生图 + ControlNet (Canny) 测试 def test_text2img_controlnet(self): print("\n" + "="*50) print("TEST 3: 文生图 + ControlNet (Text2Img with ControlNet Canny)") print("="*50) print("注意: 如果默认 Checkpoint 是 SD1.5, 下面配置的 SDXL Canny 模型可能会导致访问拒绝!") test_img_url = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png" payload = { "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "prompt": "simple white line art, cat, black background", "negativePrompt": "lowres, bad anatomy", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1, "controlNet": [{ "unitOrder": 1, "sourceImage": test_img_url, "width": 512, "height": 512, "preprocessor": 1, # 1 = Canny "model": self.sdxl_canny, "controlWeight": 1.0, "startingControlStep": 0.0, "endingControlStep": 1.0, "pixelPerfect": 1, "controlMode": 0, "annotationParameters": { "canny": { "preprocessorResolution": 512, "lowThreshold": 100, "highThreshold": 200 } } }] } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 4. 文生图 + 边缘 (SoftEdge/HED) def test_text2img_softedge(self): print("\n" + "="*50) print("TEST 4: 文生图 + SoftEdge (软边缘 控制网)") print("="*50) # 测试用图片 test_img_url = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png" payload = { "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "prompt": "Soft edge artwork, beautiful soft lighting, cat", "negativePrompt": "lowres, bad anatomy", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1, "controlNet": [{ "unitOrder": 1, "sourceImage": test_img_url, "width": 512, "height": 512, "preprocessor": 5, # 5 = HED / 软边缘 "model": self.sdxl_softedge, "controlWeight": 1.0, "startingControlStep": 0.0, "endingControlStep": 1.0, "pixelPerfect": 1, "controlMode": 0, "annotationParameters": { "hed": { "preprocessorResolution": 512 } } }] } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 5. 文生图 + 线稿 (Lineart) def test_text2img_lineart(self): print("\n" + "="*50) print("TEST 5: 文生图 + Lineart (线稿 控制网)") print("="*50) test_img_url = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png" payload = { "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "prompt": "Detailed coloring, colorful fantasy style, masterpiece, cat", "negativePrompt": "lowres", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1, "controlNet": [{ "unitOrder": 1, "sourceImage": test_img_url, "width": 512, "height": 512, "preprocessor": 32, # 32 = Lineart Standard "model": self.sdxl_lineart, "controlWeight": 1.0, "startingControlStep": 0.0, "endingControlStep": 1.0, "pixelPerfect": 1, "controlMode": 0, "annotationParameters": { "lineart": { "preprocessorResolution": 512 } } }] } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 6. 文生图 + 骨骼 (OpenPose) def test_text2img_openpose(self): print("\n" + "="*50) print("TEST 6: 文生图 + OpenPose (骨骼 控制网)") print("="*50) # 测试用图片（最好使用含有人物动作的图） test_img_url = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png" payload = { "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "prompt": "1girl, dancing pose, beautiful dress", "negativePrompt": "lowres", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1, "controlNet": [{ "unitOrder": 1, "sourceImage": test_img_url, "width": 512, "height": 512, "preprocessor": 14, # 14 = OpenPose Full "model": self.sdxl_openpose, "controlWeight": 1.0, "startingControlStep": 0.0, "endingControlStep": 1.0, "pixelPerfect": 1, "controlMode": 0, "annotationParameters": { "openposeFull": { "preprocessorResolution": 512 } } }] } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 7. 局部重绘 (Inpaint Mode 4) def test_inpaint_mode4(self): print("\n" + "="*50) print("TEST 7: 局部重绘 (Inpaint Mode 4 蒙版重绘)") print("="*50) test_img_url = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png" test_mask_url = test_img_url # 仅作演示，实际应用中应当是一个黑底白色的蒙版图片 payload = { # Inpainting 可能需要特殊模板，如不兼容请参考文档 "templateUuid": TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, "mode": 4, # 4 = Inpaint 蒙版重绘 "sourceImage": test_img_url, "denoisingStrength": 0.5, "prompt": "A completely different background, sci-fi city", "negativePrompt": "lowres", "sampler": 15, "steps": 20, "cfgScale": 7.0, "width": 512, "height": 512, "imgCount": 1, "inpaintParam": { "maskImage": test_mask_url, "maskBlur": 4, "inpaintArea": 0 } } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) # 8. 人像换脸 (InstantID) def test_instantid_faceswap(self): print("\n" + "="*50) print("TEST 8: 人像换脸 (InstantID)") print("="*50) # Note: InstantID faceswap uses a UNIQUE templateUuid INSTANT_ID_TEMPLATE_UUID = "7d888009f81d4252a7c458c874cd017f" face_img_url = "https://liblibai-online.liblib.cloud/img/081e9f07d9bd4c2ba090efde163518f9/49943c0b-4d79-4e2f-8c55-bc1e5b8c69d8.png" pose_img_url = "https://liblibai-online.liblib.cloud/img/081e9f07d9bd4c2ba090efde163518f9/e713676d-baaa-4dac-99b9-d5d814a29f9f.png" payload = { "templateUuid": INSTANT_ID_TEMPLATE_UUID, "generateParams": { "checkPointId": DEFAULT_CHECKPOINT_ID, # 仅 XL模型支持人像换脸 "prompt": "Asian portrait,A young woman wearing a green baseball cap, close shot, background is coffee store, masterpiece, best quality, ultra resolution", "width": 768, "height": 1152, "sampler": 20, "steps": 35, "cfgScale": 2.0, "imgCount": 1, "controlNet": [ { "unitOrder": 1, # 第一步：先识别要用的人像人脸 "sourceImage": face_img_url, "width": 1080, "height": 1432 }, { "unitOrder": 2, # 第二步：再识别要参考的人物面部朝向 "sourceImage": pose_img_url, "width": 1024, "height": 1024 } ] } } task_id = self._submit_task(payload) self._wait_and_print_result(task_id) if __name__ == "__main__": try: runner = LibLibTestRunner() print("选择要测试的模式:") print("1. 纯文生图 (Text2Img)") print("2. 纯图生图 (Img2Img)") print("3. 硬边缘控制 (Canny ControlNet)") print("4. 软边缘控制 (SoftEdge ControlNet)") print("5. 线稿控制 (Lineart ControlNet)") print("6. 骨骼控制 (OpenPose ControlNet)") print("7. 局部重绘 (Inpaint Mode=4)") print("8. 人像换脸 (InstantID)") print("9. 全部测试 (All)") if len(sys.argv) > 1: choice = sys.argv[1] else: # 默认测试文生图以验证 API Key 最基本权限 choice = "1" if choice == "1": runner.test_text2img() elif choice == "2": runner.test_img2img() elif choice == "3": runner.test_text2img_controlnet() elif choice == "4": runner.test_text2img_softedge() elif choice == "5": runner.test_text2img_lineart() elif choice == "6": runner.test_text2img_openpose() elif choice == "7": runner.test_inpaint_mode4() elif choice == "8": runner.test_instantid_faceswap() elif choice == "9": runner.test_text2img() runner.test_img2img() runner.test_text2img_controlnet() # 省略后面的测试避免同时发生太多请求... else: print("Unknown choice.") except Exception as e: print(f"\n[FATAL ERROR] {e}")