llmgw/internal/comfy/nodes.go

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, withoutNamed(freeze, reqs), 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, " | ")
}

// withoutNamed is ComfyUI's packages less those a node names itself: of a
// package a node asks for by name, the node's range decides — TripoSR's
// checkpoint loads with transformers 4 and not 5 — and of everything else,
// what ComfyUI runs. A node that names a package must name a range ComfyUI
// can run with too: the copy it gets is first on ComfyUI's path.
func withoutNamed(freeze []byte, reqs []string) []byte {
	named := map[string]bool{}
	for _, r := range reqs {
		named[packageName(r)] = true
	}
	var out []string
	for _, line := range strings.Split(string(freeze), "\n") {
		if strings.TrimSpace(line) == "" || named[packageName(line)] {
			continue
		}
		out = append(out, line)
	}
	return []byte(strings.Join(out, "\n") + "\n")
}

// packageName is a requirement's package, as PyPI compares them.
func packageName(spec string) string {
	s := strings.TrimSpace(spec)
	if i := strings.IndexAny(s, "<>=!~;[ @"); i >= 0 {
		s = s[:i]
	}
	return strings.NewReplacer("_", "-", ".", "-").Replace(strings.ToLower(s))
}