Luminary/internal/provider/comfyui/client.go

package comfyui

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

type Client struct {
	BaseURL    string
	HTTPClient *http.Client
}

func New(baseURL string) *Client {
	return &Client{
		BaseURL:    strings.TrimRight(baseURL, "/"),
		HTTPClient: &http.Client{Timeout: 120 * time.Second},
	}
}

type PromptResponse struct {
	PromptID string `json:"prompt_id"`
	Number   int    `json:"number"`
}

func (c *Client) SubmitPrompt(ctx context.Context, prompt map[string]any, clientID string) (*PromptResponse, error) {
	body, _ := json.Marshal(map[string]any{
		"prompt":    prompt,
		"client_id": clientID,
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/prompt", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("comfyui prompt %d: %s", resp.StatusCode, string(data))
	}
	var out PromptResponse
	if err := json.Unmarshal(data, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

func (c *Client) Interrupt(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/interrupt", nil)
	if err != nil {
		return err
	}
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		data, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("comfyui interrupt %d: %s", resp.StatusCode, string(data))
	}
	return nil
}

type HistoryEntry struct {
	Outputs map[string]HistoryNodeOutput `json:"outputs"`
	Status  map[string]any               `json:"status"`
}

type HistoryNodeOutput struct {
	Images []OutputImage `json:"images"`
}

type OutputImage struct {
	Filename  string `json:"filename"`
	Subfolder string `json:"subfolder"`
	Type      string `json:"type"`
}

func (c *Client) GetHistory(ctx context.Context, promptID string) (map[string]HistoryEntry, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.BaseURL+"/history/"+promptID, nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("comfyui history %d: %s", resp.StatusCode, string(data))
	}
	var out map[string]HistoryEntry
	if err := json.Unmarshal(data, &out); err != nil {
		return nil, err
	}
	return out, nil
}

func (c *Client) DownloadView(ctx context.Context, filename, subfolder, imageType string) ([]byte, string, error) {
	q := url.Values{}
	q.Set("filename", filename)
	if subfolder != "" {
		q.Set("subfolder", subfolder)
	}
	if imageType == "" {
		imageType = "output"
	}
	q.Set("type", imageType)
	u := c.BaseURL + "/view?" + q.Encode()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
	if err != nil {
		return nil, "", err
	}
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, "", err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, "", err
	}
	if resp.StatusCode >= 400 {
		return nil, "", fmt.Errorf("comfyui view %d", resp.StatusCode)
	}
	mime := resp.Header.Get("Content-Type")
	if mime == "" {
		mime = "image/png"
	}
	return data, mime, nil
}

func (c *Client) UploadImage(ctx context.Context, data []byte, filename string) (string, error) {
	if len(data) == 0 {
		return "", fmt.Errorf("empty image payload")
	}
	if strings.TrimSpace(filename) == "" {
		filename = "input.png"
	}
	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)
	_ = w.WriteField("type", "input")
	_ = w.WriteField("overwrite", "true")
	part, err := w.CreateFormFile("image", filename)
	if err != nil {
		return "", err
	}
	if _, err := part.Write(data); err != nil {
		return "", err
	}
	if err := w.Close(); err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/upload/image", &buf)
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", w.FormDataContentType())
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 400 {
		return "", fmt.Errorf("comfyui upload %d: %s", resp.StatusCode, string(body))
	}
	var out struct {
		Name string `json:"name"`
	}
	if err := json.Unmarshal(body, &out); err != nil {
		return "", err
	}
	return out.Name, nil
}

func (c *Client) HealthCheck(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.BaseURL+"/system_stats", nil)
	if err != nil {
		return err
	}
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		return fmt.Errorf("status %d", resp.StatusCode)
	}
	return nil
}

type ProgressEvent struct {
	Type      string
	Data      map[string]any
	Node      string
	Value     float64
	Max       float64
	PromptID  string
	ErrorText string
}

func (c *Client) wsURL(clientID string) string {
	u, _ := url.Parse(c.BaseURL)
	scheme := "ws"
	if u.Scheme == "https" {
		scheme = "wss"
	}
	return fmt.Sprintf("%s://%s/ws?clientId=%s", scheme, u.Host, url.QueryEscape(clientID))
}

func (c *Client) Watch(ctx context.Context, clientID string, onEvent func(ProgressEvent)) error {
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, c.wsURL(clientID), nil)
	if err != nil {
		return err
	}
	defer conn.Close()
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				return
			}
			var envelope struct {
				Type string          `json:"type"`
				Data json.RawMessage `json:"data"`
			}
			if err := json.Unmarshal(msg, &envelope); err != nil {
				continue
			}
			ev := ProgressEvent{Type: envelope.Type}
			var data map[string]any
			_ = json.Unmarshal(envelope.Data, &data)
			ev.Data = data
			if envelope.Type == "progress" {
				if n, ok := data["node"].(string); ok {
					ev.Node = n
				}
				if v, ok := data["value"].(float64); ok {
					ev.Value = v
				}
				if m, ok := data["max"].(float64); ok {
					ev.Max = m
				}
			}
			if envelope.Type == "executing" {
				if n, ok := data["node"].(string); ok {
					ev.Node = n
				}
				if pid, ok := data["prompt_id"].(string); ok {
					ev.PromptID = pid
				}
			}
			if envelope.Type == "execution_error" {
				if msgs, ok := data["exception_message"].(string); ok {
					ev.ErrorText = msgs
				}
			}
			onEvent(ev)
		}
	}()
	select {
	case <-ctx.Done():
		_ = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		<-done
		return ctx.Err()
	case <-done:
		return nil
	}
}

func CollectOutputImages(history map[string]HistoryEntry, promptID string) []OutputImage {
	entry, ok := history[promptID]
	if !ok {
		return nil
	}
	var images []OutputImage
	for _, out := range entry.Outputs {
		images = append(images, out.Images...)
	}
	return images
}