Python调用ComfyOne示例

# -*- coding: utf-8 -*-
import requests
import json
import base64
from PIL import Image
import io
import time
import os

# This is a demo script showcasing the basic functionality of ComfyOne API
# It demonstrates how to:
# - Register a new backend instance with the ComfyOne service
# - Handle API responses and error cases
# - Implement retry logic for robust API communication
#
# Usage:
#   Simply run this script directly: python generate_image.py
#   The script will create a demo instance and attempt to register it with the service

# Constants
BASE_URL = "https://pandora-server-cf.onethingai.com"

# 请求 comfyone API
def request_comfyone_api(api: str, payload: dict, method: str, api_key: str) -> dict | None:
    """
    向ComfyOne API发送请求的通用函数
    
    Args:
        api: comfyone API路径
        payload: 请求体数据，可以为None
        method: HTTP请求方法(GET, POST等)
        api_key: onethingai 开发者密钥
        
    Returns:
        dict: API响应数据
        None: 请求失败时返回None
    """
    url = f"{BASE_URL}/{api}"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    try:
        retry_count = 0
        max_retries = 3
        while retry_count < max_retries:
            try:
                if payload:
                    response = requests.request(method, url, headers=headers, json=payload, timeout=5)
                else:
                    response = requests.request(method, url, headers=headers, timeout=5)
                response.raise_for_status()
                return response.json()
            except requests.exceptions.Timeout:
                retry_count += 1
                if retry_count == max_retries:
                    raise
                time.sleep(2 ** retry_count)  # Exponential backoff
                continue
            except requests.exceptions.RequestException as e:
                print(str(e))
                retry_count += 1
                if retry_count == max_retries:
                    return {"error": str(e)}
                time.sleep(2 ** retry_count)
                continue
    except Exception as e:
        print(f"Error requesting {api}: {e}")
        return None

# 查询可用后端
def get_available_backends(api_key: str) -> dict | None:
    """
    查询所有可用的ComfyOne后端服务实例
    
    Args:
        api_key (str): onethingai开发者密钥
        
    Returns:
        dict: 包含可用后端列表的响应数据
        None: 查询失败时返回None
    """
    api = "v1/backends"
    headers = {
        "Authorization": f"Bearer {api_key}"
    }
    return request_comfyone_api(api, None, "GET", api_key)


# 注册实例
def register_comfyone_backend(instance_id: str, api_key: str) -> str | None:
    """
    注册一个新的实例作为ComfyOne后端服务。
    该函数将提供的实例ID添加到ComfyOne后端管理系统中，
    使其能够处理图像生成请求和工作流处理。
    
    参数:
        instance_id (str): 要注册的实例的唯一标识符
        api_key (str): onethingai开发者密钥
        
    返回:
        str或None: 如果成功注册为后端则返回实例ID，失败则返回None
    """
    api = "v1/backends"
    headers = {
        "Content-Type": "application/json", 
        "Authorization": f"Bearer {api_key}"
    }
    payload = {
        "instance_id": instance_id
    }
    return request_comfyone_api(api, payload, "POST", api_key)

# 创建工作流
def create_workflow(api_key: str, name: str, inputs: dict, outputs: dict, workflow_file: str):
    """
    创建一个新的工作流。
    该函数将创建一个新的工作流配置，包含输入参数、输出节点和工作流定义。
    
    参数:
        api_key (str): onethingai开发者密钥
        name (str): 工作流名称
        inputs (dict): 工作流输入参数配置
        outputs (dict): 工作流输出节点配置
        workflow_file (str): 工作流定义文件路径
        
    返回:
        dict: API响应数据，包含工作流ID等信息
        None: 创建失败时返回None
    """
    api = "v1/workflows"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    with open(workflow_file, 'r', encoding='utf-8-sig') as f:
        #print(f.read())
        workflow_data = json.load(f)
    payload = {
        "name": name,
        "inputs": inputs,
        "outputs": outputs,
        "workflow": workflow_data
    }
    return request_comfyone_api(api, payload, "POST", api_key)


def prompt_comfyone(payload: dict, api_key: str) -> str | None:
    """
    向ComfyOne API发送prompt请求。
    该函数用于向已创建的工作流发送具体的生成参数。

    参数:
        payload (dict): prompt请求参数,包含工作流ID和具体的生成参数
        api_key (str): onethingai开发者密钥:w
    
    返回:
        str或None: 如果请求成功返回任务ID,失败返回None
    """
    # API endpoint
    api = f"v1/prompts"
    # Request headers
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    return request_comfyone_api(api, payload, "POST", api_key)

# Usage example
if __name__ == "__main__":
    # 以下演示了ComfyOne API的基本使用流程:
    # 1. 设置API密钥和实例ID，确保API密钥和实例ID正确，注意密钥安全，不要泄露
    # 2. 查询可用后端,如果实例未注册则进行注册，确保至少有一个后端可用
    # 3. 创建工做流
    # 4. 构造prompt参数并调用API生成结果

    API_KEY = "onethingai-api-key"  # 替换为你的实际API密钥
    INSTANCE_ID = "onethingai-instance-id"  # 替换为你的实际实例ID

    inputs = [
        {'id': '5', 'type': 'number', 'name': 'height'}, # 输入节点的字段配置
        {'id': '5', 'type': 'number', 'name': 'width'}]  # 注意： 输入节点字段配置需要和工作流json文件中一致
    
    outputs = ['9'] # 需要监控的输出节点的ID

    # 查询可用后端
    backends = get_available_backends(API_KEY)
    print(backends)

    # 如果实例没有注册到ComfyOne，则注册
    is_registered = False
    for backend in backends['data']:
        if backend['name'] == INSTANCE_ID:
            is_registered = True
            break
    if not is_registered:
        register_result = register_comfyone_backend(INSTANCE_ID, API_KEY)
        if register_result is not None and register_result['code'] == 0:
            print(f"注册实例成功: {register_result['data']['name']}")
        else:
            print(f"注册实例失败: {register_result['code']}, {register_result['msg']}")
            exit(1)

    # 创建工作流, 创建工作流后， 工作流ID会返回
    create_workflow_result = create_workflow(API_KEY, "test", inputs, outputs, "test_flow.json")

    # 生成图片
    if create_workflow_result is not None and create_workflow_result['code'] == 0:
        prompt_params = {
        "workflow_id": create_workflow_result['data']['id'], # 工作流ID，需要创建工作流然后再Prompt
       "inputs": [
                {
                "id": "5",
                "params": {
                    "width":  1024,
                    "height": 1024
                }
            }
        ],
        "free_cache": False
    }
        prompt_result = prompt_comfyone(prompt_params, API_KEY)
        if prompt_result is not None and prompt_result['code'] == 0:    
            print(prompt_result)
        else:
            print(f"生成图片失败: {prompt_result['code']}, {prompt_result['msg']}")
            exit(1)
    else:
        print(f"创建工作流失败: {create_workflow_result['code']}, {create_workflow_result['msg']}")
        exit(1)