diff --git a/BACKLOG.md b/BACKLOG.md index 1429421f..4641e9bd 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -1364,7 +1364,14 @@ by the workflow's `{{checkpoint}}`. each other with a 30 s cold start each time) is the other session's work (llmgw-b7), built on this. -### H4. Custom nodes with packages that are built, not downloaded (S) +### H4. Custom nodes with packages that are built, not downloaded (S) — built 2026-10-11 +Built as written, with `uv build --no-build-isolation` against ComfyUI's +Python, the card's capability asked of torch in the replica (so +CUDA_VISIBLE_DEVICES is the card it will run on), and wheels kept under +the engines' cache by commit, `sm` and directory. Also +`comfyui.huggingface_repo`: a node that loads its weights by repository +name (SF3D's) finds them in the engines' Hugging Face cache, linked by the +gate, so no token is on the node. - `nodes[].build: [subdir, ...]`: directories of the node's repository that are Python packages with C++/CUDA extensions (SF3D's texture_baker and uv_unwrapper). The gate builds each as a wheel with the engine's diff --git a/CHANGELOG.md b/CHANGELOG.md index 86ee085b..1c806c60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,11 @@ APIs change between minors without a deprecation period. `@`; a ComfyUI profile may carry its own `comfyui` workflow, checkpoint and nodes, and any profile its own `source`. One catalog entry then runs fp8 where the card has it and fp16 on Turing. +- **Custom nodes that compile**: `nodes[].build` names directories of a + node that are C++/CUDA packages; the gate builds each as a wheel for the + card the replica runs on, keeps it on the node, and installs it beside + the replica. `comfyui.huggingface_repo` shows a model's weights to a node + that loads them by repository name. - **A licence a machine can match**: `license.id`, `commercial`, `excluded_regions`, `revenue_cap_usd`, `attribution`. - **On-demand models in `/v1/models`**, with `loaded` saying whether one is diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 6afe6585..746e0f53 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -8790,6 +8790,10 @@ }, "description": "Nodes are custom nodes the workflow needs, put in each replica's own custom_nodes at the commit named." }, + "huggingface_repo": { + "type": "string", + "description": "HuggingFaceRepo is the repository a custom node loads the weights by name from (Stable Fast 3D's loader asks for stabilityai/stable-fast-3d): the replica's weights are shown to the node there, in the Hugging Face cache of the engines' user, so nothing is fetched and no token is needed when it runs." + }, "environment": { "type": "string", "description": "Environment names the models that may share one ComfyUI replica: every model of an environment agrees on the custom nodes they both name, so one replica can carry all of them and ComfyUI moves between their weights itself. Empty is an environment of the model alone." @@ -8814,6 +8818,13 @@ "type": "string" }, "description": "Requirements are the PyPI packages the node needs, by name and version range only. They are installed beside the replica, never into ComfyUI's environment: a package named here at the range named, and every other one at the version ComfyUI already runs. One named here is first on ComfyUI's path, so its range must suit ComfyUI too." + }, + "build": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Build are directories of the repository that are Python packages to be compiled rather than downloaded: C++ and CUDA extensions no index has (Stable Fast 3D's texture_baker and uv_unwrapper). Each is built as a wheel with ComfyUI's Python for the card the replica runs on, kept on the node for the next replica on the same kind of card, and installed beside the replica like its requirements." } } }, diff --git a/docs/reference.md b/docs/reference.md index d547e522..d63f5a89 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -67,6 +67,7 @@ rather than repeated. | `output` | `string` | [ComfyUI](#comfyui) | | `inputs` | `[]Input` | [ComfyUI](#comfyui) | | `nodes` | `[]CustomNode` | [ComfyUI](#comfyui) | +| `huggingface_repo` | `string` | [ComfyUI](#comfyui) | | `environment` | `string` | [ComfyUI](#comfyui) | | `dir` | `string` | [ComfyUIConfig](#comfyuiconfig) | | `python` | `string` | [ComfyUIConfig](#comfyuiconfig) | @@ -84,6 +85,7 @@ rather than repeated. | `repo` | `string` | [CustomNode](#customnode) | | `commit` | `string` | [CustomNode](#customnode) | | `requirements` | `[]string` | [CustomNode](#customnode) | +| `build` | `[]string` | [CustomNode](#customnode) | | `id` | `string` | [Endpoint](#endpoint) | | `url` | `string` | [Endpoint](#endpoint) | | `api_key` | `string` | [Endpoint](#endpoint) | @@ -824,6 +826,16 @@ Inputs are what the workflow takes from a request at Nodes are custom nodes the workflow needs, put in each replica's own custom_nodes at the commit named. +### `huggingface_repo` + +`string` + +HuggingFaceRepo is the repository a custom node loads the weights by +name from (Stable Fast 3D's loader asks for stabilityai/stable-fast-3d): +the replica's weights are shown to the node there, in the Hugging Face +cache of the engines' user, so nothing is fetched and no token is +needed when it runs. + ### `environment` `string` @@ -958,6 +970,17 @@ ComfyUI's environment: a package named here at the range named, and every other one at the version ComfyUI already runs. One named here is first on ComfyUI's path, so its range must suit ComfyUI too. +### `build` + +`[]string` + +Build are directories of the repository that are Python packages to +be compiled rather than downloaded: C++ and CUDA extensions no index +has (Stable Fast 3D's texture_baker and uv_unwrapper). Each is built +as a wheel with ComfyUI's Python for the card the replica runs on, +kept on the node for the next replica on the same kind of card, and +installed beside the replica like its requirements. + ## Endpoint ### `id` diff --git a/internal/comfy/gate.go b/internal/comfy/gate.go index 7e696d7b..63c81591 100644 --- a/internal/comfy/gate.go +++ b/internal/comfy/gate.go @@ -3,6 +3,7 @@ package comfy import ( "context" "crypto/rand" + "crypto/sha1" "crypto/subtle" "encoding/hex" "encoding/json" @@ -114,6 +115,11 @@ func RunGate(ctx context.Context, cfg GateConfig) error { if err := LinkWeights(cfg.Weights, models); err != nil { return err } + if repo := os.Getenv(HFRepoEnv); repo != "" { + if err := LinkHFCache(cfg.Weights, repo, hfHubDir()); err != nil { + return fmt.Errorf("showing the weights as %s: %w", repo, err) + } + } extra := filepath.Join(state, "extra_model_paths.yaml") if err := os.WriteFile(extra, []byte(extraModelPaths(models)), 0o644); err != nil { //nolint:gosec // not secret return err @@ -383,3 +389,71 @@ func relay(r io.Reader, w io.Writer) { } } } + +// HFRepoEnv is where the agent hands the gate the repository a custom node +// loads the weights by name from. +const HFRepoEnv = "LLMGW_COMFYUI_HF_REPO" + +// hfHubDir is the Hugging Face cache of the user the gate runs as. +func hfHubDir() string { + if d := os.Getenv("HF_HUB_CACHE"); d != "" { + return d + } + if d := os.Getenv("HF_HOME"); d != "" { + return filepath.Join(d, "hub") + } + home, _ := os.UserHomeDir() + return filepath.Join(home, ".cache", "huggingface", "hub") +} + +// LinkHFCache shows the weights to a node that loads them by repository +// name: a snapshot of repo in the Hugging Face cache at hub, each file a link +// to the weights, refs/main naming it. huggingface_hub then answers from the +// cache, and a gated repository needs no token on the node. +func LinkHFCache(weights, repo, hub string) error { + owner, name, ok := strings.Cut(repo, "/") + if !ok || owner == "" || name == "" || strings.Contains(name, "/") || strings.Contains(repo, "..") || strings.Contains(repo, "\\") { + return fmt.Errorf("%q is not owner/name", repo) + } + weights, err := filepath.Abs(weights) + if err != nil { + return err + } + // The snapshot is named as a commit is, which is what refs/main must + // hold, and by the weights, so other weights are another snapshot. + sum := sha1.Sum([]byte(weights)) //nolint:gosec // a name, not a secret + revision := hex.EncodeToString(sum[:]) + info, err := os.Stat(weights) + if err != nil { + return err + } + files := []string{weights} + if info.IsDir() { + entries, err := os.ReadDir(weights) + if err != nil { + return err + } + files = files[:0] + for _, e := range entries { + if !e.IsDir() && !strings.HasPrefix(e.Name(), ".") { + files = append(files, filepath.Join(weights, e.Name())) + } + } + } + root := filepath.Join(hub, "models--"+owner+"--"+name) + snap := filepath.Join(root, "snapshots", revision) + // Made again each time: links to weights that moved must not stay. + _ = os.RemoveAll(snap) + if err := os.MkdirAll(snap, 0o755); err != nil { + return err + } + for _, f := range files { + if err := os.Symlink(f, filepath.Join(snap, FileName(f))); err != nil { + return err + } + } + if err := os.MkdirAll(filepath.Join(root, "refs"), 0o755); err != nil { + return err + } + return os.WriteFile(filepath.Join(root, "refs", "main"), []byte(revision), 0o644) //nolint:gosec // not secret +} diff --git a/internal/comfy/nodes.go b/internal/comfy/nodes.go index 4d4edafa..2fcae15a 100644 --- a/internal/comfy/nodes.go +++ b/internal/comfy/nodes.go @@ -11,6 +11,9 @@ import ( "path" "path/filepath" "regexp" + "runtime" + "slices" + "strconv" "strings" "time" ) @@ -25,6 +28,8 @@ type Node struct { Repo string `json:"repo"` Commit string `json:"commit"` Requirements []string `json:"requirements,omitempty"` + // Build are directories of the node that are packages to compile. + Build []string `json:"build,omitempty"` } //go:embed all:nodes @@ -65,12 +70,17 @@ func installNodes(ctx context.Context, state, python string, nodes []Node, uv st } ctx, cancel := context.WithTimeout(ctx, 20*time.Minute) defer cancel() - var reqs []string + var reqs, built []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...) + wheels, err := buildNode(ctx, n, dir, python, uv, out) + if err != nil { + return "", err + } + built = append(built, wheels...) continue } out(fmt.Sprintf("custom node %s at %s", n.Repo, n.Commit[:12])) @@ -86,13 +96,25 @@ func installNodes(ctx context.Context, state, python string, nodes []Node, uv st } } reqs = append(reqs, n.Requirements...) + wheels, err := buildNode(ctx, n, dir, python, uv, out) + if err != nil { + return "", err + } + built = append(built, wheels...) } - if len(reqs) == 0 { + if len(reqs) == 0 && len(built) == 0 { return "", nil } if uv == "" { return "", fmt.Errorf("custom nodes need %s, and no uv is installed to install them with", strings.Join(reqs, ", ")) } + // A built package is named by its wheel; ComfyUI's own copy, if one was + // ever put there, must not pin it to another. + named := slices.Clone(reqs) + for _, w := range built { + named = append(named, wheelPackage(w)) + } + reqs = append(reqs, built...) site := filepath.Join(state, "site") marker := filepath.Join(site, ".llmgw-requirements") want := strings.Join(reqs, "\n") @@ -108,7 +130,7 @@ func installNodes(ctx context.Context, state, python string, nodes []Node, uv st 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 + if err := os.WriteFile(constraints, withoutNamed(freeze, named), 0o644); err != nil { //nolint:gosec // not secret return "", err } _ = os.RemoveAll(site) @@ -190,3 +212,123 @@ func packageName(spec string) string { } return strings.NewReplacer("_", "-", ".", "-").Replace(strings.ToLower(s)) } + +// Packages a node's repository has that no index does — C++ and CUDA +// extensions — are built here, for the card the replica runs on: a CUDA +// kernel compiled for one architecture does not load on another, and a +// site's cards are not the PoC's. A build is minutes; its wheel is kept on +// the node, by the node's commit, the directory and the architecture, for +// the next replica on the same kind of card. + +// cudaArch is the compute capability of the card the replica sees, as +// TORCH_CUDA_ARCH_LIST wants it ("7.5"). A var for the tests. +var cudaArch = func(ctx context.Context, python string) (string, error) { + cmd := exec.CommandContext(ctx, python, "-c", //nolint:gosec // the agent's own Python + "import torch; print('%d.%d' % torch.cuda.get_device_capability())") + b, err := cmd.Output() + if err != nil { + return "", fmt.Errorf("asking torch for the card's compute capability: %v", err) + } + return strings.TrimSpace(string(b)), nil +} + +// buildWheel builds the package in dir as a wheel into outDir, with +// ComfyUI's Python and what it already has (torch, setuptools): a CUDA +// extension builds against the torch it will load with. A var for the tests. +var buildWheel = func(ctx context.Context, uv, python, dir, outDir string, env []string) ([]byte, error) { + cmd := exec.CommandContext(ctx, uv, "build", "--wheel", "--no-build-isolation", //nolint:gosec // the catalog's node, at its commit + "--python", python, "--out-dir", outDir, dir) + cmd.Env = env + return cmd.CombinedOutput() +} + +// buildCache is where built wheels are kept: under the engines' home, which +// outlives a replica. +func buildCache() string { + if d, err := os.UserCacheDir(); err == nil { + return filepath.Join(d, "llmgw", "node-builds") + } + return filepath.Join(os.TempDir(), "llmgw-node-builds") +} + +// buildNode builds each of a node's packages that are compiled, or finds the +// wheel built before for this commit and card, and returns the wheels. +func buildNode(ctx context.Context, n Node, dir, python, uv string, out func(string)) ([]string, error) { + if len(n.Build) == 0 { + return nil, nil + } + if uv == "" { + return nil, fmt.Errorf("custom node %s has packages to build, and no uv is installed to build them with", n.Repo) + } + arch, err := cudaArch(ctx, python) + if err != nil { + return nil, fmt.Errorf("custom node %s: %w", n.Repo, err) + } + cudaHome := os.Getenv("CUDA_HOME") + if cudaHome == "" { + for _, c := range []string{"/usr/local/cuda"} { + if _, err := os.Stat(filepath.Join(c, "bin", "nvcc")); err == nil { + cudaHome = c + } + } + } + var wheels []string + for _, sub := range n.Build { + cached := filepath.Join(buildCache(), n.Commit[:12], "sm"+strings.ReplaceAll(arch, ".", ""), nodeDir.ReplaceAllString(sub, "_")) + if w, ok := firstWheel(cached); ok { + wheels = append(wheels, w) + continue + } + if cudaHome == "" { + return nil, fmt.Errorf("custom node %s: %s is compiled, and this node has no CUDA toolkit (no nvcc in "+ + "CUDA_HOME or /usr/local/cuda) to compile it with", n.Repo, sub) + } + out(fmt.Sprintf("building %s of %s for compute %s", sub, n.Repo, arch)) + // Built beside where it is kept and moved there whole: two replicas + // starting at once on one node must not see half a wheel. + if err := os.MkdirAll(filepath.Dir(cached), 0o755); err != nil { + return nil, err + } + tmp, err := os.MkdirTemp(filepath.Dir(cached), ".build-") + if err != nil { + return nil, err + } + env := append(os.Environ(), "TORCH_CUDA_ARCH_LIST="+arch, "CUDA_HOME="+cudaHome, + "PATH="+filepath.Join(cudaHome, "bin")+string(os.PathListSeparator)+os.Getenv("PATH"), + "UV_NO_CONFIG=1", "MAX_JOBS="+strconv.Itoa(max(runtime.NumCPU(), 1))) + b, err := buildWheel(ctx, uv, python, filepath.Join(dir, sub), tmp, env) + if err != nil { + _ = os.RemoveAll(tmp) + return nil, fmt.Errorf("custom node %s: building %s: %v: %s", n.Repo, sub, err, lastLines(b)) + } + if _, ok := firstWheel(tmp); !ok { + _ = os.RemoveAll(tmp) + return nil, fmt.Errorf("custom node %s: building %s made no wheel: %s", n.Repo, sub, lastLines(b)) + } + if err := os.Rename(tmp, cached); err != nil { + _ = os.RemoveAll(tmp) + // Another replica got there first: its wheel is as good. + if _, ok := firstWheel(cached); !ok { + return nil, err + } + } + w, _ := firstWheel(cached) + wheels = append(wheels, w) + } + return wheels, nil +} + +func firstWheel(dir string) (string, bool) { + m, _ := filepath.Glob(filepath.Join(dir, "*.whl")) + if len(m) == 0 { + return "", false + } + return m[0], true +} + +// wheelPackage is the package a wheel file is of: its name up to the first +// '-', as PyPI compares names. +func wheelPackage(wheel string) string { + name, _, _ := strings.Cut(filepath.Base(wheel), "-") + return packageName(name) +} diff --git a/internal/comfy/nodes_test.go b/internal/comfy/nodes_test.go new file mode 100644 index 00000000..7a9369db --- /dev/null +++ b/internal/comfy/nodes_test.go @@ -0,0 +1,165 @@ +package comfy + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strings" + "sync/atomic" + "testing" +) + +// fakeUV is a uv that answers pip freeze with ComfyUI's packages and records +// what pip install was given. +func fakeUV(t *testing.T, freeze string) (uv, log string) { + t.Helper() + dir := t.TempDir() + log = filepath.Join(dir, "install.log") + uv = filepath.Join(dir, "uv") + script := "#!/bin/sh\nif [ \"$2\" = freeze ]; then printf '%s' '" + freeze + "'; exit 0; fi\n" + + "echo \"$@\" >> " + log + "\n" + + "prev=; for a in \"$@\"; do if [ \"$prev\" = --target ]; then mkdir -p \"$a\"; fi; prev=$a; done\nexit 0\n" + if err := os.WriteFile(uv, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + return uv, log +} + +func gitRepo(t *testing.T, files map[string]string) (repo, commit string) { + t.Helper() + 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") + for name, body := range files { + _ = os.MkdirAll(filepath.Dir(filepath.Join(repo, name)), 0o755) + _ = os.WriteFile(filepath.Join(repo, name), []byte(body), 0o644) + } + run("add", ".") + run("commit", "-qm", "one") + return repo, run("rev-parse", "HEAD") +} + +func TestANodesCompiledPackagesAreBuiltForTheCardOnceAndInstalledBesideTheReplica(t *testing.T) { + repo, commit := gitRepo(t, map[string]string{"__init__.py": "", "texture_baker/setup.py": "", "uv_unwrapper/setup.py": ""}) + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + t.Setenv("HOME", t.TempDir()) + cuda := t.TempDir() + _ = os.MkdirAll(filepath.Join(cuda, "bin"), 0o755) + _ = os.WriteFile(filepath.Join(cuda, "bin", "nvcc"), nil, 0o755) + t.Setenv("CUDA_HOME", cuda) + + var builds atomic.Int32 + var envSeen []string + oldArch, oldBuild := cudaArch, buildWheel + t.Cleanup(func() { cudaArch, buildWheel = oldArch, oldBuild }) + cudaArch = func(context.Context, string) (string, error) { return "7.5", nil } + buildWheel = func(_ context.Context, _, _, dir, out string, env []string) ([]byte, error) { + builds.Add(1) + envSeen = env + name := strings.ReplaceAll(filepath.Base(dir), "_", "-") + return nil, os.WriteFile(filepath.Join(out, strings.ReplaceAll(name, "-", "_")+"-0.0.1-cp312-cp312-linux_x86_64.whl"), []byte("w"), 0o644) + } + // ComfyUI's environment has an old hand-installed copy, which must not + // pin the built one. + uv, log := fakeUV(t, "numpy==2.5.3\ntexture-baker @ file:///home/someone/texture_baker-0.0.1.whl\n") + node := Node{Repo: "file://" + repo, Commit: commit, Requirements: []string{"trimesh"}, + Build: []string{"texture_baker", "uv_unwrapper"}} + + state := t.TempDir() + site, err := installNodes(t.Context(), state, "python3", []Node{node}, uv, func(string) {}) + if err != nil { + t.Fatal(err) + } + if site == "" || builds.Load() != 2 { + t.Fatalf("site %q, %d builds", site, builds.Load()) + } + if !strings.Contains(strings.Join(envSeen, " "), "TORCH_CUDA_ARCH_LIST=7.5") { + t.Error("the build was not told the card's architecture") + } + installed, _ := os.ReadFile(log) + if !strings.Contains(string(installed), "trimesh") || !strings.Contains(string(installed), "sm75/texture_baker/texture_baker-0.0.1") || + !strings.Contains(string(installed), "uv_unwrapper-0.0.1") { + t.Errorf("installed %s", installed) + } + constraints, _ := os.ReadFile(filepath.Join(state, "constraints.txt")) + if strings.Contains(string(constraints), "texture-baker") || !strings.Contains(string(constraints), "numpy==2.5.3") { + t.Errorf("constraints %s", constraints) + } + + // A second replica on the same kind of card finds the wheels built. + if _, err := installNodes(t.Context(), t.TempDir(), "python3", []Node{node}, uv, func(string) {}); err != nil { + t.Fatal(err) + } + if builds.Load() != 2 { + t.Errorf("built again: %d builds", builds.Load()) + } + // Another kind of card builds its own. + cudaArch = func(context.Context, string) (string, error) { return "8.9", nil } + if _, err := installNodes(t.Context(), t.TempDir(), "python3", []Node{node}, uv, func(string) {}); err != nil { + t.Fatal(err) + } + if builds.Load() != 4 { + t.Errorf("an Ada used Turing's wheels: %d builds", builds.Load()) + } +} + +func TestANodeToCompileOnAMachineWithoutCUDASaysSo(t *testing.T) { + repo, commit := gitRepo(t, map[string]string{"__init__.py": "", "ext/setup.py": ""}) + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + t.Setenv("HOME", t.TempDir()) + t.Setenv("CUDA_HOME", "") + if _, err := os.Stat("/usr/local/cuda/bin/nvcc"); err == nil { + t.Skip("this machine has a CUDA toolkit") + } + oldArch := cudaArch + t.Cleanup(func() { cudaArch = oldArch }) + cudaArch = func(context.Context, string) (string, error) { return "8.0", nil } + uv, _ := fakeUV(t, "") + _, err := installNodes(t.Context(), t.TempDir(), "python3", []Node{{Repo: "file://" + repo, Commit: commit, Build: []string{"ext"}}}, + uv, func(string) {}) + if err == nil || !strings.Contains(err.Error(), "no CUDA toolkit") { + t.Errorf("%v", err) + } +} + +func TestWeightsANodeLoadsByRepositoryNameAreInTheHuggingFaceCache(t *testing.T) { + weights := t.TempDir() + _ = os.WriteFile(filepath.Join(weights, "config.yaml"), []byte("c"), 0o644) + _ = os.WriteFile(filepath.Join(weights, "model.safetensors"), []byte("m"), 0o644) + hub := t.TempDir() + if err := LinkHFCache(weights, "stabilityai/stable-fast-3d", hub); err != nil { + t.Fatal(err) + } + root := filepath.Join(hub, "models--stabilityai--stable-fast-3d") + ref, _ := os.ReadFile(filepath.Join(root, "refs", "main")) + if len(ref) != 40 { + t.Fatalf("refs/main %q is not a commit's form", ref) + } + b, err := os.ReadFile(filepath.Join(root, "snapshots", string(ref), "model.safetensors")) + if err != nil || string(b) != "m" { + t.Errorf("the weights through the cache: %q %v", b, err) + } + // Again, as a replica started a second time does. + if err := LinkHFCache(weights, "stabilityai/stable-fast-3d", hub); err != nil { + t.Errorf("a second time: %v", err) + } + for _, bad := range []string{"noslash", "a/b/c", "../x", "a/.."} { + if err := LinkHFCache(weights, bad, hub); err == nil { + t.Errorf("%q taken for a repository", bad) + } + } +} diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 1673e8fa..730d7a52 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -608,6 +608,12 @@ type ComfyUI struct { // 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"` + // HuggingFaceRepo is the repository a custom node loads the weights by + // name from (Stable Fast 3D's loader asks for stabilityai/stable-fast-3d): + // the replica's weights are shown to the node there, in the Hugging Face + // cache of the engines' user, so nothing is fetched and no token is + // needed when it runs. + HuggingFaceRepo string `json:"huggingface_repo,omitempty" koanf:"huggingface_repo"` // Environment names the models that may share one ComfyUI replica: every // model of an environment agrees on the custom nodes they both name, so // one replica can carry all of them and ComfyUI moves between their @@ -628,12 +634,24 @@ type CustomNode struct { // every other one at the version ComfyUI already runs. One named here // is first on ComfyUI's path, so its range must suit ComfyUI too. Requirements []string `json:"requirements,omitempty" koanf:"requirements"` + // Build are directories of the repository that are Python packages to + // be compiled rather than downloaded: C++ and CUDA extensions no index + // has (Stable Fast 3D's texture_baker and uv_unwrapper). Each is built + // as a wheel with ComfyUI's Python for the card the replica runs on, + // kept on the node for the next replica on the same kind of card, and + // installed beside the replica like its requirements. + Build []string `json:"build,omitempty" koanf:"build"` } // NodesEnv carries a ComfyUI model's custom nodes to the agent, which hands // them to the gate. const NodesEnv = "LLMGW_COMFYUI_NODES" +// HFRepoEnv carries the repository a node loads the weights by name from. +const HFRepoEnv = "LLMGW_COMFYUI_HF_REPO" + +var hfRepo = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$`) + 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*)*$`) @@ -648,6 +666,12 @@ func (n CustomNode) validate(model string) error { 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 _, b := range n.Build { + clean := path.Clean(b) + if b == "" || path.IsAbs(b) || clean == "." || clean == ".." || strings.HasPrefix(clean, "../") || strings.ContainsAny(b, "\\ ") { + return fmt.Errorf("model %q: custom node %s build %q is not a directory inside the repository", model, n.Repo, b) + } + } for _, r := range n.Requirements { // A package and a version range: an option or a URL would install // from somewhere other than PyPI. @@ -720,6 +744,9 @@ func (m Model) validateComfyUI() error { if err := validOutput(m.Name, m.comfyOutput()); err != nil { return err } + if m.ComfyUI != nil && m.ComfyUI.HuggingFaceRepo != "" && !hfRepo.MatchString(m.ComfyUI.HuggingFaceRepo) { + return fmt.Errorf("model %q: comfyui.huggingface_repo %q is not owner/name", m.Name, m.ComfyUI.HuggingFaceRepo) + } if out := m.comfyOutput(); out != OutputPictures && w == "" { return fmt.Errorf("model %q: a model that makes %s names its workflow: no built-in makes them", m.Name, out) } @@ -1264,12 +1291,19 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) { env := p.Env if nodes := m.ComfyNodesFor(p); p.Engine == EngineComfyUI && len(nodes) > 0 { nodes, _ := json.Marshal(nodes) - env = maps.Clone(p.Env) + env = maps.Clone(env) if env == nil { env = map[string]string{} } env[NodesEnv] = string(nodes) } + if p.Engine == EngineComfyUI && m.ComfyUI != nil && m.ComfyUI.HuggingFaceRepo != "" { + env = maps.Clone(env) + if env == nil { + env = map[string]string{} + } + env[HFRepoEnv] = m.ComfyUI.HuggingFaceRepo + } return engines.Profile{ Name: p.Name, Engine: p.Engine, diff --git a/internal/control/catalog/environments.go b/internal/control/catalog/environments.go index 25d708c0..0686f299 100644 --- a/internal/control/catalog/environments.go +++ b/internal/control/catalog/environments.go @@ -110,6 +110,14 @@ func (c *Catalog) environments() ([]Environment, error) { nodeAt[env][repo], nodeFrom[env][repo] = n.Commit, name e.Nodes = append(e.Nodes, CustomNode{Repo: n.Repo, Commit: n.Commit}) } + i := slices.IndexFunc(e.Nodes, func(x CustomNode) bool { + return strings.TrimSuffix(strings.TrimRight(x.Repo, "/"), ".git") == repo + }) + for _, b := range n.Build { + if !slices.Contains(e.Nodes[i].Build, b) { + e.Nodes[i].Build = append(e.Nodes[i].Build, b) + } + } for _, r := range n.Requirements { pkg, spec := requirementName(r), normalizeSpec(r) if prev, seen := pkgSpec[env][pkg]; seen { @@ -123,9 +131,6 @@ func (c *Catalog) environments() ([]Environment, error) { pkgSpec[env][pkg], pkgFrom[env][pkg] = spec, name // The requirement goes with the node that first named it: // what the replica installs is the union, and each once. - i := slices.IndexFunc(e.Nodes, func(x CustomNode) bool { - return strings.TrimSuffix(strings.TrimRight(x.Repo, "/"), ".git") == repo - }) e.Nodes[i].Requirements = append(e.Nodes[i].Requirements, strings.TrimSpace(r)) } } diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 4658a21f..18aca251 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -73,6 +73,7 @@ var Reference = []Block{ {Key: "output", Type: "string", Doc: "Output is what the workflow makes: pictures (the default), asked for\nat /v1/images/generations; meshes, asked for as a job at\n/v1/3d/generations from a picture; or video, audio or files. Every\nkind is asked for as a job at /v1/workflows/runs too.", Nested: ""}, {Key: "inputs", Type: "[]Input", Doc: "Inputs are what the workflow takes from a request at\n/v1/workflows/runs, each the placeholder of its name.", 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: ""}, + {Key: "huggingface_repo", Type: "string", Doc: "HuggingFaceRepo is the repository a custom node loads the weights by\nname from (Stable Fast 3D's loader asks for stabilityai/stable-fast-3d):\nthe replica's weights are shown to the node there, in the Hugging Face\ncache of the engines' user, so nothing is fetched and no token is\nneeded when it runs.", Nested: ""}, {Key: "environment", Type: "string", Doc: "Environment names the models that may share one ComfyUI replica: every\nmodel of an environment agrees on the custom nodes they both name, so\none replica can carry all of them and ComfyUI moves between their\nweights itself. Empty is an environment of the model alone.", Nested: ""}, }}, {Package: "engines", Name: "ComfyUIConfig", Title: "ComfyUIConfig", Doc: "", Fields: []Field{ @@ -96,6 +97,7 @@ var Reference = []Block{ {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, by name and version\nrange only. They are installed beside the replica, never into\nComfyUI's environment: a package named here at the range named, and\nevery other one at the version ComfyUI already runs. One named here\nis first on ComfyUI's path, so its range must suit ComfyUI too.", Nested: ""}, + {Key: "build", Type: "[]string", Doc: "Build are directories of the repository that are Python packages to\nbe compiled rather than downloaded: C++ and CUDA extensions no index\nhas (Stable Fast 3D's texture_baker and uv_unwrapper). Each is built\nas a wheel with ComfyUI's Python for the card the replica runs on,\nkept on the node for the next replica on the same kind of card, and\ninstalled beside the replica like its requirements.", Nested: ""}, }}, {Package: "config", Name: "Endpoint", Title: "Endpoint", Doc: "", Fields: []Field{ {Key: "id", Type: "string", Doc: "", Nested: ""},