Files
llmgw/internal/comfy/comfy.go
T
Ilker ManapandClaude Opus 5.5 cfc636e328
CI / test (push) Failing after 6m28s
A mesh's base colour can be brought to the picture it was made from
Stable Fast 3D takes the light out of the picture it sees, and a dark
basalt (median 0.17) came back a light stone dimly lit (0.67). The save
node takes albedo_match, the picture and its mask, and scales the base
colour texture in linear light by the ratio of the object's medians, so
marks keep their contrast; texels nothing maps to are left alone. On the
game team's rock it is 0.18.

A workflow's {{name|false}} default now fills a BOOLEAN input with a
boolean rather than the word.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-14 00:51:06 +02:00

790 lines
22 KiB
Go

// Package comfy has images made by ComfyUI: a workflow filled in with what
// was asked for, queued, waited on, and its pictures fetched.
//
// ComfyUI speaks no OpenAI. A workflow is a graph of nodes that is queued with
// POST /prompt, is done when GET /history/{id} lists it, and leaves its
// pictures to be read one by one from GET /view. What the graph is depends on
// the model — FLUX schnell is not FLUX dev, and neither is SDXL — so the graph
// is a template the catalog names, with placeholders where a request's
// prompt, size and seed go.
package comfy
import (
"bytes"
"context"
"crypto/rand"
"embed"
"encoding/binary"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"math/big"
"mime/multipart"
"net/http"
"net/url"
"path"
"slices"
"strconv"
"strings"
"time"
)
//go:embed workflows/*.json
var builtin embed.FS
// DefaultWorkflow is the template used when the catalog names none.
const DefaultWorkflow = "flux-schnell"
// Builtins are the workflows that come with llmgw, by name.
func Builtins() []string {
entries, _ := builtin.ReadDir("workflows")
out := make([]string, 0, len(entries))
for _, e := range entries {
out = append(out, strings.TrimSuffix(e.Name(), ".json"))
}
return out
}
// Template is a workflow in ComfyUI's API format, with placeholders.
//
// A placeholder is a string value "{{name}}", or "{{name|default}}" for one a
// request may leave out. One that is the whole value becomes the value
// itself, a number where the node wants a number; one inside a longer string
// is written into it. The names are those of Params.
type Template struct {
graph map[string]any
}
// Load reads a template: a built-in's name, or the workflow's JSON itself.
func Load(nameOrJSON string) (*Template, error) {
s := strings.TrimSpace(nameOrJSON)
if s == "" {
s = DefaultWorkflow
}
raw := []byte(s)
if !strings.HasPrefix(s, "{") {
b, err := builtin.ReadFile("workflows/" + s + ".json")
if err != nil {
return nil, fmt.Errorf("no built-in workflow %q (there are %s)", s, strings.Join(Builtins(), ", "))
}
raw = b
}
var g map[string]any
if err := json.Unmarshal(raw, &g); err != nil {
return nil, fmt.Errorf("the workflow is not JSON: %w", err)
}
if len(g) == 0 {
return nil, errors.New("the workflow has no nodes")
}
for id, n := range g {
node, ok := n.(map[string]any)
if !ok || node["class_type"] == nil {
// The UI's own save format has a "nodes" list instead, and is
// not something /prompt accepts.
return nil, fmt.Errorf("node %q has no class_type: the workflow must be saved in API format", id)
}
}
return &Template{graph: g}, nil
}
// Params are what a request decides. A zero number is one not given, which
// leaves it to the template's default.
type Params struct {
Checkpoint string
Prompt string
NegativePrompt string
Width, Height int
N int
Seed int64
Steps int
Guidance float64
// Image is the input picture, by the name Upload gave it in ComfyUI's
// input directory: {{image}}.
Image string
// Extra are a workflow's own placeholders, by name: {{resolution}}.
Extra map[string]any
}
func (p Params) value(name string) (any, bool) {
// A workflow's own input of a built-in's name is the request's: an input
// named prompt is the prompt a run was given.
if v, ok := p.Extra[name]; ok {
return v, true
}
switch name {
case "checkpoint":
return p.Checkpoint, p.Checkpoint != ""
case "prompt":
return p.Prompt, true
case "negative_prompt":
return p.NegativePrompt, true
case "width":
return p.Width, p.Width > 0
case "height":
return p.Height, p.Height > 0
case "n":
return max(p.N, 1), true
case "seed":
return p.Seed, true
case "steps":
return p.Steps, p.Steps > 0
case "guidance":
return p.Guidance, p.Guidance > 0
case "image":
return p.Image, p.Image != ""
}
return nil, false
}
// Builtin are the placeholders Params fills by itself, which a workflow may
// use without declaring them as inputs.
var Builtin = []string{"checkpoint", "prompt", "negative_prompt", "width", "height", "n", "seed", "steps", "guidance", "image"}
// Placeholders are the names a template's placeholders use, each once,
// sorted, and whether each has a default of its own.
func (t *Template) Placeholders() map[string]bool {
out := map[string]bool{}
var walk func(v any)
walk = func(v any) {
switch v := v.(type) {
case map[string]any:
for _, x := range v {
walk(x)
}
case []any:
for _, x := range v {
walk(x)
}
case string:
for s := v; ; {
i := strings.Index(s, "{{")
j := strings.Index(s, "}}")
if i < 0 || j < i {
return
}
name, _, hasDef := strings.Cut(s[i+2:j], "|")
name = strings.TrimSpace(name)
out[name] = out[name] || hasDef
s = s[j+2:]
}
}
}
walk(t.graph)
return out
}
// Fill is the workflow for one request.
func (t *Template) Fill(p Params) (map[string]any, error) {
var missing []string
var fill func(v any) any
fill = func(v any) any {
switch v := v.(type) {
case map[string]any:
out := make(map[string]any, len(v))
for k, x := range v {
out[k] = fill(x)
}
return out
case []any:
out := make([]any, len(v))
for i, x := range v {
out[i] = fill(x)
}
return out
case string:
return fillString(v, p, &missing)
}
return v
}
g := fill(t.graph).(map[string]any)
if len(missing) > 0 {
slices.Sort(missing)
return nil, fmt.Errorf("the workflow wants %s, which the request did not give and the workflow has no default for",
strings.Join(slices.Compact(missing), ", "))
}
return g, nil
}
func fillString(s string, p Params, missing *[]string) any {
if !strings.Contains(s, "{{") {
return s
}
resolve := func(ph string) (any, bool) {
name, def, hasDef := strings.Cut(ph, "|")
name = strings.TrimSpace(name)
if v, ok := p.value(name); ok {
return v, true
}
if hasDef {
def = strings.TrimSpace(def)
if f, err := strconv.ParseFloat(def, 64); err == nil {
return f, true
}
// A BOOLEAN input refuses the word: {{albedo_match|false}}.
if b, err := strconv.ParseBool(def); err == nil && (def == "true" || def == "false") {
return b, true
}
return def, true
}
*missing = append(*missing, name)
return nil, false
}
// The whole value: what the node gets is the value, not its text.
if strings.HasPrefix(s, "{{") && strings.HasSuffix(s, "}}") && strings.Count(s, "{{") == 1 {
v, _ := resolve(s[2 : len(s)-2])
return v
}
var b strings.Builder
for {
i := strings.Index(s, "{{")
j := strings.Index(s, "}}")
if i < 0 || j < i {
b.WriteString(s)
return b.String()
}
b.WriteString(s[:i])
if v, ok := resolve(s[i+2 : j]); ok {
fmt.Fprint(&b, v)
}
s = s[j+2:]
}
}
// RandomSeed is a seed for a request that named none. ComfyUI's are 64-bit,
// but JSON numbers are read as doubles on the way: kept to 2^53.
func RandomSeed() int64 {
n, err := rand.Int(rand.Reader, big.NewInt(1<<53))
if err != nil {
return time.Now().UnixNano() & (1<<53 - 1)
}
return n.Int64()
}
// Image is one file a workflow made: a picture, a mesh.
type Image struct {
Data []byte
ContentType string
// Name is the file's own, as ComfyUI saved it.
Name string
}
// Error is a failure, said as what it means for the request.
type Error struct {
Msg string
// Retryable is a failure of this server, which another one might not
// have: it could not be reached, or it answered 5xx.
Retryable bool
}
func (e *Error) Error() string { return e.Msg }
// Client queues workflows and waits for them.
type Client struct {
HTTP *http.Client
// Poll is how often a queued workflow is asked after.
Poll time.Duration
}
// Generate runs a workflow on the ComfyUI at base and returns its pictures,
// in the order of the nodes that made them. A context that ends takes the
// workflow out of the queue, or interrupts it if it is running, so a client
// that went away does not leave the card working for nobody.
func (c *Client) Generate(ctx context.Context, base, apiKey string, workflow map[string]any) ([]Image, error) {
return c.GenerateWatch(ctx, base, apiKey, workflow, nil)
}
// progressEvery is how often a workflow's progress is asked after, beside
// its history: a few requests a second is for the end, not for the news.
const progressEvery = time.Second
// GenerateWatch is Generate, saying how far along the workflow is as it
// changes: queued and how many ahead, then the node running and its step.
func (c *Client) GenerateWatch(ctx context.Context, base, apiKey string, workflow map[string]any,
onProgress func(Progress)) ([]Image, error) {
base = strings.TrimRight(base, "/")
client := make([]byte, 8)
_, _ = rand.Read(client)
body, err := json.Marshal(map[string]any{"prompt": workflow, "client_id": "llmgw-" + hex.EncodeToString(client)})
if err != nil {
return nil, err
}
var queued struct {
PromptID string `json:"prompt_id"`
Error json.RawMessage `json:"error"`
NodeErrors json.RawMessage `json:"node_errors"`
}
status, err := c.do(ctx, http.MethodPost, base+"/prompt", apiKey, body, &queued)
if err != nil {
return nil, err
}
if status != http.StatusOK || queued.PromptID == "" {
// A workflow ComfyUI will not run: a checkpoint it does not have, a
// node it does not know. The catalog's to fix, not the client's, and
// no other replica set up the same way would take it either.
return nil, &Error{Msg: "ComfyUI refused the workflow: " + refusal(queued.Error, queued.NodeErrors)}
}
id := queued.PromptID
poll := c.Poll
if poll <= 0 {
poll = 250 * time.Millisecond
}
t := time.NewTicker(poll)
defer t.Stop()
var last Progress
var asked time.Time
for {
select {
case <-ctx.Done():
c.cancel(context.WithoutCancel(ctx), base, apiKey, id)
return nil, ctx.Err()
case <-t.C:
}
if onProgress != nil && time.Since(asked) >= progressEvery {
asked = time.Now()
if p, ok := c.progress(ctx, base, apiKey, id, workflow); ok && p != last {
last = p
onProgress(p)
}
}
var hist map[string]historyEntry
if _, err := c.do(ctx, http.MethodGet, base+"/history/"+url.PathEscape(id), apiKey, nil, &hist); err != nil {
if ctx.Err() != nil {
continue // said by the select
}
return nil, err
}
h, ok := hist[id]
if !ok {
continue
}
if h.Status.StatusStr == "error" {
return nil, &Error{Msg: "the workflow failed in ComfyUI: " + h.failure()}
}
if !h.Status.Completed {
continue
}
return c.fetch(ctx, base, apiKey, h)
}
}
type historyEntry struct {
Status struct {
StatusStr string `json:"status_str"`
Completed bool `json:"completed"`
Messages [][]json.RawMessage `json:"messages"`
} `json:"status"`
// Outputs are each node's, by kind: "images" for a picture, "3d" for
// the mesh a 3D node saves, and whatever else a custom node names.
Outputs map[string]map[string]json.RawMessage `json:"outputs"`
}
// fileRef is a file an output names, to be read from /view.
type fileRef struct {
Filename string `json:"filename"`
Subfolder string `json:"subfolder"`
Type string `json:"type"`
}
// failure is what ComfyUI said went wrong, from its execution_error.
func (h historyEntry) failure() string {
for _, m := range h.Status.Messages {
if len(m) < 2 || string(m[0]) != `"execution_error"` {
continue
}
var e struct {
NodeType string `json:"node_type"`
Message string `json:"exception_message"`
}
if json.Unmarshal(m[1], &e) == nil && e.Message != "" {
return strings.TrimSpace(e.NodeType + ": " + strings.TrimSpace(e.Message))
}
}
return "no reason given"
}
func (c *Client) fetch(ctx context.Context, base, apiKey string, h historyEntry) ([]Image, error) {
nodes := make([]string, 0, len(h.Outputs))
for id := range h.Outputs {
nodes = append(nodes, id)
}
// Numerically where they are numbers, which is how ComfyUI names them.
slices.SortFunc(nodes, func(a, b string) int {
x, ea := strconv.Atoi(a)
y, eb := strconv.Atoi(b)
if ea == nil && eb == nil {
return x - y
}
return strings.Compare(a, b)
})
var out []Image
for _, n := range nodes {
kinds := make([]string, 0, len(h.Outputs[n]))
for k := range h.Outputs[n] {
kinds = append(kinds, k)
}
slices.Sort(kinds)
for _, k := range kinds {
// Any list of files, whatever the node calls it; text and
// numbers a node reports beside them are not files.
var refs []fileRef
if json.Unmarshal(h.Outputs[n][k], &refs) != nil {
continue
}
for _, f := range refs {
if f.Filename == "" {
continue
}
q := url.Values{"filename": {f.Filename}, "subfolder": {f.Subfolder}, "type": {f.Type}}
img, err := c.get(ctx, base+"/view?"+q.Encode(), apiKey)
if err != nil {
return nil, err
}
img.Name = f.Filename
// A type the server did not say is the file's by its name,
// else by its first bytes.
if generic(img.ContentType) {
if t := typeByName(f.Filename); t != "" {
img.ContentType = t
} else {
img.ContentType = http.DetectContentType(img.Data)
}
}
if img.ContentType == "image/png" {
img.Data = withoutPNGText(img.Data)
}
out = append(out, img)
}
}
}
if len(out) == 0 {
return nil, &Error{Msg: "the workflow finished without a file: it has no node that saves one"}
}
return out, nil
}
// cancel takes a workflow out of the queue, or stops it if it is the one
// running. Best effort: the server may be gone, which also stops it.
func (c *Client) cancel(ctx context.Context, base, apiKey, id string) {
ctx, done := context.WithTimeout(ctx, 5*time.Second)
defer done()
body, _ := json.Marshal(map[string]any{"delete": []string{id}})
_, _ = c.do(ctx, http.MethodPost, base+"/queue", apiKey, body, nil)
var q struct {
Running [][]json.RawMessage `json:"queue_running"`
}
if _, err := c.do(ctx, http.MethodGet, base+"/queue", apiKey, nil, &q); err != nil {
return
}
for _, r := range q.Running {
if len(r) > 1 && string(r[1]) == strconv.Quote(id) {
_, _ = c.do(ctx, http.MethodPost, base+"/interrupt", apiKey, []byte("{}"), nil)
return
}
}
}
func (c *Client) client() *http.Client {
if c.HTTP != nil {
return c.HTTP
}
return http.DefaultClient
}
func (c *Client) request(ctx context.Context, method, u, apiKey string, body []byte) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, method, u, bytes.NewReader(body))
if err != nil {
return nil, err
}
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
if apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := c.client().Do(req)
if err != nil {
if ctx.Err() != nil {
return nil, ctx.Err()
}
return nil, &Error{Msg: "ComfyUI could not be reached: " + err.Error(), Retryable: true}
}
return resp, nil
}
// do sends a request and reads a JSON answer into out. A 5xx is an error;
// any other status is returned with whatever body could be read.
func (c *Client) do(ctx context.Context, method, u, apiKey string, body []byte, out any) (int, error) {
resp, err := c.request(ctx, method, u, apiKey, body)
if err != nil {
return 0, err
}
defer resp.Body.Close()
b, err := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
if err != nil {
return 0, &Error{Msg: "reading ComfyUI's answer: " + err.Error(), Retryable: true}
}
if resp.StatusCode >= 500 {
return resp.StatusCode, &Error{Msg: fmt.Sprintf("ComfyUI answered %d: %s", resp.StatusCode, snippet(b)), Retryable: true}
}
if out != nil && len(b) > 0 {
if err := json.Unmarshal(b, out); err != nil && resp.StatusCode == http.StatusOK {
return resp.StatusCode, &Error{Msg: "ComfyUI's answer is not JSON: " + snippet(b)}
}
}
return resp.StatusCode, nil
}
func (c *Client) get(ctx context.Context, u, apiKey string) (Image, error) {
resp, err := c.request(ctx, http.MethodGet, u, apiKey, nil)
if err != nil {
return Image{}, err
}
defer resp.Body.Close()
b, err := io.ReadAll(io.LimitReader(resp.Body, maxFile))
if err != nil {
return Image{}, &Error{Msg: "reading a file from ComfyUI: " + err.Error(), Retryable: true}
}
if resp.StatusCode != http.StatusOK {
return Image{}, &Error{Msg: fmt.Sprintf("ComfyUI answered %d for a file it made", resp.StatusCode)}
}
return Image{Data: b, ContentType: resp.Header.Get("Content-Type")}, nil
}
// generic is a type that says nothing of what a file is.
func generic(ct string) bool {
return ct == "" || strings.HasPrefix(ct, "application/octet-stream") || strings.HasPrefix(ct, "text/plain")
}
// withoutPNGText is a PNG without its text chunks. ComfyUI writes the whole
// workflow into every picture it saves — the prompt, and the checkpoint and
// nodes behind it — and what a picture says about how it was made is the
// operator's to publish, not something to hand every client. A file that is
// not a well-formed PNG is returned as it came.
func withoutPNGText(b []byte) []byte {
const sig = "\x89PNG\r\n\x1a\n"
if len(b) < len(sig) || string(b[:len(sig)]) != sig {
return b
}
out := make([]byte, 0, len(b))
out = append(out, b[:len(sig)]...)
for i := len(sig); i < len(b); {
if i+12 > len(b) {
return b
}
n := int(binary.BigEndian.Uint32(b[i:]))
end := i + 12 + n
if n < 0 || end > len(b) {
return b
}
switch string(b[i+4 : i+8]) {
case "tEXt", "iTXt", "zTXt":
default:
out = append(out, b[i:end]...)
}
i = end
}
return out
}
// refusal is what /prompt said when it would not queue a workflow.
func refusal(e, nodes json.RawMessage) string {
var top struct {
Message string `json:"message"`
Details string `json:"details"`
}
msg := snippet(e)
if json.Unmarshal(e, &top) == nil && top.Message != "" {
msg = top.Message
if top.Details != "" {
msg += ": " + top.Details
}
}
var perNode map[string]struct {
ClassType string `json:"class_type"`
Errors []struct {
Message string `json:"message"`
Details string `json:"details"`
} `json:"errors"`
}
if json.Unmarshal(nodes, &perNode) == nil {
var parts []string
for _, n := range perNode {
for _, er := range n.Errors {
parts = append(parts, strings.TrimSpace(n.ClassType+": "+er.Message+" "+er.Details))
}
}
slices.Sort(parts)
if len(parts) > 0 {
msg += " (" + strings.Join(parts, "; ") + ")"
}
}
if msg == "" {
msg = "no reason given"
}
return msg
}
func snippet(b []byte) string {
s := strings.TrimSpace(string(b))
if len(s) > 300 {
s = s[:300] + "…"
}
return s
}
// maxFile bounds one file read from ComfyUI: a mesh with its textures runs
// to tens of megabytes, a picture to a few.
const maxFile = 256 << 20
// typeByName is a file's type by its extension, for those a server sends
// as octet-stream: the meshes a 3D node saves.
func typeByName(name string) string {
switch strings.ToLower(path.Ext(name)) {
case ".glb":
return "model/gltf-binary"
case ".gltf":
return "model/gltf+json"
case ".obj":
return "model/obj"
case ".stl":
return "model/stl"
case ".ply":
return "application/ply"
case ".png":
return "image/png"
case ".jpg", ".jpeg":
return "image/jpeg"
case ".webp":
return "image/webp"
case ".gif":
return "image/gif"
case ".mp4":
return "video/mp4"
case ".webm":
return "video/webm"
case ".mov":
return "video/quicktime"
case ".wav":
return "audio/wav"
case ".mp3":
return "audio/mpeg"
case ".flac":
return "audio/flac"
case ".ogg", ".opus":
return "audio/ogg"
}
return ""
}
// Upload puts a file in ComfyUI's input directory, where a loader node reads
// it — LoadImage, LoadAudio, LoadVideo, whatever the workflow uses: ComfyUI's
// /upload/image takes any file — and returns the name it was given there.
func (c *Client) Upload(ctx context.Context, base, apiKey, name string, data []byte) (string, error) {
base = strings.TrimRight(base, "/")
var body bytes.Buffer
mw := multipart.NewWriter(&body)
fw, err := mw.CreateFormFile("image", name)
if err != nil {
return "", err
}
if _, err := fw.Write(data); err != nil {
return "", err
}
_ = mw.WriteField("type", "input")
_ = mw.WriteField("overwrite", "true")
if err := mw.Close(); err != nil {
return "", err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/upload/image", &body)
if err != nil {
return "", err
}
req.Header.Set("Content-Type", mw.FormDataContentType())
if apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := c.client().Do(req)
if err != nil {
if ctx.Err() != nil {
return "", ctx.Err()
}
return "", &Error{Msg: "ComfyUI could not be reached: " + err.Error(), Retryable: true}
}
defer resp.Body.Close()
b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
if resp.StatusCode >= 500 {
return "", &Error{Msg: fmt.Sprintf("ComfyUI answered %d to the picture: %s", resp.StatusCode, snippet(b)), Retryable: true}
}
var out struct {
Name string `json:"name"`
Subfolder string `json:"subfolder"`
}
if resp.StatusCode != http.StatusOK || json.Unmarshal(b, &out) != nil || out.Name == "" {
return "", &Error{Msg: fmt.Sprintf("ComfyUI did not take the picture (%d): %s", resp.StatusCode, snippet(b))}
}
if out.Subfolder != "" {
return out.Subfolder + "/" + out.Name, nil
}
return out.Name, nil
}
// progress is a workflow's: its place in the queue, from ComfyUI, and the
// node running and its step, from the gate when there is one — a ComfyUI
// outside the cluster has no gate, and says only where it is queued.
func (c *Client) progress(ctx context.Context, base, apiKey, id string, workflow map[string]any) (Progress, bool) {
p := Progress{PromptID: id, State: "queued"}
var q struct {
Running [][]json.RawMessage `json:"queue_running"`
Pending [][]json.RawMessage `json:"queue_pending"`
}
if _, err := c.do(ctx, http.MethodGet, base+"/queue", apiKey, nil, &q); err != nil {
return p, false
}
quoted := strconv.Quote(id)
queued := false
for _, r := range q.Running {
if len(r) > 1 && string(r[1]) == quoted {
p.State = "running"
}
}
if p.State != "running" {
// Pending in the order ComfyUI will take them: by number.
mine := -1.0
var numbers []float64
for _, r := range q.Pending {
if len(r) < 2 {
continue
}
var n float64
_ = json.Unmarshal(r[0], &n)
if string(r[1]) == quoted {
mine, queued = n, true
continue
}
numbers = append(numbers, n)
}
if queued {
p.Ahead = len(q.Running)
for _, n := range numbers {
if n < mine {
p.Ahead++
}
}
return p, true
}
}
var g Progress
if st, err := c.do(ctx, http.MethodGet, base+"/progress/"+url.PathEscape(id), apiKey, nil, &g); err == nil &&
st == http.StatusOK && g.State != "" {
p = g
p.PromptID = id
if node, ok := workflow[p.Node].(map[string]any); ok {
p.NodeClass, _ = node["class_type"].(string)
}
}
return p, true
}