diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 2f9a16fc..33549f4e 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -8599,6 +8599,34 @@ "output": { "type": "string", "description": "Output is what the workflow makes: pictures (the default), asked for at /v1/images/generations, or meshes, asked for as a job at /v1/3d/generations from a picture." + }, + "nodes": { + "type": "array", + "items": { + "$ref": "#/components/schemas/catalog.CustomNode" + }, + "description": "Nodes are custom nodes the workflow needs, put in each replica's own custom_nodes at the commit named." + } + } + }, + "catalog.CustomNode": { + "type": "object", + "description": "CustomNode is a ComfyUI node from a git repository.", + "properties": { + "repo": { + "type": "string", + "description": "Repo is the repository, https only." + }, + "commit": { + "type": "string", + "description": "Commit is the full commit the node is used at: a branch moves, and a node that changed under a model is a model that changed unasked." + }, + "requirements": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Requirements are the PyPI packages the node needs that ComfyUI does not have, by name and version range only. They are installed beside the replica, never into ComfyUI's environment, at the versions that environment already has of anything it shares." } } }, diff --git a/docs/reference.md b/docs/reference.md index 01a44bab..9f31aa0f 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -65,6 +65,7 @@ rather than repeated. | `workflow` | `string` | [ComfyUI](#comfyui) | | `checkpoint` | `string` | [ComfyUI](#comfyui) | | `output` | `string` | [ComfyUI](#comfyui) | +| `nodes` | `[]CustomNode` | [ComfyUI](#comfyui) | | `dir` | `string` | [ComfyUIConfig](#comfyuiconfig) | | `python` | `string` | [ComfyUIConfig](#comfyuiconfig) | | `interval` | `time.Duration` | [Control](#control) | @@ -78,6 +79,9 @@ rather than repeated. | `eval_dir` | `string` | [Control](#control) | | `enabled` | `bool` | [Control](#control) | | `lease_ttl` | `time.Duration` | [Control](#control) | +| `repo` | `string` | [CustomNode](#customnode) | +| `commit` | `string` | [CustomNode](#customnode) | +| `requirements` | `[]string` | [CustomNode](#customnode) | | `id` | `string` | [Endpoint](#endpoint) | | `url` | `string` | [Endpoint](#endpoint) | | `api_key` | `string` | [Endpoint](#endpoint) | @@ -782,6 +786,13 @@ Output is what the workflow makes: pictures (the default), asked for at /v1/images/generations, or meshes, asked for as a job at /v1/3d/generations from a picture. +### `nodes` + +`[]CustomNode` + +Nodes are custom nodes the workflow needs, put in each replica's own +custom_nodes at the commit named. + ## ComfyUIConfig ### `dir` @@ -882,6 +893,30 @@ see the claim. Long enough to survive a slow query or a brief network stall, short enough that a standby takes over before anyone has finished reading the alert. Zero takes the default. +## CustomNode + +### `repo` + +`string` + +Repo is the repository, https only. + +### `commit` + +`string` + +Commit is the full commit the node is used at: a branch moves, and a +node that changed under a model is a model that changed unasked. + +### `requirements` + +`[]string` + +Requirements are the PyPI packages the node needs that ComfyUI does +not have, by name and version range only. They are installed beside +the replica, never into ComfyUI's environment, at the versions that +environment already has of anything it shares. + ## Endpoint ### `id` diff --git a/internal/agent/engines/comfyui.go b/internal/agent/engines/comfyui.go index e1597833..a434e8c1 100644 --- a/internal/agent/engines/comfyui.go +++ b/internal/agent/engines/comfyui.go @@ -124,5 +124,14 @@ func (c *ComfyUI) Commands(s ShardSpec) ([]Command, error) { if s.APIKey != "" { env["LLMGW_ENGINE_KEY"] = s.APIKey } + // A model's custom nodes, which only the catalog sets — a profile's env + // may not name the variable — and so taken before baseEnv drops it. + if nodes := p.Env[comfyNodesEnv]; nodes != "" { + env[comfyNodesEnv] = nodes + } return []Command{{Name: "comfyui", Path: bin, Args: args, Env: materializeEnv(env)}}, nil } + +// comfyNodesEnv is the custom nodes the catalog gives a ComfyUI model, as +// JSON; catalog.NodesEnv and comfy.NodesEnv by the same name. +const comfyNodesEnv = "LLMGW_COMFYUI_NODES" diff --git a/internal/comfy/comfy.go b/internal/comfy/comfy.go index 2e30905f..98e4757d 100644 --- a/internal/comfy/comfy.go +++ b/internal/comfy/comfy.go @@ -247,6 +247,17 @@ type Client struct { // 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) @@ -277,6 +288,8 @@ func (c *Client) Generate(ctx context.Context, base, apiKey string, workflow map } t := time.NewTicker(poll) defer t.Stop() + var last Progress + var asked time.Time for { select { case <-ctx.Done(): @@ -284,6 +297,13 @@ func (c *Client) Generate(ctx context.Context, base, apiKey string, workflow map 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 { @@ -638,3 +658,60 @@ func (c *Client) Upload(ctx context.Context, base, apiKey, name string, data []b } 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 +} diff --git a/internal/comfy/gate.go b/internal/comfy/gate.go index f965e72a..42bb80f1 100644 --- a/internal/comfy/gate.go +++ b/internal/comfy/gate.go @@ -2,7 +2,10 @@ package comfy import ( "context" + "crypto/rand" "crypto/subtle" + "encoding/hex" + "encoding/json" "errors" "fmt" "io" @@ -92,6 +95,18 @@ func RunGate(ctx context.Context, cfg GateConfig) error { return err } } + out := cfg.Output + if out == nil { + out = os.Stdout + } + nodes, err := parseNodes(os.Getenv(NodesEnv)) + if err != nil { + return err + } + site, err := installNodes(ctx, state, cfg.Python, nodes, findUV(), func(s string) { fmt.Fprintln(out, "[gate] "+s) }) + if err != nil { + return err + } models := filepath.Join(state, "models") if err := LinkWeights(cfg.Weights, models); err != nil { return err @@ -112,9 +127,12 @@ func RunGate(ctx context.Context, cfg GateConfig) error { child := exec.Command(cfg.Python, args...) //nolint:gosec // the agent's own configuration child.Dir = cfg.Dir child.Env = os.Environ() - out := cfg.Output - if out == nil { - out = os.Stdout + if site != "" { + pp := site + if old := os.Getenv("PYTHONPATH"); old != "" { + pp += string(os.PathListSeparator) + old + } + child.Env = append(child.Env, "PYTHONPATH="+pp) } // ComfyUI writes to a pipe of the gate's, never to what the gate was // given. That is the agent's, and an agent that restarts and adopts the @@ -141,9 +159,16 @@ func RunGate(ctx context.Context, cfg GateConfig) error { }() upstream := &url.URL{Scheme: "http", Host: net.JoinHostPort("127.0.0.1", strconv.Itoa(inner))} + // The workflows' progress, heard on ComfyUI's websocket under an id of + // the gate's own. + book := newProgressBook() + cid := make([]byte, 8) + _, _ = rand.Read(cid) + clientID := "llmgw-gate-" + hex.EncodeToString(cid) + go book.watch(ctx, upstream, clientID, log) srv := &http.Server{ Addr: cfg.Listen, - Handler: NewGate(upstream, cfg.Key, cfg.Model), + Handler: newGate(upstream, cfg.Key, cfg.Model, book, clientID), ReadHeaderTimeout: 10 * time.Second, } served := make(chan error, 1) @@ -184,6 +209,12 @@ func shutdown(srv *http.Server) { // NewGate is the gate's handler, in front of the ComfyUI at upstream. func NewGate(upstream *url.URL, key, model string) http.Handler { + return newGate(upstream, key, model, newProgressBook(), "") +} + +// newGate is the gate with the book of its workflows' progress, kept under +// clientID; an empty one leaves the workflows' own client ids alone. +func newGate(upstream *url.URL, key, model string, book *progressBook, clientID string) http.Handler { proxy := httputil.NewSingleHostReverseProxy(upstream) base := proxy.Director proxy.Director = func(r *http.Request) { @@ -232,6 +263,20 @@ func NewGate(upstream *url.URL, key, model string) http.Handler { fmt.Fprintf(w, `{"object":"list","data":[{"id":%q,"object":"model","owned_by":"comfyui"}]}`, model) return } + if id, ok := strings.CutPrefix(r.URL.Path, "/progress/"); ok && r.Method == http.MethodGet { + p, found := book.get(id) + if !found { + http.Error(w, "no word of that workflow", http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(p) + return + } + if r.URL.Path == "/prompt" && r.Method == http.MethodPost && clientID != "" { + book.promptThroughGate(proxy, clientID)(w, r) + return + } for _, p := range gatePaths { if r.URL.Path == p || strings.HasSuffix(p, "/") && strings.HasPrefix(r.URL.Path, p) { proxy.ServeHTTP(w, r) diff --git a/internal/comfy/nodes.go b/internal/comfy/nodes.go new file mode 100644 index 00000000..b13f8418 --- /dev/null +++ b/internal/comfy/nodes.go @@ -0,0 +1,138 @@ +package comfy + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path" + "path/filepath" + "regexp" + "strings" + "time" +) + +// Custom nodes, put in a replica's own base directory at the commit the +// catalog names, their packages installed beside the replica rather than +// into ComfyUI's environment: one model's node cannot break another +// model's replicas, or the pictures every replica makes. + +// Node is a custom node as the catalog gives it. +type Node struct { + Repo string `json:"repo"` + Commit string `json:"commit"` + Requirements []string `json:"requirements,omitempty"` +} + +// NodesEnv is where the agent hands the gate a model's custom nodes. +const NodesEnv = "LLMGW_COMFYUI_NODES" + +var nodeDir = regexp.MustCompile(`[^A-Za-z0-9._-]`) + +// installNodes puts each node in state/custom_nodes at its commit and its +// packages in state/site, and returns the PYTHONPATH ComfyUI is started +// with: empty when no node wants a package. Run as the engines' user, in +// the replica's own directory, before ComfyUI starts. +func installNodes(ctx context.Context, state, python string, nodes []Node, uv string, out func(string)) (string, error) { + if len(nodes) == 0 { + return "", nil + } + ctx, cancel := context.WithTimeout(ctx, 20*time.Minute) + defer cancel() + var reqs []string + for _, n := range nodes { + name := nodeDir.ReplaceAllString(strings.TrimSuffix(path.Base(n.Repo), ".git"), "_") + dir := filepath.Join(state, "custom_nodes", name) + if head, err := gitOutput(ctx, dir, "rev-parse", "HEAD"); err == nil && head == n.Commit { + reqs = append(reqs, n.Requirements...) + continue + } + out(fmt.Sprintf("custom node %s at %s", n.Repo, n.Commit[:12])) + _ = os.RemoveAll(dir) + for _, args := range [][]string{ + {"init", "-q", dir}, + {"-C", dir, "remote", "add", "origin", n.Repo}, + {"-C", dir, "fetch", "-q", "--depth", "1", "origin", n.Commit}, + {"-C", dir, "checkout", "-q", "FETCH_HEAD"}, + } { + if b, err := exec.CommandContext(ctx, "git", args...).CombinedOutput(); err != nil { //nolint:gosec // the catalog's repository, checked as https + return "", fmt.Errorf("custom node %s at %s: git %s: %v: %s", n.Repo, n.Commit, args[len(args)-2], err, lastLines(b)) + } + } + reqs = append(reqs, n.Requirements...) + } + if len(reqs) == 0 { + return "", nil + } + if uv == "" { + return "", fmt.Errorf("custom nodes need %s, and no uv is installed to install them with", strings.Join(reqs, ", ")) + } + site := filepath.Join(state, "site") + marker := filepath.Join(site, ".llmgw-requirements") + want := strings.Join(reqs, "\n") + if b, err := os.ReadFile(marker); err == nil && string(b) == want { + return site, nil + } + out("installing " + strings.Join(reqs, " ") + " for the custom nodes") + // What ComfyUI's environment has, as constraints: a package both want is + // installed at the version ComfyUI runs, so the replica's copy and + // ComfyUI's never disagree. + freeze, err := exec.CommandContext(ctx, uv, "pip", "freeze", "--python", python).Output() //nolint:gosec // the agent's own uv + if err != nil { + return "", fmt.Errorf("reading ComfyUI's packages: %w", err) + } + constraints := filepath.Join(state, "constraints.txt") + if err := os.WriteFile(constraints, freeze, 0o644); err != nil { //nolint:gosec // not secret + return "", err + } + _ = os.RemoveAll(site) + args := append([]string{"pip", "install", "--python", python, "--target", site, "-c", constraints, "--quiet"}, reqs...) + cmd := exec.CommandContext(ctx, uv, args...) //nolint:gosec // requirements checked by the catalog as package specs + cmd.Env = append(os.Environ(), "UV_NO_CONFIG=1", "UV_CACHE_DIR="+filepath.Join(state, ".uv-cache")) + if b, err := cmd.CombinedOutput(); err != nil { + return "", fmt.Errorf("installing the custom nodes' packages (%s): %v: %s", strings.Join(reqs, " "), err, lastLines(b)) + } + // The replica's copies come first on ComfyUI's path, but of anything + // both have they are the same version, so which is found is the same. + if err := os.WriteFile(marker, []byte(want), 0o644); err != nil { //nolint:gosec // not secret + return "", err + } + return site, nil +} + +// parseNodes reads the nodes the agent handed over. +func parseNodes(s string) ([]Node, error) { + if strings.TrimSpace(s) == "" { + return nil, nil + } + var nodes []Node + if err := json.Unmarshal([]byte(s), &nodes); err != nil { + return nil, fmt.Errorf("%s: %w", NodesEnv, err) + } + return nodes, nil +} + +// findUV is the uv the agent installed engines with. +func findUV() string { + if m, _ := filepath.Glob("/opt/llmgw/engines/uv/*/*/uv"); len(m) > 0 { + return m[len(m)-1] + } + if p, err := exec.LookPath("uv"); err == nil { + return p + } + return "" +} + +func gitOutput(ctx context.Context, dir string, args ...string) (string, error) { + b, err := exec.CommandContext(ctx, "git", append([]string{"-C", dir}, args...)...).Output() //nolint:gosec // fixed arguments + return strings.TrimSpace(string(b)), err +} + +func lastLines(b []byte) string { + lines := strings.Split(strings.TrimSpace(string(b)), "\n") + if len(lines) > 4 { + lines = lines[len(lines)-4:] + } + return strings.Join(lines, " | ") +} diff --git a/internal/comfy/progress.go b/internal/comfy/progress.go new file mode 100644 index 00000000..51e6a29a --- /dev/null +++ b/internal/comfy/progress.go @@ -0,0 +1,223 @@ +package comfy + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "golang.org/x/net/websocket" +) + +// How far along a workflow is, which ComfyUI says only on its websocket and +// only to the client that queued the workflow. The gate keeps one websocket +// open to its ComfyUI, under a client id of its own, puts that id on every +// workflow it lets through, and keeps the last word on each: the node +// running and its step out of how many. The router asks at +// /progress/{prompt_id} while it waits, and a job says it. + +// Progress is a workflow's state as the gate last heard it. +type Progress struct { + PromptID string `json:"prompt_id,omitempty"` + // State is queued, running, done or failed. + State string `json:"state"` + // Ahead is how many workflows wait before it, while it is queued. + Ahead int `json:"ahead,omitempty"` + // Node is the workflow's node running, by its id; NodeClass its kind. + Node string `json:"node,omitempty"` + NodeClass string `json:"node_class,omitempty"` + // Step of Steps is the running node's own count: a sampler's steps. + Step int `json:"step,omitempty"` + Steps int `json:"steps,omitempty"` + // Percent is Step of Steps, where a node counts. + Percent int `json:"percent,omitempty"` + Error string `json:"error,omitempty"` +} + +// progressBook is the last word on each of the most recent workflows. +type progressBook struct { + mu sync.Mutex + by map[string]*Progress + order []string +} + +// keptPrompts bounds the book: a gate serves one replica, which runs one +// workflow at a time and queues a few. +const keptPrompts = 256 + +func newProgressBook() *progressBook { return &progressBook{by: map[string]*Progress{}} } + +func (b *progressBook) update(id string, f func(*Progress)) { + if id == "" { + return + } + b.mu.Lock() + defer b.mu.Unlock() + p, ok := b.by[id] + if !ok { + p = &Progress{PromptID: id, State: "queued"} + b.by[id] = p + b.order = append(b.order, id) + for len(b.order) > keptPrompts { + delete(b.by, b.order[0]) + b.order = b.order[1:] + } + } + f(p) +} + +func (b *progressBook) get(id string) (Progress, bool) { + b.mu.Lock() + defer b.mu.Unlock() + p, ok := b.by[id] + if !ok { + return Progress{}, false + } + return *p, true +} + +// read takes one message off ComfyUI's websocket into the book. +func (b *progressBook) read(msg []byte) { + var m struct { + Type string `json:"type"` + Data struct { + PromptID string `json:"prompt_id"` + Node json.RawMessage `json:"node"` + Value int `json:"value"` + Max int `json:"max"` + Message string `json:"exception_message"` + NodeType string `json:"node_type"` + } `json:"data"` + } + if json.Unmarshal(msg, &m) != nil { + return + } + d := m.Data + var node string + _ = json.Unmarshal(d.Node, &node) + switch m.Type { + case "execution_start": + b.update(d.PromptID, func(p *Progress) { p.State, p.Ahead = "running", 0 }) + case "executing": + b.update(d.PromptID, func(p *Progress) { + if node == "" { + // No node: the workflow is through. + p.State, p.Node, p.Step, p.Steps, p.Percent = "done", "", 0, 0, 100 + return + } + p.State, p.Ahead, p.Node, p.Step, p.Steps, p.Percent = "running", 0, node, 0, 0, 0 + }) + case "progress": + b.update(d.PromptID, func(p *Progress) { + p.State, p.Step, p.Steps = "running", d.Value, d.Max + if node != "" { + p.Node = node + } + if d.Max > 0 { + p.Percent = min(100, d.Value*100/d.Max) + } + }) + case "execution_success": + b.update(d.PromptID, func(p *Progress) { p.State, p.Percent = "done", 100 }) + case "execution_error", "execution_interrupted": + b.update(d.PromptID, func(p *Progress) { + p.State = "failed" + p.Error = strings.TrimSpace(d.NodeType + ": " + d.Message) + if m.Type == "execution_interrupted" { + p.Error = "interrupted" + } + }) + } +} + +// watch keeps a websocket open to the ComfyUI at upstream under clientID, +// reading into the book, until ctx ends: opened again when it closes, as it +// does whenever ComfyUI restarts. +func (b *progressBook) watch(ctx context.Context, upstream *url.URL, clientID string, log *slog.Logger) { + ws := "ws://" + upstream.Host + "/ws?clientId=" + url.QueryEscape(clientID) + wait := time.Second + for ctx.Err() == nil { + conn, err := websocket.Dial(ws, "", "http://"+upstream.Host) + if err == nil { + wait = time.Second + go func() { <-ctx.Done(); conn.Close() }() + for { + var msg []byte + if err := websocket.Message.Receive(conn, &msg); err != nil { + break + } + // Binary frames are previews of pictures; only text is news. + if len(msg) > 0 && msg[0] == '{' { + b.read(msg) + } + } + conn.Close() + } else if log != nil { + log.Debug("ComfyUI's websocket is not open yet", "err", err) + } + select { + case <-ctx.Done(): + return + case <-time.After(wait): + } + wait = min(wait*2, 30*time.Second) + } +} + +// withClient is a /prompt body with the gate's client id in it, so the +// progress of the workflow it queues comes to the gate's websocket. +func withClient(body []byte, clientID string) []byte { + var m map[string]json.RawMessage + if json.Unmarshal(body, &m) != nil { + return body + } + id, _ := json.Marshal(clientID) + m["client_id"] = id + out, err := json.Marshal(m) + if err != nil { + return body + } + return out +} + +// promptThroughGate forwards a /prompt with the gate's client id, and books +// the workflow it queued as queued. +func (b *progressBook) promptThroughGate(proxy http.Handler, clientID string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(io.LimitReader(r.Body, 64<<20)) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + body = withClient(body, clientID) + r.Body = io.NopCloser(bytes.NewReader(body)) + r.ContentLength = int64(len(body)) + rec := &captured{ResponseWriter: w} + proxy.ServeHTTP(rec, r) + var out struct { + PromptID string `json:"prompt_id"` + } + if json.Unmarshal(rec.body.Bytes(), &out) == nil && out.PromptID != "" { + b.update(out.PromptID, func(*Progress) {}) + } + } +} + +// captured passes a response on and keeps a copy of the first of it. +type captured struct { + http.ResponseWriter + body bytes.Buffer +} + +func (c *captured) Write(p []byte) (int, error) { + if c.body.Len() < 64<<10 { + c.body.Write(p) + } + return c.ResponseWriter.Write(p) +} diff --git a/internal/comfy/progress_test.go b/internal/comfy/progress_test.go new file mode 100644 index 00000000..5e763f9f --- /dev/null +++ b/internal/comfy/progress_test.go @@ -0,0 +1,183 @@ +package comfy + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +func TestTheBookReadsComfyUIsWebsocket(t *testing.T) { + b := newProgressBook() + for _, m := range []string{ + `{"type":"status","data":{"status":{"exec_info":{"queue_remaining":1}}}}`, + `{"type":"execution_start","data":{"prompt_id":"p1"}}`, + `{"type":"executing","data":{"node":"5","prompt_id":"p1"}}`, + `{"type":"progress","data":{"value":12,"max":30,"prompt_id":"p1","node":"5"}}`, + } { + b.read([]byte(m)) + } + p, _ := b.get("p1") + if p.State != "running" || p.Node != "5" || p.Step != 12 || p.Steps != 30 || p.Percent != 40 { + t.Errorf("%+v", p) + } + b.read([]byte(`{"type":"executing","data":{"node":null,"prompt_id":"p1"}}`)) + if p, _ := b.get("p1"); p.State != "done" || p.Percent != 100 { + t.Errorf("finished: %+v", p) + } + b.read([]byte(`{"type":"execution_error","data":{"prompt_id":"p2","node_type":"KSampler","exception_message":"CUDA out of memory"}}`)) + if p, _ := b.get("p2"); p.State != "failed" || p.Error != "KSampler: CUDA out of memory" { + t.Errorf("failed: %+v", p) + } + for i := range keptPrompts + 10 { + b.update(fmt.Sprint("x", i), func(*Progress) {}) + } + if _, ok := b.get("p1"); ok || len(b.by) != keptPrompts { + t.Errorf("the book kept %d", len(b.by)) + } +} + +func TestTheGateHearsTheWorkflowsItLetsThrough(t *testing.T) { + var mu sync.Mutex + var sent map[string]any + comfy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/prompt" { + mu.Lock() + _ = json.NewDecoder(r.Body).Decode(&sent) + mu.Unlock() + fmt.Fprint(w, `{"prompt_id":"p9"}`) + } + })) + defer comfy.Close() + u, _ := url.Parse(comfy.URL) + book := newProgressBook() + gate := httptest.NewServer(newGate(u, "", "m", book, "llmgw-gate-1")) + defer gate.Close() + resp, err := http.Post(gate.URL+"/prompt", "application/json", strings.NewReader(`{"prompt":{},"client_id":"router"}`)) + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if !strings.Contains(string(b), "p9") { + t.Fatalf("the answer did not come through: %s", b) + } + mu.Lock() + if sent["client_id"] != "llmgw-gate-1" { + t.Errorf("ComfyUI was sent client %v: its progress would go to nobody listening", sent["client_id"]) + } + mu.Unlock() + book.read([]byte(`{"type":"progress","data":{"value":3,"max":4,"prompt_id":"p9","node":"2"}}`)) + resp, _ = http.Get(gate.URL + "/progress/p9") + var p Progress + _ = json.NewDecoder(resp.Body).Decode(&p) + resp.Body.Close() + if p.State != "running" || p.Percent != 75 { + t.Errorf("/progress: %+v", p) + } + if resp, _ := http.Get(gate.URL + "/progress/unknown"); resp.StatusCode != http.StatusNotFound { + t.Errorf("an unknown workflow: %d", resp.StatusCode) + } +} + +func TestAWaitingWorkflowSaysHowManyAreAheadAndThenHowFarAlong(t *testing.T) { + var mu sync.Mutex + stage := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + s := stage + mu.Unlock() + switch r.URL.Path { + case "/prompt": + fmt.Fprint(w, `{"prompt_id":"me"}`) + case "/queue": + if s == 0 { + fmt.Fprint(w, `{"queue_running":[[1,"a"]],"queue_pending":[[3,"me"],[2,"b"],[4,"c"]]}`) + } else { + fmt.Fprint(w, `{"queue_running":[[3,"me"]],"queue_pending":[]}`) + } + case "/progress/me": + fmt.Fprint(w, `{"state":"running","node":"5","step":6,"steps":12,"percent":50}`) + case "/history/me": + if s < 2 { + fmt.Fprint(w, `{}`) + return + } + fmt.Fprint(w, `{"me":{"status":{"status_str":"success","completed":true},"outputs":{"9":{"images":[{"filename":"a.png","type":"temp"}]}}}}`) + case "/view": + w.Header().Set("Content-Type", "image/png") + fmt.Fprint(w, "png") + } + })) + defer srv.Close() + var seen []Progress + c := &Client{Poll: time.Millisecond} + done := make(chan error, 1) + go func() { + _, err := c.GenerateWatch(t.Context(), srv.URL, "", map[string]any{"5": map[string]any{"class_type": "KSampler"}}, + func(p Progress) { + mu.Lock() + seen = append(seen, p) + stage++ + mu.Unlock() + }) + done <- err + }() + if err := <-done; err != nil { + t.Fatal(err) + } + if len(seen) < 2 || seen[0].State != "queued" || seen[0].Ahead != 2 { + t.Fatalf("%+v: queued behind the one running and the one numbered before it, wanted first", seen) + } + if seen[1].State != "running" || seen[1].Percent != 50 || seen[1].NodeClass != "KSampler" { + t.Errorf("then %+v", seen[1]) + } +} + +func TestANodeIsPutAtItsCommit(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("no git") + } + repo := t.TempDir() + run := func(args ...string) string { + cmd := exec.Command("git", args...) + cmd.Dir = repo + cmd.Env = append(os.Environ(), "GIT_AUTHOR_NAME=t", "GIT_AUTHOR_EMAIL=t@t", "GIT_COMMITTER_NAME=t", "GIT_COMMITTER_EMAIL=t@t") + b, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("git %v: %v %s", args, err, b) + } + return strings.TrimSpace(string(b)) + } + run("init", "-q") + run("config", "uploadpack.allowReachableSHA1InWant", "true") + _ = os.WriteFile(filepath.Join(repo, "__init__.py"), []byte("v1\n"), 0o644) + run("add", ".") + run("commit", "-qm", "one") + first := run("rev-parse", "HEAD") + _ = os.WriteFile(filepath.Join(repo, "__init__.py"), []byte("v2\n"), 0o644) + run("commit", "-qam", "two") + + state := t.TempDir() + site, err := installNodes(t.Context(), state, "python3", []Node{{Repo: "file://" + repo, Commit: first}}, "", func(string) {}) + if err != nil { + t.Fatal(err) + } + b, _ := os.ReadFile(filepath.Join(state, "custom_nodes", filepath.Base(repo), "__init__.py")) + if string(b) != "v1\n" || site != "" { + t.Errorf("the node is %q, site %q: wanted the commit named, not the branch's head", b, site) + } + if _, err := installNodes(t.Context(), state, "python3", []Node{{Repo: "file://" + repo, Commit: first, + Requirements: []string{"trimesh"}}}, "", func(string) {}); err == nil || !strings.Contains(err.Error(), "no uv") { + t.Errorf("packages without a uv: %v", err) + } +} diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index e5201c2f..39a3209e 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -7,8 +7,10 @@ import ( "encoding/hex" "encoding/json" "fmt" + "maps" "net/url" "path" + "regexp" "slices" "strings" "time" @@ -550,6 +552,51 @@ type ComfyUI struct { // at /v1/images/generations, or meshes, asked for as a job at // /v1/3d/generations from a picture. Output string `json:"output,omitempty" koanf:"output"` + // Nodes are custom nodes the workflow needs, put in each replica's own + // custom_nodes at the commit named. + Nodes []CustomNode `json:"nodes,omitempty" koanf:"nodes"` +} + +// CustomNode is a ComfyUI node from a git repository. +type CustomNode struct { + // Repo is the repository, https only. + Repo string `json:"repo" koanf:"repo"` + // Commit is the full commit the node is used at: a branch moves, and a + // node that changed under a model is a model that changed unasked. + Commit string `json:"commit" koanf:"commit"` + // Requirements are the PyPI packages the node needs that ComfyUI does + // not have, by name and version range only. They are installed beside + // the replica, never into ComfyUI's environment, at the versions that + // environment already has of anything it shares. + Requirements []string `json:"requirements,omitempty" koanf:"requirements"` +} + +// NodesEnv carries a ComfyUI model's custom nodes to the agent, which hands +// them to the gate. +const NodesEnv = "LLMGW_COMFYUI_NODES" + +var ( + fullCommit = regexp.MustCompile(`^[0-9a-f]{40}$`) + packageSpec = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*(\[[A-Za-z0-9,._-]+\])?\s*([<>=!~]=?\s*[A-Za-z0-9.*+!_-]+\s*,?\s*)*$`) +) + +func (n CustomNode) validate(model string) error { + u, err := url.Parse(n.Repo) + if err != nil || u.Scheme != "https" || u.Host == "" || u.User != nil { + return fmt.Errorf("model %q: custom node %q is not an https repository", model, n.Repo) + } + if !fullCommit.MatchString(n.Commit) { + return fmt.Errorf("model %q: custom node %s needs its full commit (40 hex), not %q: a branch moves", + model, n.Repo, n.Commit) + } + for _, r := range n.Requirements { + // A package and a version range: an option or a URL would install + // from somewhere other than PyPI. + if !packageSpec.MatchString(strings.TrimSpace(r)) { + return fmt.Errorf("model %q: custom node %s requirement %q is not a package with a version range", model, n.Repo, r) + } + } + return nil } // ComfyWorkflow is what a ComfyUI model's replicas are sent: the workflow @@ -604,6 +651,13 @@ func (m Model) validateComfyUI() error { if err := checkWorkflow(m.Name, w); err != nil { return err } + if m.ComfyUI != nil { + for _, n := range m.ComfyUI.Nodes { + if err := n.validate(m.Name); err != nil { + return err + } + } + } if out := m.comfyOutput(); out != OutputPictures && out != OutputMeshes { return fmt.Errorf("model %q: comfyui.output %q is not pictures or meshes", m.Name, out) } @@ -1126,6 +1180,15 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) { if p.Name != name { continue } + env := p.Env + if p.Engine == EngineComfyUI && m.ComfyUI != nil && len(m.ComfyUI.Nodes) > 0 { + nodes, _ := json.Marshal(m.ComfyUI.Nodes) + env = maps.Clone(p.Env) + if env == nil { + env = map[string]string{} + } + env[NodesEnv] = string(nodes) + } return engines.Profile{ Name: p.Name, Engine: p.Engine, @@ -1135,7 +1198,7 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) { GPUsPerNode: p.GPUsPerNode, Nodes: p.Nodes, Args: slices.Clone(p.Args), - Env: p.Env, + Env: env, ModelPath: m.SourceOf(p), ServedModelName: m.Name, Adapters: m.engineAdapters(), diff --git a/internal/control/catalog/comfyui_test.go b/internal/control/catalog/comfyui_test.go index dfde5890..149d93b5 100644 --- a/internal/control/catalog/comfyui_test.go +++ b/internal/control/catalog/comfyui_test.go @@ -49,3 +49,35 @@ func TestAComfyUIModelIsRefusedWhatItCannotDo(t *testing.T) { t.Errorf("a set with its checkpoint named was refused: %v", err) } } + +func TestACustomNodeIsPinnedAndItsPackagesArePackages(t *testing.T) { + node := func(n CustomNode) Model { + return comfyModel("hf:stabilityai/TripoSR/model.ckpt", &ComfyUI{Workflow: "w.json", Output: OutputMeshes, + Nodes: []CustomNode{n}}) + } + ok := CustomNode{Repo: "https://github.com/flowtyone/ComfyUI-Flowty-TripoSR", + Commit: "0123456789abcdef0123456789abcdef01234567", Requirements: []string{"omegaconf", "trimesh>=4,<5", "scikit-image"}} + cat, err := New([]Model{node(ok)}) + if err != nil { + t.Fatal(err) + } + m, _ := cat.Resolve("flux") + p, _ := m.EngineProfile("gpu") + if !strings.Contains(p.Env[NodesEnv], "ComfyUI-Flowty-TripoSR") || !m.MakesMeshes() || m.MakesImages() { + t.Errorf("env %v meshes %v pictures %v", p.Env, m.MakesMeshes(), m.MakesImages()) + } + for name, n := range map[string]CustomNode{ + "a branch": {Repo: ok.Repo, Commit: "main"}, + "http": {Repo: "http://example.com/n", Commit: ok.Commit}, + "an option": {Repo: ok.Repo, Commit: ok.Commit, Requirements: []string{"--index-url=https://evil"}}, + "a url": {Repo: ok.Repo, Commit: ok.Commit, Requirements: []string{"x @ https://evil/x.whl"}}, + } { + if _, err := New([]Model{node(n)}); err == nil { + t.Errorf("%s: accepted", name) + } + } + noWorkflow := comfyModel("hf:a/b/c.ckpt", &ComfyUI{Output: OutputMeshes}) + if _, err := New([]Model{noWorkflow}); err == nil || !strings.Contains(err.Error(), "names its workflow") { + t.Errorf("a mesh model without a workflow: %v", err) + } +} diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 4a0e6fc1..cabf1ac4 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -71,6 +71,7 @@ var Reference = []Block{ {Key: "workflow", Type: "string", Doc: "Workflow is a built-in's name (flux-schnell, the default, or\nflux-dev) or a JSON file on the control plane, in ComfyUI's API\nformat.", Nested: ""}, {Key: "checkpoint", Type: "string", Doc: "Checkpoint is the file the workflow's {{checkpoint}} names: the\nsource's own name unless set, which a source that is a set of files\nmust.", Nested: ""}, {Key: "output", Type: "string", Doc: "Output is what the workflow makes: pictures (the default), asked for\nat /v1/images/generations, or meshes, asked for as a job at\n/v1/3d/generations from a picture.", Nested: ""}, + {Key: "nodes", Type: "[]CustomNode", Doc: "Nodes are custom nodes the workflow needs, put in each replica's own\ncustom_nodes at the commit named.", Nested: ""}, }}, {Package: "engines", Name: "ComfyUIConfig", Title: "ComfyUIConfig", Doc: "", Fields: []Field{ {Key: "dir", Type: "string", Doc: "Dir is ComfyUI's directory, where main.py is. Readable by the user\nengines run as.", Nested: ""}, @@ -89,6 +90,11 @@ var Reference = []Block{ {Key: "enabled", Type: "bool", Doc: "Enabled runs the control plane even when configuration declares no\nmodels.\n\nWithout it, a deployment whose catalog lives entirely in the database\ncannot start: control plane mode is inferred from the config file, so an\nempty models block means \"router only\", the process refuses to start for\nwant of anything to serve, and the models it would have loaded are never\nread. That is a cluster that cannot be managed at runtime because it is\nmanaged at runtime. AgentListen cannot carry this signal — it has a\ndefault, so every deployment would become a control plane.", Nested: ""}, {Key: "lease_ttl", Type: "time.Duration", Doc: "LeaseTTL is how long this control plane's claim to be the one placing\nwork stands without being renewed.\n\nOnly meaningful with a shared store, where a second control plane can\nsee the claim. Long enough to survive a slow query or a brief network\nstall, short enough that a standby takes over before anyone has finished\nreading the alert. Zero takes the default.", Nested: ""}, }}, + {Package: "catalog", Name: "CustomNode", Title: "CustomNode", Doc: "", Fields: []Field{ + {Key: "repo", Type: "string", Doc: "Repo is the repository, https only.", Nested: ""}, + {Key: "commit", Type: "string", Doc: "Commit is the full commit the node is used at: a branch moves, and a\nnode that changed under a model is a model that changed unasked.", Nested: ""}, + {Key: "requirements", Type: "[]string", Doc: "Requirements are the PyPI packages the node needs that ComfyUI does\nnot have, by name and version range only. They are installed beside\nthe replica, never into ComfyUI's environment, at the versions that\nenvironment already has of anything it shares.", Nested: ""}, + }}, {Package: "config", Name: "Endpoint", Title: "Endpoint", Doc: "", Fields: []Field{ {Key: "id", Type: "string", Doc: "", Nested: ""}, {Key: "url", Type: "string", Doc: "", Nested: ""}, diff --git a/internal/router/meshes.go b/internal/router/meshes.go index c835b4b4..d2e0a963 100644 --- a/internal/router/meshes.go +++ b/internal/router/meshes.go @@ -143,7 +143,7 @@ func (r *Router) generateMeshes(ctx context.Context, w http.ResponseWriter, req if err != nil { return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) } - files, err := r.comfy.Generate(ctx, ep.URL, ep.APIKey, workflow) + files, err := r.comfy.GenerateWatch(ctx, ep.URL, ep.APIKey, workflow, r.jobProgress(ctx, req)) if err != nil { return fail(err) } @@ -159,6 +159,13 @@ func (r *Router) generateMeshes(ctx context.Context, w http.ResponseWriter, req "its last node should save a .glb", ep.ID, len(files)) } data := mesh.Data + if mesh.ContentType == "model/obj" { + glb, err := objToGLB(data) + if err != nil { + return proxy.Outcome{}, fmt.Errorf("%s: the mesh it saved: %w", ep.ID, err) + } + data, mesh.ContentType = glb, "model/gltf-binary" + } m := r.modelMap()[req.Model] if m.License != nil && mesh.ContentType == "model/gltf-binary" { if with, err := glbWithLicense(data, *m.License); err == nil { @@ -238,3 +245,27 @@ func glbWithLicense(glb []byte, lic catalog.License) ([]byte, error) { out = append(out, js...) return append(out, rest...), nil } + +// jobProgress is where a job's request says how far along it is: into the +// job, for whoever asks after it. Nil for a request that is no job's — and +// for one that only says it is: the job must be the same key's, so nobody +// writes into another's progress by naming it in a tag. +func (r *Router) jobProgress(ctx context.Context, req *parsedRequest) func(comfy.Progress) { + id := req.Tags["job"] + if id == "" || r.store == nil || id != req.RequestID { + return nil + } + j, err := r.store.JobByID(ctx, id) + if err != nil || j.KeyID != req.KeyID { + return nil + } + return func(p comfy.Progress) { + b, err := json.Marshal(p) + if err != nil { + return + } + if err := r.store.SetJobProgress(context.WithoutCancel(ctx), id, string(b), p.PromptID); err != nil { + r.log.Warn("a job's progress could not be kept", "job", id, "err", err) + } + } +} diff --git a/internal/router/meshes_test.go b/internal/router/meshes_test.go index 3fad7fee..76cbef09 100644 --- a/internal/router/meshes_test.go +++ b/internal/router/meshes_test.go @@ -245,3 +245,45 @@ func TestA3DModelIsAskedOfAModelThatMakesThem(t *testing.T) { t.Errorf("/v1/models: %s", rec.Body.String()) } } + +func TestAnOBJLeavesTheGatewayAsAGLB(t *testing.T) { + // TripoSR saves an OBJ with a colour after each position, and a quad + // here stands for the polygons other nodes write. + obj := "# made by a test\nv 0 0 0 1 0 0\nv 1 0 0 0 1 0\nv 1 1 0 0 0 1\nv 0 1 0 1 1 1\nf 1 2 3 4\nf -4 -2 -1\n" + glb, err := objToGLB([]byte(obj)) + if err != nil { + t.Fatal(err) + } + if string(glb[:4]) != "glTF" || int(binary.LittleEndian.Uint32(glb[8:])) != len(glb) { + t.Fatal("not a glTF binary of its own length") + } + n := binary.LittleEndian.Uint32(glb[12:]) + var doc struct { + Accessors []struct { + Count int `json:"count"` + Max []float64 `json:"max"` + } `json:"accessors"` + Meshes []struct { + Primitives []struct { + Attributes map[string]int `json:"attributes"` + } `json:"primitives"` + } `json:"meshes"` + } + if err := json.Unmarshal(glb[20:20+n], &doc); err != nil { + t.Fatal(err) + } + if doc.Accessors[0].Count != 4 || doc.Accessors[1].Count != 9 || doc.Accessors[0].Max[1] != 1 { + t.Errorf("accessors %+v: four vertices and three triangles wanted", doc.Accessors) + } + if _, ok := doc.Meshes[0].Primitives[0].Attributes["COLOR_0"]; !ok { + t.Error("the vertices' colours were lost") + } + if _, err := glbWithLicense(glb, catalog.License{Name: "MIT"}); err != nil { + t.Errorf("the licence cannot be written into it: %v", err) + } + for _, bad := range []string{"v 0 0 0\n", "v 0 0 0\nf 1 2 3\n", "f x y z\n"} { + if _, err := objToGLB([]byte(bad)); err == nil { + t.Errorf("%q made a GLB", bad) + } + } +} diff --git a/internal/router/objglb.go b/internal/router/objglb.go new file mode 100644 index 00000000..03bcc062 --- /dev/null +++ b/internal/router/objglb.go @@ -0,0 +1,146 @@ +package router + +import ( + "bufio" + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "math" + "strconv" + "strings" +) + +// objToGLB turns a Wavefront OBJ — what a 3D node saves when it saves no +// GLB — into a glTF binary: one mesh of its vertices, their colours where +// the file gives them (TripoSR's do, after each position), and its faces +// as triangles. A GLB is the one file a game engine, a browser and the +// licence notice all take, so every mesh leaves the gateway as one. +func objToGLB(obj []byte) ([]byte, error) { + var pos, col []float32 + var idx []uint32 + colours := true + sc := bufio.NewScanner(bytes.NewReader(obj)) + sc.Buffer(make([]byte, 0, 64<<10), 1<<20) + for sc.Scan() { + f := strings.Fields(sc.Text()) + if len(f) == 0 { + continue + } + switch f[0] { + case "v": + if len(f) < 4 { + return nil, errors.New("a vertex with fewer than three coordinates") + } + for _, s := range f[1:4] { + v, err := strconv.ParseFloat(s, 32) + if err != nil { + return nil, fmt.Errorf("a vertex coordinate %q: %w", s, err) + } + pos = append(pos, float32(v)) + } + if len(f) >= 7 { + for _, s := range f[4:7] { + v, err := strconv.ParseFloat(s, 32) + if err != nil { + return nil, fmt.Errorf("a vertex colour %q: %w", s, err) + } + col = append(col, float32(v)) + } + } else { + colours = false + } + case "f": + n := uint32(len(pos) / 3) //nolint:gosec // bounded by the file read + var face []uint32 + for _, s := range f[1:] { + ref, _, _ := strings.Cut(s, "/") + i, err := strconv.Atoi(ref) + if err != nil || i == 0 { + return nil, fmt.Errorf("a face index %q", s) + } + if i < 0 { + i = int(n) + i + 1 + } + if i < 1 || uint32(i) > n { //nolint:gosec // checked positive + return nil, fmt.Errorf("a face index %d past the %d vertices", i, n) + } + face = append(face, uint32(i-1)) //nolint:gosec // checked positive + } + // A polygon as a fan of triangles. + for k := 1; k+1 < len(face); k++ { + idx = append(idx, face[0], face[k], face[k+1]) + } + } + } + if err := sc.Err(); err != nil { + return nil, err + } + if len(pos) == 0 || len(idx) == 0 { + return nil, errors.New("the OBJ has no faces") + } + if !colours || len(col) != len(pos) { + col = nil + } + + var bin bytes.Buffer + view := func(data any) (offset, length int) { + offset = bin.Len() + _ = binary.Write(&bin, binary.LittleEndian, data) + length = bin.Len() - offset + for bin.Len()%4 != 0 { + bin.WriteByte(0) + } + return offset, length + } + lo := []float32{math.MaxFloat32, math.MaxFloat32, math.MaxFloat32} + hi := []float32{-math.MaxFloat32, -math.MaxFloat32, -math.MaxFloat32} + for i, v := range pos { + lo[i%3], hi[i%3] = min(lo[i%3], v), max(hi[i%3], v) + } + po, pl := view(pos) + io, il := view(idx) + views := []map[string]any{ + {"buffer": 0, "byteOffset": po, "byteLength": pl, "target": 34962}, + {"buffer": 0, "byteOffset": io, "byteLength": il, "target": 34963}, + } + accessors := []map[string]any{ + {"bufferView": 0, "componentType": 5126, "count": len(pos) / 3, "type": "VEC3", "min": lo, "max": hi}, + {"bufferView": 1, "componentType": 5125, "count": len(idx), "type": "SCALAR"}, + } + attrs := map[string]any{"POSITION": 0} + if col != nil { + co, cl := view(col) + views = append(views, map[string]any{"buffer": 0, "byteOffset": co, "byteLength": cl, "target": 34962}) + accessors = append(accessors, map[string]any{"bufferView": 2, "componentType": 5126, "count": len(col) / 3, "type": "VEC3"}) + attrs["COLOR_0"] = 2 + } + doc := map[string]any{ + "asset": map[string]any{"version": "2.0", "generator": "llmgw"}, + "scene": 0, + "scenes": []any{map[string]any{"nodes": []int{0}}}, + "nodes": []any{map[string]any{"mesh": 0}}, + "meshes": []any{map[string]any{"primitives": []any{map[string]any{"attributes": attrs, "indices": 1}}}}, + "accessors": accessors, + "bufferViews": views, + "buffers": []any{map[string]any{"byteLength": bin.Len()}}, + } + js, err := json.Marshal(doc) + if err != nil { + return nil, err + } + for len(js)%4 != 0 { + js = append(js, ' ') + } + out := make([]byte, 0, 28+len(js)+bin.Len()) + out = append(out, "glTF"...) + out = binary.LittleEndian.AppendUint32(out, 2) + out = binary.LittleEndian.AppendUint32(out, uint32(28+len(js)+bin.Len())) //nolint:gosec // bounded by the file read + out = binary.LittleEndian.AppendUint32(out, uint32(len(js))) //nolint:gosec // bounded by the file read + out = append(out, "JSON"...) + out = append(out, js...) + out = binary.LittleEndian.AppendUint32(out, uint32(bin.Len())) //nolint:gosec // bounded by the file read + out = append(out, "BIN\x00"...) + return append(out, bin.Bytes()...), nil +}