Go调用ComfyOne示例

示例工作流文件test_flow.json （与Go脚本文件，放在同一目录下）

package main

import (
    "bytes"
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "os"
    "time"
)

// Constants
const baseURL = "https://pandora-server-cf.onethingai.com"

// Response structs
type APIResponse struct {
    Code int         `json:"code"`
    Msg  string      `json:"msg"`
    Data interface{} `json:"data"`
}

type Backend struct {
    Name string `json:"name"`
}

// Helper function for API requests
func requestComfyOneAPI(api string, payload interface{}, method string, apiKey string) (*APIResponse, error) {
    url := fmt.Sprintf("%s/%s", baseURL, api)
    maxRetries := 3

    var jsonData []byte
    var err error
    if payload != nil {
        jsonData, err = json.Marshal(payload)
        if err != nil {
            return nil, fmt.Errorf("error marshaling payload: %v", err)
        }
    }

    for retry := 0; retry < maxRetries; retry++ {
        client := &http.Client{
            Timeout: 5 * time.Second,
        }

        var req *http.Request
        if payload != nil {
            req, err = http.NewRequest(method, url, bytes.NewBuffer(jsonData))
        } else {
            req, err = http.NewRequest(method, url, nil)
        }
        if err != nil {
            return nil, err
        }

        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))

        resp, err := client.Do(req)
        if err != nil {
            if retry == maxRetries-1 {
                return nil, err
            }
            time.Sleep(time.Duration(1<<retry) * time.Second)
            continue
        }
        defer resp.Body.Close()

        body, err := io.ReadAll(resp.Body)
        if err != nil {
            return nil, err
        }

        var apiResp APIResponse
        if err := json.Unmarshal(body, &apiResp); err != nil {
            return nil, err
        }

        return &apiResp, nil
    }

    return nil, fmt.Errorf("max retries exceeded")
}

// Main API functions
func getAvailableBackends(apiKey string) (*APIResponse, error) {
    return requestComfyOneAPI("v1/backends", nil, "GET", apiKey)
}

func registerComfyOneBackend(instanceID string, apiKey string) (*APIResponse, error) {
    payload := map[string]string{
        "instance_id": instanceID,
    }
    return requestComfyOneAPI("v1/backends", payload, "POST", apiKey)
}

func createWorkflow(name string, inputs interface{}, outputs interface{}, workflowFile string, apiKey string) (*APIResponse, error) {
    // Read workflow file
    workflowData, err := os.ReadFile(workflowFile)
    if err != nil {
        return nil, err
    }
    // Remove UTF-8 BOM if present (0xEF,0xBB,0xBF)
    if len(workflowData) >= 3 && workflowData[0] == 0xEF && workflowData[1] == 0xBB && workflowData[2] == 0xBF {
        workflowData = workflowData[3:]
    }

    var workflow interface{}
    if err := json.Unmarshal(workflowData, &workflow); err != nil {
        return nil, err
    }
    payload := map[string]interface{}{
        "name":     name,
        "inputs":   inputs,
        "outputs":  outputs,
        "workflow": workflow,
    }

    return requestComfyOneAPI("v1/workflows", payload, "POST", apiKey)
}

func promptComfyOne(payload interface{}, apiKey string) (*APIResponse, error) {
    return requestComfyOneAPI("v1/prompts", payload, "POST", apiKey)
}

func main() {
    apiKey := "onethingai-api-key"         // Replace with your actual API key
    instanceID := "onethingai-instance-id" // Replace with your actual instance ID

    inputs := []map[string]interface{}{
        {"id": "5", "type": "number", "name": "height"},
        {"id": "5", "type": "number", "name": "width"},
    }
    outputs := []string{"9"}

    // Query available backends
    backends, err := getAvailableBackends(apiKey)
    if err != nil {
        fmt.Printf("Error getting backends: %v\n", err)
        os.Exit(1)
    }

    // Check if instance is registered
    isRegistered := false
    if backendList, ok := backends.Data.([]interface{}); ok {
        for _, b := range backendList {
            if backend, ok := b.(map[string]interface{}); ok {
                if backend["name"] == instanceID {
                    isRegistered = true
                    break
                }
            }
        }
    }

    // Register if not already registered
    if !isRegistered {
        registerResult, err := registerComfyOneBackend(instanceID, apiKey)
        if err != nil || registerResult.Code != 0 {
            fmt.Printf("Failed to register instance: %v\n", err)
            os.Exit(1)
        }
        fmt.Printf("Successfully registered instance\n")
    }

    // Create workflow
    workflowResult, err := createWorkflow("test", inputs, outputs, "test_flow.json", apiKey)
    if err != nil || workflowResult.Code != 0 {
        fmt.Printf("Failed to create workflow: %v\n", err)
        os.Exit(1)
    }

    // Generate image
    workflowData := workflowResult.Data.(map[string]interface{})
    promptParams := map[string]interface{}{
        "workflow_id": workflowData["id"],
        "inputs": []map[string]interface{}{
            {
                "id": "5",
                "params": map[string]interface{}{
                    "width":  1024,
                    "height": 1024,
                },
            },
        },
        "free_cache": false,
    }

    promptResult, err := promptComfyOne(promptParams, apiKey)
    if err != nil || promptResult.Code != 0 {
        fmt.Printf("Failed to generate image: %v\n", err)
        os.Exit(1)
    }

    fmt.Printf("Success: %+v\n", promptResult)
}