From 2f06c0e51354a82b3b2f5e3161bb66b920453118 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Fri, 9 Oct 2026 21:15:50 +0200 Subject: [PATCH] Pictures from ComfyUI at /v1/images/generations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An external model with api: comfyui is a ComfyUI server. A request for pictures goes through the same key, allowlist, limit, admission and guard checks as any other; at the replica it becomes a workflow — a built-in (flux-schnell, flux-dev) or a JSON file in API format, with {{prompt}}, {{width}}, {{seed}} and the like filled in — queued, waited on, and its pictures sent back as b64_json. A client that leaves takes its workflow off the queue or interrupts it. ComfyUI is never scraped for metrics. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/admin-api.openapi.json | 12 + docs/reference.md | 36 ++ internal/comfy/comfy.go | 494 +++++++++++++++++++++ internal/comfy/comfy_test.go | 198 +++++++++ internal/comfy/workflows/flux-dev.json | 12 + internal/comfy/workflows/flux-schnell.json | 11 + internal/config/config.go | 5 +- internal/control/catalog/catalog.go | 39 ++ internal/control/control.go | 23 +- internal/docgen/reference_gen.go | 4 + internal/router/endpoint/endpoint.go | 13 +- internal/router/handlers.go | 12 +- internal/router/images.go | 200 +++++++++ internal/router/images_test.go | 153 +++++++ internal/router/router.go | 15 +- internal/router/scrape/poller.go | 4 +- 16 files changed, 1221 insertions(+), 10 deletions(-) create mode 100644 internal/comfy/comfy.go create mode 100644 internal/comfy/comfy_test.go create mode 100644 internal/comfy/workflows/flux-dev.json create mode 100644 internal/comfy/workflows/flux-schnell.json create mode 100644 internal/router/images.go create mode 100644 internal/router/images_test.go diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index a0a1b6e..44e689f 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -8565,6 +8565,10 @@ }, "embeddings": { "type": "boolean" + }, + "images": { + "type": "boolean", + "description": "Images is a model that makes pictures, asked for at /v1/images/generations. Implied by an external model whose API is ComfyUI's." } } }, @@ -8605,6 +8609,14 @@ "peer": { "type": "string", "description": "Peer names the federated llmgw cluster this model is offered by, for a model that federation added (`llmgw ctl federation add`) rather than one written by hand. Such a model is named PEER/MODEL, comes and goes with the peer's own list, and is sent only while the peer is up." + }, + "api": { + "type": "string", + "description": "API is what the provider speaks: OpenAI's, unless this says comfyui — a ComfyUI server, asked for pictures at /v1/images/generations and sent the workflow below with the request's prompt in it. Model is then the checkpoint file the workflow loads, as ComfyUI lists it." + }, + "workflow": { + "type": "string", + "description": "Workflow is the ComfyUI workflow, in its API format: a built-in's name (flux-schnell, the default, or flux-dev) or a JSON file on the control plane." } } }, diff --git a/docs/reference.md b/docs/reference.md index fb1384e..e8fa916 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -60,6 +60,7 @@ rather than repeated. | `tools` | `bool` | [Capabilities](#capabilities) | | `vision` | `bool` | [Capabilities](#capabilities) | | `embeddings` | `bool` | [Capabilities](#capabilities) | +| `images` | `bool` | [Capabilities](#capabilities) | | `interval` | `time.Duration` | [Control](#control) | | `base_backoff` | `time.Duration` | [Control](#control) | | `max_backoff` | `time.Duration` | [Control](#control) | @@ -76,6 +77,7 @@ rather than repeated. | `api_key` | `string` | [Endpoint](#endpoint) | | `models` | `[]string` | [Endpoint](#endpoint) | | `engine` | `string` | [Endpoint](#endpoint) | +| `workflow` | `string` | [Endpoint](#endpoint) | | `pool` | `string` | [Endpoint](#endpoint) | | `weight` | `float64` | [Endpoint](#endpoint) | | `priority` | `int` | [Endpoint](#endpoint) | @@ -90,6 +92,8 @@ rather than repeated. | `model` | `string` | [External](#external) | | `api_key_file` | `string` | [External](#external) | | `peer` | `string` | [External](#external) | +| `api` | `string` | [External](#external) | +| `workflow` | `string` | [External](#external) | | `budgets` | `[]FacilityBudget` | [Facility](#facility) | | `margin` | `float64` | [Facility](#facility) | | `gpu_watts` | `float64` | [Facility](#facility) | @@ -716,6 +720,14 @@ token is MaxTTFTRatio times this model's. Zero is 50, 0.05 and 1.5. `bool` +### `images` + +`bool` + +Images is a model that makes pictures, asked for at +/v1/images/generations. Implied by an external model whose API is +ComfyUI's. + ## Control ### `interval` @@ -829,6 +841,13 @@ with one. Engine names the inference engine, which decides what its metrics are called. Empty means vLLM's names. +### `workflow` + +`string` + +Workflow is the ComfyUI workflow for engine: comfyui, a built-in's +name or its JSON; flux-schnell unless set. + ### `pool` `string` @@ -920,6 +939,23 @@ a model that federation added (`llmgw ctl federation add`) rather than one written by hand. Such a model is named PEER/MODEL, comes and goes with the peer's own list, and is sent only while the peer is up. +### `api` + +`string` + +API is what the provider speaks: OpenAI's, unless this says comfyui — +a ComfyUI server, asked for pictures at /v1/images/generations and +sent the workflow below with the request's prompt in it. Model is then +the checkpoint file the workflow loads, as ComfyUI lists it. + +### `workflow` + +`string` + +Workflow is the ComfyUI workflow, in its API format: a built-in's name +(flux-schnell, the default, or flux-dev) or a JSON file on the control +plane. + ## Facility ### `budgets` diff --git a/internal/comfy/comfy.go b/internal/comfy/comfy.go new file mode 100644 index 0000000..c2fc066 --- /dev/null +++ b/internal/comfy/comfy.go @@ -0,0 +1,494 @@ +// 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/hex" + "encoding/json" + "errors" + "fmt" + "io" + "math/big" + "net/http" + "net/url" + "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 +} + +func (p Params) value(name string) (any, bool) { + 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 + } + return nil, false +} + +// 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 + } + 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 picture made. +type Image struct { + Data []byte + ContentType 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) { + 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() + for { + select { + case <-ctx.Done(): + c.cancel(context.WithoutCancel(ctx), base, apiKey, id) + return nil, ctx.Err() + case <-t.C: + } + 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 map[string]struct { + Images []struct { + Filename string `json:"filename"` + Subfolder string `json:"subfolder"` + Type string `json:"type"` + } `json:"images"` + } `json:"outputs"` +} + +// 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 { + for _, im := range h.Outputs[n].Images { + q := url.Values{"filename": {im.Filename}, "subfolder": {im.Subfolder}, "type": {im.Type}} + img, err := c.get(ctx, base+"/view?"+q.Encode(), apiKey) + if err != nil { + return nil, err + } + out = append(out, img) + } + } + if len(out) == 0 { + return nil, &Error{Msg: "the workflow finished without a picture: it has no SaveImage or PreviewImage node"} + } + 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, 64<<20)) + if err != nil { + return Image{}, &Error{Msg: "reading a picture from ComfyUI: " + err.Error(), Retryable: true} + } + if resp.StatusCode != http.StatusOK { + return Image{}, &Error{Msg: fmt.Sprintf("ComfyUI answered %d for a picture it made", resp.StatusCode)} + } + ct := resp.Header.Get("Content-Type") + if ct == "" || ct == "application/octet-stream" { + ct = http.DetectContentType(b) + } + return Image{Data: b, ContentType: ct}, nil +} + +// 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 +} diff --git a/internal/comfy/comfy_test.go b/internal/comfy/comfy_test.go new file mode 100644 index 0000000..bbf8fe0 --- /dev/null +++ b/internal/comfy/comfy_test.go @@ -0,0 +1,198 @@ +package comfy + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" +) + +func TestTheBuiltInWorkflowsFillIn(t *testing.T) { + for _, name := range Builtins() { + tpl, err := Load(name) + if err != nil { + t.Fatalf("%s: %v", name, err) + } + g, err := tpl.Fill(Params{Checkpoint: "flux.safetensors", Prompt: `a "red" fox`, Width: 768, Height: 512, N: 2, Seed: 42}) + if err != nil { + t.Fatalf("%s: %v", name, err) + } + b, _ := json.Marshal(g) + s := string(b) + for _, want := range []string{`"ckpt_name":"flux.safetensors"`, `"text":"a \"red\" fox"`, `"width":768`, + `"height":512`, `"batch_size":2`, `"seed":42`} { + if !strings.Contains(s, want) { + t.Errorf("%s: %s not in %s", name, want, s) + } + } + if strings.Contains(s, "{{") { + t.Errorf("%s: a placeholder was left: %s", name, s) + } + } +} + +func TestAStepCountGivenReplacesTheDefault(t *testing.T) { + tpl, _ := Load("flux-schnell") + g, _ := tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512}) + if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 4.0 { + t.Errorf("default steps = %v, want 4", steps) + } + g, _ = tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512, Steps: 8}) + if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 8 { + t.Errorf("steps = %v, want 8", steps) + } +} + +func TestAPlaceholderWithoutAValueOrADefaultIsSaid(t *testing.T) { + tpl, err := Load(`{"1":{"class_type":"X","inputs":{"a":"{{checkpoint}}","b":"size {{width}}x{{height}}"}}}`) + if err != nil { + t.Fatal(err) + } + if _, err := tpl.Fill(Params{Width: 5}); err == nil || !strings.Contains(err.Error(), "checkpoint, height") { + t.Fatalf("err = %v", err) + } + g, err := tpl.Fill(Params{Checkpoint: "c", Width: 5, Height: 6}) + if err != nil { + t.Fatal(err) + } + if b := g["1"].(map[string]any)["inputs"].(map[string]any)["b"]; b != "size 5x6" { + t.Errorf("b = %v", b) + } +} + +func TestAWorkflowInTheUIsFormatIsRefused(t *testing.T) { + if _, err := Load(`{"nodes":[{"id":1}],"links":[]}`); err == nil || !strings.Contains(err.Error(), "API format") { + t.Fatalf("err = %v", err) + } + if _, err := Load("sdxl-turbo"); err == nil { + t.Fatal("an unknown built-in was loaded") + } +} + +// fakeComfy is just enough of ComfyUI: /prompt queues, /history answers +// after a few asks, /view serves a picture. +type fakeComfy struct { + mu sync.Mutex + asked int + fail string + refuse bool + deleted []string + queued map[string]any + readyAt int + srv *httptest.Server + interrupted bool +} + +func newFakeComfy(t *testing.T) *fakeComfy { + f := &fakeComfy{readyAt: 2} + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + f.mu.Lock() + defer f.mu.Unlock() + switch { + case r.Method == http.MethodPost && r.URL.Path == "/prompt": + if f.refuse { + w.WriteHeader(http.StatusBadRequest) + fmt.Fprint(w, `{"error":{"message":"Prompt outputs failed validation","details":""},"node_errors":{"1":{"class_type":"CheckpointLoaderSimple","errors":[{"message":"Value not in list","details":"ckpt_name: 'x' not in []"}]}}}`) + return + } + var b struct { + Prompt map[string]any `json:"prompt"` + } + _ = json.NewDecoder(r.Body).Decode(&b) + f.queued = b.Prompt + fmt.Fprint(w, `{"prompt_id":"p1","number":1,"node_errors":{}}`) + case r.URL.Path == "/history/p1": + f.asked++ + if f.asked < f.readyAt { + fmt.Fprint(w, `{}`) + return + } + if f.fail != "" { + fmt.Fprintf(w, `{"p1":{"status":{"status_str":"error","completed":false,"messages":[["execution_start",{}],["execution_error",{"node_type":"KSampler","exception_message":%q}]]},"outputs":{}}}`, f.fail) + return + } + fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true,"messages":[]},"outputs":{"7":{"images":[{"filename":"a.png","subfolder":"","type":"temp"},{"filename":"b.png","subfolder":"","type":"temp"}]}}}}`) + case r.URL.Path == "/view": + w.Header().Set("Content-Type", "image/png") + fmt.Fprint(w, "PNG:"+r.URL.Query().Get("filename")) + case r.URL.Path == "/queue" && r.Method == http.MethodPost: + var b struct { + Delete []string `json:"delete"` + } + _ = json.NewDecoder(r.Body).Decode(&b) + f.deleted = append(f.deleted, b.Delete...) + case r.URL.Path == "/queue": + fmt.Fprint(w, `{"queue_running":[[1,"p1",{},{},[]]],"queue_pending":[]}`) + case r.URL.Path == "/interrupt": + f.interrupted = true + default: + http.NotFound(w, r) + } + })) + t.Cleanup(f.srv.Close) + return f +} + +func TestGenerateWaitsForTheWorkflowAndFetchesItsPictures(t *testing.T) { + f := newFakeComfy(t) + c := &Client{Poll: time.Millisecond} + images, err := c.Generate(t.Context(), f.srv.URL, "", map[string]any{"1": map[string]any{"class_type": "X"}}) + if err != nil { + t.Fatal(err) + } + if len(images) != 2 || string(images[0].Data) != "PNG:a.png" || images[1].ContentType != "image/png" { + t.Fatalf("images = %+v", images) + } + if f.queued["1"] == nil { + t.Error("the workflow was not what was queued") + } +} + +func TestAWorkflowThatFailsSaysWhy(t *testing.T) { + f := newFakeComfy(t) + f.fail = "CUDA out of memory" + _, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{}) + if err == nil || !strings.Contains(err.Error(), "KSampler: CUDA out of memory") { + t.Fatalf("err = %v", err) + } +} + +func TestAWorkflowComfyUIRefusesIsNotRetried(t *testing.T) { + f := newFakeComfy(t) + f.refuse = true + _, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{}) + ce, ok := err.(*Error) + if !ok || ce.Retryable || !strings.Contains(ce.Msg, "Value not in list") { + t.Fatalf("err = %#v", err) + } +} + +func TestAComfyUIThatCannotBeReachedIsRetryable(t *testing.T) { + f := newFakeComfy(t) + f.srv.Close() + _, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{}) + if ce, ok := err.(*Error); !ok || !ce.Retryable { + t.Fatalf("err = %#v", err) + } +} + +func TestAClientThatLeavesTakesItsWorkflowOffTheCard(t *testing.T) { + f := newFakeComfy(t) + f.readyAt = 1 << 30 + ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond) + defer cancel() + _, err := (&Client{Poll: time.Millisecond}).Generate(ctx, f.srv.URL, "", map[string]any{}) + if err == nil { + t.Fatal("no error") + } + f.mu.Lock() + defer f.mu.Unlock() + if len(f.deleted) != 1 || f.deleted[0] != "p1" || !f.interrupted { + t.Errorf("deleted %v, interrupted %v", f.deleted, f.interrupted) + } +} diff --git a/internal/comfy/workflows/flux-dev.json b/internal/comfy/workflows/flux-dev.json new file mode 100644 index 0000000..b231573 --- /dev/null +++ b/internal/comfy/workflows/flux-dev.json @@ -0,0 +1,12 @@ +{ + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "{{checkpoint}}"}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{prompt}}", "clip": ["1", 1]}}, + "3": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{negative_prompt}}", "clip": ["1", 1]}}, + "8": {"class_type": "FluxGuidance", "inputs": {"guidance": "{{guidance|3.5}}", "conditioning": ["2", 0]}}, + "4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": "{{width}}", "height": "{{height}}", "batch_size": "{{n}}"}}, + "5": {"class_type": "KSampler", "inputs": { + "seed": "{{seed}}", "steps": "{{steps|20}}", "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, + "model": ["1", 0], "positive": ["8", 0], "negative": ["3", 0], "latent_image": ["4", 0]}}, + "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}}, + "7": {"class_type": "PreviewImage", "inputs": {"images": ["6", 0]}} +} diff --git a/internal/comfy/workflows/flux-schnell.json b/internal/comfy/workflows/flux-schnell.json new file mode 100644 index 0000000..387b94f --- /dev/null +++ b/internal/comfy/workflows/flux-schnell.json @@ -0,0 +1,11 @@ +{ + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "{{checkpoint}}"}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{prompt}}", "clip": ["1", 1]}}, + "3": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{negative_prompt}}", "clip": ["1", 1]}}, + "4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": "{{width}}", "height": "{{height}}", "batch_size": "{{n}}"}}, + "5": {"class_type": "KSampler", "inputs": { + "seed": "{{seed}}", "steps": "{{steps|4}}", "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, + "model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0]}}, + "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}}, + "7": {"class_type": "PreviewImage", "inputs": {"images": ["6", 0]}} +} diff --git a/internal/config/config.go b/internal/config/config.go index 0bb6d8f..582c25c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -723,7 +723,10 @@ type Endpoint struct { Models []string `koanf:"models"` // Engine names the inference engine, which decides what its metrics are // called. Empty means vLLM's names. - Engine string `koanf:"engine"` + Engine string `koanf:"engine"` + // Workflow is the ComfyUI workflow for engine: comfyui, a built-in's + // name or its JSON; flux-schnell unless set. + Workflow string `koanf:"workflow"` Pool string `koanf:"pool"` Weight float64 `koanf:"weight"` Priority int `koanf:"priority"` diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 06371b2..0f52fed 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -14,6 +14,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" "gitea.mynodes.xyz/ilker/llmgw/internal/control/facility" "gitea.mynodes.xyz/ilker/llmgw/internal/control/placement" "gitea.mynodes.xyz/ilker/llmgw/internal/guard" @@ -457,6 +458,15 @@ type Capabilities struct { Tools bool `json:"tools,omitempty" koanf:"tools"` Vision bool `json:"vision,omitempty" koanf:"vision"` Embeddings bool `json:"embeddings,omitempty" koanf:"embeddings"` + // Images is a model that makes pictures, asked for at + // /v1/images/generations. Implied by an external model whose API is + // ComfyUI's. + Images bool `json:"images,omitempty" koanf:"images"` +} + +// MakesImages reports whether a model answers /v1/images/generations. +func (m Model) MakesImages() bool { + return m.Capabilities.Images || m.External != nil && m.External.API == APIComfyUI } // Catalog is the set of models the cluster should serve. @@ -1101,8 +1111,20 @@ type External struct { // than one written by hand. Such a model is named PEER/MODEL, comes and // goes with the peer's own list, and is sent only while the peer is up. Peer string `json:"peer,omitempty" koanf:"peer"` + // API is what the provider speaks: OpenAI's, unless this says comfyui — + // a ComfyUI server, asked for pictures at /v1/images/generations and + // sent the workflow below with the request's prompt in it. Model is then + // the checkpoint file the workflow loads, as ComfyUI lists it. + API string `json:"api,omitempty" koanf:"api"` + // Workflow is the ComfyUI workflow, in its API format: a built-in's name + // (flux-schnell, the default, or flux-dev) or a JSON file on the control + // plane. + Workflow string `json:"workflow,omitempty" koanf:"workflow"` } +// APIComfyUI is External.API for a ComfyUI server. +const APIComfyUI = "comfyui" + // IsPeer reports whether a model is offered by a federated cluster. func (m Model) IsPeer() bool { return m.External != nil && m.External.Peer != "" } @@ -1117,6 +1139,23 @@ func (e *External) validate(model string, profiles int) error { if strings.TrimSpace(e.Model) == "" { return fmt.Errorf("model %q: external model needs the name the provider knows it by", model) } + switch e.API { + case "", "openai": + if e.Workflow != "" { + return fmt.Errorf("model %q: a workflow is for api: comfyui", model) + } + case APIComfyUI: + if e.Peer != "" { + return fmt.Errorf("model %q: a federated model speaks OpenAI's API, not ComfyUI's", model) + } + // A file is read where the router runs; a name is checked here. + if w := e.Workflow; w != "" && !strings.ContainsAny(w, "/.") && !slices.Contains(comfy.Builtins(), w) { + return fmt.Errorf("model %q: no built-in workflow %q (there are %s; or name a JSON file)", + model, w, strings.Join(comfy.Builtins(), ", ")) + } + default: + return fmt.Errorf("model %q: external api %q is not openai or comfyui", model, e.API) + } return nil } diff --git a/internal/control/control.go b/internal/control/control.go index 7630919..2d8260c 100644 --- a/internal/control/control.go +++ b/internal/control/control.go @@ -9,6 +9,7 @@ import ( "encoding/hex" "fmt" "log/slog" + "os" "strings" "sync" "time" @@ -412,12 +413,27 @@ func (c *Controller) externalEndpoints() []endpoint.Info { if m.IsPeer() { id = "peer:" + m.Name } + engine, workflow := endpoint.EngineExternal, "" + if m.External.API == catalog.APIComfyUI { + engine, workflow = endpoint.EngineComfyUI, m.External.Workflow + // A file is read here, on the control plane the catalog says it + // is on, and sent as the workflow itself. + if strings.ContainsAny(workflow, "/.") { + b, err := os.ReadFile(workflow) + if err != nil { + c.log.Warn("a ComfyUI model's workflow could not be read; not offered", "model", m.Name, "err", err) + continue + } + workflow = string(b) + } + } out = append(out, endpoint.Info{ ID: id, URL: strings.TrimRight(m.External.URL, "/"), Models: append([]string{m.Name}, m.Aliases...), APIKey: key, - Engine: endpoint.EngineExternal, Pool: "external", Weight: 1, + Engine: engine, Pool: "external", Weight: 1, UpstreamModel: m.External.Model, Peer: m.External.Peer, PeerCluster: cluster, PeerTags: tags, + Workflow: workflow, }) } return out @@ -427,8 +443,9 @@ func (c *Controller) externalEndpoints() []endpoint.Info { func fingerprint(infos []endpoint.Info) string { var b strings.Builder for _, i := range infos { - fmt.Fprintf(&b, "%s|%s|%s|%s|%v|%g|%d|%s|%s|%s|%t|%x;", i.ID, i.URL, i.Engine, i.Pool, i.Models, i.Weight, - i.Priority, i.UpstreamModel, i.Peer, i.PeerCluster, i.PeerTags, sha256.Sum256([]byte(i.APIKey))) + fmt.Fprintf(&b, "%s|%s|%s|%s|%v|%g|%d|%s|%s|%s|%t|%x|%x;", i.ID, i.URL, i.Engine, i.Pool, i.Models, i.Weight, + i.Priority, i.UpstreamModel, i.Peer, i.PeerCluster, i.PeerTags, sha256.Sum256([]byte(i.APIKey)), + sha256.Sum256([]byte(i.Workflow))) } return b.String() } diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 19ae5a0..0afa759 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -64,6 +64,7 @@ var Reference = []Block{ {Key: "tools", Type: "bool", Doc: "", Nested: ""}, {Key: "vision", Type: "bool", Doc: "", Nested: ""}, {Key: "embeddings", Type: "bool", Doc: "", Nested: ""}, + {Key: "images", Type: "bool", Doc: "Images is a model that makes pictures, asked for at\n/v1/images/generations. Implied by an external model whose API is\nComfyUI's.", Nested: ""}, }}, {Package: "config", Name: "Control", Title: "Control", Doc: "", Fields: []Field{ {Key: "interval", Type: "time.Duration", Doc: "Interval is how often desired and actual state are compared.", Nested: ""}, @@ -84,6 +85,7 @@ var Reference = []Block{ {Key: "api_key", Type: "string", Doc: "APIKey is sent to the engine as a bearer token, for an engine started\nwith one.", Nested: ""}, {Key: "models", Type: "[]string", Doc: "", Nested: ""}, {Key: "engine", Type: "string", Doc: "Engine names the inference engine, which decides what its metrics are\ncalled. Empty means vLLM's names.", Nested: ""}, + {Key: "workflow", Type: "string", Doc: "Workflow is the ComfyUI workflow for engine: comfyui, a built-in's\nname or its JSON; flux-schnell unless set.", Nested: ""}, {Key: "pool", Type: "string", Doc: "", Nested: ""}, {Key: "weight", Type: "float64", Doc: "", Nested: ""}, {Key: "priority", Type: "int", Doc: "", Nested: ""}, @@ -104,6 +106,8 @@ var Reference = []Block{ {Key: "model", Type: "string", Doc: "Model is what the provider calls it; requests are rewritten to it.", Nested: ""}, {Key: "api_key_file", Type: "string", Doc: "APIKeyFile holds the provider's key, on the control plane. A file\nrather than the key itself, so the key is in neither the catalog, the\nstore nor the audit log; empty for a provider that asks for none.", Nested: ""}, {Key: "peer", Type: "string", Doc: "Peer names the federated llmgw cluster this model is offered by, for\na model that federation added (`llmgw ctl federation add`) rather\nthan one written by hand. Such a model is named PEER/MODEL, comes and\ngoes with the peer's own list, and is sent only while the peer is up.", Nested: ""}, + {Key: "api", Type: "string", Doc: "API is what the provider speaks: OpenAI's, unless this says comfyui —\na ComfyUI server, asked for pictures at /v1/images/generations and\nsent the workflow below with the request's prompt in it. Model is then\nthe checkpoint file the workflow loads, as ComfyUI lists it.", Nested: ""}, + {Key: "workflow", Type: "string", Doc: "Workflow is the ComfyUI workflow, in its API format: a built-in's name\n(flux-schnell, the default, or flux-dev) or a JSON file on the control\nplane.", Nested: ""}, }}, {Package: "config", Name: "Facility", Title: "Facility", Doc: "", Fields: []Field{ {Key: "budgets", Type: "[]FacilityBudget", Doc: "Budgets are the domains with a power cap.", Nested: ""}, diff --git a/internal/router/endpoint/endpoint.go b/internal/router/endpoint/endpoint.go index bfdb397..049322d 100644 --- a/internal/router/endpoint/endpoint.go +++ b/internal/router/endpoint/endpoint.go @@ -48,12 +48,23 @@ type Info struct { Peer string PeerCluster string PeerTags bool + // Workflow is the ComfyUI workflow sent for an endpoint whose engine is + // EngineComfyUI: a built-in's name or the workflow's JSON. + Workflow string } // EngineExternal marks an endpoint that is a provider outside the cluster. // It publishes no metrics, so it is never scraped and never reads as stale. const EngineExternal = "external" +// EngineComfyUI marks a ComfyUI server. The router speaks ComfyUI's API to +// it, not OpenAI's, and it publishes no metrics either. +const EngineComfyUI = "comfyui" + +// Scraped reports whether an engine publishes metrics the router reads. One +// that does not is never read as stale for having none. +func Scraped(engine string) bool { return engine != EngineExternal && engine != EngineComfyUI } + // Serves reports whether the replica serves the given model. func (i Info) Serves(model string) bool { return slices.Contains(i.Models, model) @@ -257,7 +268,7 @@ func (r *Registry) Snapshots(model string) []Snapshot { continue } var age time.Duration - stale := l.info.Engine != EngineExternal + stale := Scraped(l.info.Engine) if !l.scrapedAt.IsZero() { age = now.Sub(l.scrapedAt) stale = age > r.cfg.StaleAfter diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 3b8c97a..c112574 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -102,6 +102,8 @@ func (r *Router) Handler() http.Handler { mux.HandleFunc("POST /v1/score", r.handleInference("/v1/score")) mux.HandleFunc("POST /v1/audio/transcriptions", r.handleInference("/v1/audio/transcriptions")) mux.HandleFunc("POST /v1/audio/translations", r.handleInference("/v1/audio/translations")) + // Pictures, made by ComfyUI from a workflow the catalog names. + mux.HandleFunc("POST "+imagesPath, r.handleInference(imagesPath)) // Answered by the guardrails' judge model. mux.HandleFunc("POST /v1/moderations", r.handleModerations) // Codex, and whatever else speaks only the Responses API, translated to @@ -477,6 +479,14 @@ func (r *Router) handleInference(path string) http.HandlerFunc { return } + // A picture asked of a model that writes text, or the other way + // round, is refused before it waits for either. + if why := r.imageRouteRefusal(path, parsed.Model, parsed.UpstreamBody); why != "" { + writeError(w, http.StatusBadRequest, "invalid_request_error", why) + r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0) + return + } + // Before authentication work and before admission: a request that // cannot fit should not consume a quota reservation or a queue slot on // its way to being refused. @@ -1121,7 +1131,7 @@ func (r *Router) handleModels(w http.ResponseWriter, req *http.Request) { } entry["capabilities"] = map[string]bool{ "tools": m.Capabilities.Tools, "vision": m.Capabilities.Vision, - "embeddings": m.Capabilities.Embeddings, + "embeddings": m.Capabilities.Embeddings, "images": m.MakesImages(), } data = append(data, entry) } diff --git a/internal/router/images.go b/internal/router/images.go new file mode 100644 index 0000000..940c751 --- /dev/null +++ b/internal/router/images.go @@ -0,0 +1,200 @@ +package router + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" +) + +// /v1/images/generations, OpenAI's shape, made by ComfyUI: the request goes +// through the same checks as any other — key, allowlist, limits, admission, +// guardrails on its prompt — and only at the replica is it turned into a +// workflow, queued and waited on, and its pictures sent back as base64. + +const imagesPath = "/v1/images/generations" + +// imageRequest is what a request for pictures asks. Beyond OpenAI's fields, +// those a diffusion model has and DALL·E does not: seed, steps, +// negative_prompt and guidance. +type imageRequest struct { + Prompt string `json:"prompt"` + N int `json:"n"` + Size string `json:"size"` + ResponseFormat string `json:"response_format"` + Seed *int64 `json:"seed"` + Steps int `json:"steps"` + NegativePrompt string `json:"negative_prompt"` + Guidance *float64 `json:"guidance"` + + width, height int +} + +// The bounds of a request. A FLUX latent is 1/16 of the picture, so a side +// that is not a multiple of 16 is not one it can make; the rest keep one +// request from holding a card for the length of a film. +const ( + maxImages = 4 + minSide = 256 + maxSide = 2048 + maxSteps = 100 + defaultSide = 1024 +) + +func parseImageRequest(body []byte) (imageRequest, error) { + var in imageRequest + if err := json.Unmarshal(body, &in); err != nil { + return in, err + } + if strings.TrimSpace(in.Prompt) == "" { + return in, errors.New("prompt is required") + } + switch { + case in.N == 0: + in.N = 1 + case in.N < 0 || in.N > maxImages: + return in, fmt.Errorf("n is between 1 and %d", maxImages) + } + switch in.ResponseFormat { + case "", "b64_json": + case "url": + // A URL would have to be served by somebody and kept for some time, + // and a picture made from a prompt is the prompt's owner's: it is + // not kept here. + return in, errors.New(`response_format "url" is not offered: the picture comes back as b64_json and is kept nowhere`) + default: + return in, fmt.Errorf("response_format %q is not b64_json", in.ResponseFormat) + } + in.width, in.height = defaultSide, defaultSide + if s := strings.TrimSpace(in.Size); s != "" && s != "auto" { + w, h, ok := strings.Cut(strings.ToLower(s), "x") + wi, err1 := strconv.Atoi(w) + hi, err2 := strconv.Atoi(h) + if !ok || err1 != nil || err2 != nil { + return in, fmt.Errorf("size %q is not WIDTHxHEIGHT", s) + } + in.width, in.height = wi, hi + } + for _, side := range []int{in.width, in.height} { + if side < minSide || side > maxSide || side%16 != 0 { + return in, fmt.Errorf("size %dx%d: each side is a multiple of 16 from %d to %d", in.width, in.height, minSide, maxSide) + } + } + if in.Steps < 0 || in.Steps > maxSteps { + return in, fmt.Errorf("steps is between 1 and %d", maxSteps) + } + if in.Guidance != nil && (*in.Guidance <= 0 || *in.Guidance > 30) { + return in, errors.New("guidance is above 0 and at most 30") + } + return in, nil +} + +// imageRouteRefusal says why a request cannot be answered at the path it was +// sent to, before it takes a place in a queue: a picture asked of a model +// that writes text, or text of one that makes pictures. Empty when it can, +// and for a model the catalog does not know, which a hand-configured +// endpoint may serve as it likes. +func (r *Router) imageRouteRefusal(path, model string, body []byte) string { + m, known := r.modelMap()[model] + if path == imagesPath { + if known && !m.MakesImages() { + return fmt.Sprintf("%s does not make pictures", m.Name) + } + if _, err := parseImageRequest(body); err != nil { + return err.Error() + } + return "" + } + if known && m.External != nil && m.External.API == catalog.APIComfyUI { + return fmt.Sprintf("%s makes pictures: ask for them at %s", m.Name, imagesPath) + } + return "" +} + +// generateImages answers a request at a ComfyUI replica. It returns what the +// proxy would: a failure of this replica that another may not have is +// retryable, and nothing is written until the pictures are all in hand. +func (r *Router) generateImages(ctx context.Context, w http.ResponseWriter, req *parsedRequest, ep endpoint.Info, body []byte) (proxy.Outcome, error) { + if req.Path != imagesPath { + return proxy.Outcome{}, fmt.Errorf("%s is a ComfyUI server: it answers %s only", ep.ID, imagesPath) + } + in, err := parseImageRequest(body) + if err != nil { + return proxy.Outcome{}, err + } + var upstream struct { + Model string `json:"model"` + } + _ = json.Unmarshal(body, &upstream) + tpl, err := comfy.Load(ep.Workflow) + if err != nil { + return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) + } + seed := comfy.RandomSeed() + if in.Seed != nil { + seed = *in.Seed + } + p := comfy.Params{ + Checkpoint: upstream.Model, Prompt: in.Prompt, NegativePrompt: in.NegativePrompt, + Width: in.width, Height: in.height, N: in.N, Seed: seed, Steps: in.Steps, + } + if in.Guidance != nil { + p.Guidance = *in.Guidance + } + workflow, err := tpl.Fill(p) + if err != nil { + return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) + } + + if t := r.cfg.Timeouts.Total; t > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, t) + defer cancel() + } + start := time.Now() + images, err := r.comfy.Generate(ctx, ep.URL, ep.APIKey, workflow) + if err != nil { + var ce *comfy.Error + switch { + case errors.As(err, &ce): + return proxy.Outcome{Retryable: ce.Retryable}, err + case errors.Is(err, context.Canceled) && ctxCauseIsClient(ctx): + return proxy.Outcome{}, proxy.ErrClientGone + case errors.Is(err, context.DeadlineExceeded): + return proxy.Outcome{}, fmt.Errorf("%s made no picture within %s", ep.ID, r.cfg.Timeouts.Total) + } + return proxy.Outcome{}, err + } + + data := make([]map[string]any, 0, len(images)) + for _, im := range images { + data = append(data, map[string]any{"b64_json": base64.StdEncoding.EncodeToString(im.Data)}) + } + out, err := json.Marshal(map[string]any{"created": time.Now().Unix(), "data": data}) + if err != nil { + return proxy.Outcome{}, err + } + // The seed that made them, so a picture liked can be made again. + w.Header().Set("x-llmgw-seed", strconv.FormatInt(seed, 10)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(out) + took := time.Since(start) + return proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(out)), TTFT: took, Decode: took}, nil +} + +// ctxCauseIsClient reports whether a context ended because the client went +// away, rather than because a deadline set here passed. +func ctxCauseIsClient(ctx context.Context) bool { + return !errors.Is(context.Cause(ctx), context.DeadlineExceeded) +} diff --git a/internal/router/images_test.go b/internal/router/images_test.go new file mode 100644 index 0000000..40f04b7 --- /dev/null +++ b/internal/router/images_test.go @@ -0,0 +1,153 @@ +package router + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// imageRouter has a ComfyUI model, flux, beside a text model, and returns +// what the ComfyUI server was sent. +func imageRouter(t *testing.T) (*Router, *map[string]any, *sync.Mutex) { + t.Helper() + var mu sync.Mutex + var queued map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/metrics": + t.Error("ComfyUI was scraped for metrics") + case "/prompt": + var b struct { + Prompt map[string]any `json:"prompt"` + } + _ = json.NewDecoder(r.Body).Decode(&b) + mu.Lock() + queued = b.Prompt + mu.Unlock() + fmt.Fprint(w, `{"prompt_id":"p1"}`) + case "/history/p1": + fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"7":{"images":[{"filename":"a.png","type":"temp"}]}}}}`) + case "/view": + w.Header().Set("Content-Type", "image/png") + fmt.Fprint(w, "png-bytes") + default: + http.NotFound(w, r) + } + })) + t.Cleanup(srv.Close) + + cfg := baseConfig() + cfg.Models = []catalog.Model{catModel("chat", nil), {Name: "flux", Source: "external", + External: &catalog.External{URL: srv.URL, Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}}} + rt := newTestRouter(t, cfg) + rt.comfy = &comfy.Client{Poll: time.Millisecond} + rt.SetEndpoints(t.Context(), []endpoint.Info{ + {ID: "external:flux", URL: srv.URL, Models: []string{"flux"}, Engine: endpoint.EngineComfyUI, + Pool: "external", Weight: 1, UpstreamModel: "flux1-schnell-fp8.safetensors"}, + }) + return rt, &queued, &mu +} + +func postImages(t *testing.T, h http.Handler, path string, body any) *httptest.ResponseRecorder { + t.Helper() + b, _ := json.Marshal(body) + req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(b)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec +} + +func TestAPictureIsMadeByComfyUIAndComesBackAsBase64(t *testing.T) { + rt, queued, mu := imageRouter(t) + rec := postImages(t, rt.Handler(), imagesPath, map[string]any{ + "model": "flux", "prompt": "a lighthouse at dusk", "size": "768x512", "seed": 7}) + if rec.Code != http.StatusOK { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + var out struct { + Data []struct { + B64 string `json:"b64_json"` + } `json:"data"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || len(out.Data) != 1 { + t.Fatalf("body %s: %v", rec.Body.String(), err) + } + if b, _ := base64.StdEncoding.DecodeString(out.Data[0].B64); string(b) != "png-bytes" { + t.Errorf("picture = %q", b) + } + if rec.Header().Get("x-llmgw-seed") != "7" { + t.Errorf("seed header = %q", rec.Header().Get("x-llmgw-seed")) + } + mu.Lock() + defer mu.Unlock() + wf, _ := json.Marshal(*queued) + for _, want := range []string{`"ckpt_name":"flux1-schnell-fp8.safetensors"`, `"text":"a lighthouse at dusk"`, + `"width":768`, `"height":512`, `"seed":7`} { + if !strings.Contains(string(wf), want) { + t.Errorf("%s not in the workflow sent: %s", want, wf) + } + } +} + +func TestPicturesAndTextAreAskedForAtTheirOwnPaths(t *testing.T) { + rt, _, _ := imageRouter(t) + h := rt.Handler() + for name, c := range map[string]struct { + path string + body map[string]any + want string + }{ + "text of flux": {"/v1/chat/completions", map[string]any{"model": "flux", "messages": []any{}}, "ask for them at " + imagesPath}, + "picture of chat": {imagesPath, map[string]any{"model": "chat", "prompt": "x"}, "does not make pictures"}, + "no prompt": {imagesPath, map[string]any{"model": "flux"}, "prompt is required"}, + "odd size": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "size": "1000x1000"}, "multiple of 16"}, + "url wanted": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "response_format": "url"}, "b64_json"}, + "too many at once": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "n": 9}, "n is between"}, + } { + rec := postImages(t, h, c.path, c.body) + if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), c.want) { + t.Errorf("%s: %d %s", name, rec.Code, rec.Body.String()) + } + } +} + +func TestAComfyUIModelSaysItMakesPictures(t *testing.T) { + rt, _, _ := imageRouter(t) + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + rec := httptest.NewRecorder() + rt.Handler().ServeHTTP(rec, req) + if !strings.Contains(rec.Body.String(), `"images":true`) { + t.Errorf("/v1/models: %s", rec.Body.String()) + } +} + +func TestAComfyUIModelIsCheckedInTheCatalog(t *testing.T) { + ok := catalog.Model{Name: "f", Source: "external", External: &catalog.External{ + URL: "http://10.0.0.147:8188", Model: "flux1-dev.safetensors", API: catalog.APIComfyUI, Workflow: "flux-dev"}} + if _, err := catalog.New([]catalog.Model{ok}); err != nil { + t.Fatalf("refused: %v", err) + } + for name, e := range map[string]catalog.External{ + "unknown api": {URL: "http://a:1", Model: "m", API: "a1111"}, + "unknown workflow": {URL: "http://a:1", Model: "m", API: catalog.APIComfyUI, Workflow: "sdxl"}, + "workflow, openai": {URL: "http://a:1", Model: "m", Workflow: "flux-dev"}, + "comfyui from peer": {URL: "http://a:1", Model: "m", API: catalog.APIComfyUI, Peer: "other"}, + } { + e := e + if _, err := catalog.New([]catalog.Model{{Name: "f", Source: "external", External: &e}}); err == nil { + t.Errorf("%s: accepted", name) + } + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 0b9fba1..0aa4d0c 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -29,6 +29,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/accounting" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" "gitea.mynodes.xyz/ilker/llmgw/internal/auth" + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" "gitea.mynodes.xyz/ilker/llmgw/internal/config" "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" @@ -53,7 +54,9 @@ type Router struct { reg *endpoint.Registry poller *scrape.Poller px *proxy.Proxy - m *metrics + // comfy asks ComfyUI replicas for pictures. + comfy *comfy.Client + m *metrics // adm is nil when admission control is disabled. adm *admission.Controller @@ -311,6 +314,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { kv: kvindex.New(), kvClient: newKVClient(), m: metrics, + comfy: &comfy.Client{HTTP: &http.Client{Transport: &http.Transport{DialContext: (&net.Dialer{Timeout: cfg.Router.Timeouts.Dial}).DialContext}}}, px: proxy.New(proxy.Config{ DialTimeout: cfg.Router.Timeouts.Dial, ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader, @@ -487,6 +491,7 @@ func toInfos(eps []config.Endpoint) []endpoint.Info { Pool: e.Pool, Weight: e.Weight, Priority: e.Priority, + Workflow: e.Workflow, }) } return out @@ -1162,7 +1167,13 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa } // A filter per attempt, for the model this attempt is served by. filter := r.outputFilter(uctx, req, model, &preq) - out, perr := r.px.Do(uctx, w, preq) + var out proxy.Outcome + var perr error + if dec.Endpoint.Engine == endpoint.EngineComfyUI { + out, perr = r.generateImages(uctx, w, req, dec.Endpoint.Info, body) + } else { + out, perr = r.px.Do(uctx, w, preq) + } releaseAdapter() uspan.SetAttributes(attribute.Int("http.response.status_code", out.Status)) endSpan(uspan, perr) diff --git a/internal/router/scrape/poller.go b/internal/router/scrape/poller.go index e0fb2f2..2e1134c 100644 --- a/internal/router/scrape/poller.go +++ b/internal/router/scrape/poller.go @@ -96,8 +96,8 @@ func (p *Poller) Sync(ctx context.Context) { want := make(map[string]endpoint.Info, len(infos)) for _, in := range infos { - if in.Engine == endpoint.EngineExternal { - continue // a provider elsewhere publishes no /metrics + if !endpoint.Scraped(in.Engine) { + continue // a provider elsewhere, or ComfyUI, publishes no /metrics } want[in.ID] = in }