Tool_Agent/tests/test_local_openpose.py

import os
import sys
import time
import base64
import requests

sys.path.append(os.path.join(os.path.dirname(__file__), '../tools/local/liblibai_controlnet'))
from liblibai_client import LibLibAIClient

TEMPLATE_UUID = "e10adc3949ba59abbe56e057f20f883e"
DEFAULT_CHECKPOINT_ID = "0ea388c7eb854be3ba3c6f65aac6bfd3" 

# 这里填入您实际能用的 OpenPose SDXL 模型的 UUID
OPENPOSE_MODEL_ID = "b6806516962f4e1599a93ac4483c3d23" # 注意：这里默认放的是Canny，请记得更换成 OpenPose UUID


def get_base64_image(image_path):
    with open(image_path, "rb") as f:
        return "data:image/png;base64," + base64.b64encode(f.read()).decode("utf-8")

def main():
    print("=" * 50)
    print("测试: 读取本地图片并自动上传 -> 使用 OpenPose 生成图片")
    print("=" * 50)

    client = LibLibAIClient()
    
    # 1. 读取本地图片并转换为 Base64
    local_image_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "features/pose_skeleton/img_1_openpose.png"))
    print(f"Reading local image: {local_image_path}")
    
    if not os.path.exists(local_image_path):
        print(f"Error: 找不到图片 {local_image_path}")
        sys.exit(1)
        
    base64_image = get_base64_image(local_image_path)
    
    # 2. 借助已经写好的 `upload_base64_image` 上传到 OSS
    print("Uploading Base64 to Liblib OSS...")
    image_url = client.upload_base64_image(base64_image)
    print(f"Uploaded successfully! Public URL: {image_url}")
    
    # 3. 发送任务
    payload = {
        "templateUuid": TEMPLATE_UUID,
        "generateParams": {
            "checkPointId": DEFAULT_CHECKPOINT_ID,
            "prompt": "1girl, highly detailed, masterpieces, beautiful face, standing pose",
            "negativePrompt": "lowres, bad anatomy",
            "sampler": 15,
            "steps": 20,
            "cfgScale": 7.0,
            "width": 512,
            "height": 512,
            "imgCount": 1,
            "controlNet": [{
                "unitOrder": 1,
                "sourceImage": image_url,
                "width": 512,
                "height": 512,
                "preprocessor": 14,        # 14 = OpenPose Full
                "model": OPENPOSE_MODEL_ID, # !!需更换为OpenPose模型UUID
                "controlWeight": 1.0,
                "startingControlStep": 0.0,
                "endingControlStep": 1.0,
                "pixelPerfect": 1,
                "controlMode": 0,
                "annotationParameters": {
                    "openposeFull": {
                        "preprocessorResolution": 512
                    }
                }
            }]
        }
    }
    
    auth_url = client.generate_auth_url("/api/generate/webui/text2img")
    print(f"Submitting OpenPose payload...")
    resp = requests.post(auth_url, json=payload, timeout=10)
    data = resp.json()
    if data.get("code") != 0:
        print(f"Submit task failed: {data.get('msg')} (code: {data.get('code')})")
        sys.exit(1)
        
    task_id = data["data"]["generateUuid"]
    print(f"Task submitted! Task ID: {task_id}")
    
    timeout = 300
    start_time = time.time()
    while time.time() - start_time < timeout:
        task_data = client.query_task_status(task_id)
        status = task_data.get("generateStatus")
        
        if status == 5:
            images = [img["imageUrl"] for img in task_data.get("images", [])]
            print(f"\n[SUCCESSS] Generated images: {images}")
            
            # --- 新增: 自动下载图片到 output 文件夹 ---
            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"openpose_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"[{i}] 图片自动下载失败: {e}")
                    
            return
        elif status in [6, 7]:
            print("\n[FAILED] Task failed.")
            return
            
        print(".", end="", flush=True)
        time.sleep(5)
        
    print("\n[TIMEOUT]")

if __name__ == "__main__":
    main()