Tool_Agent/tests/test_liblibai_workflows.py

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}")