G2: the gate keeps a websocket open to its ComfyUI under a client id of its
own, puts that id on every workflow it lets through — ComfyUI tells only
the client that queued a workflow how it goes — and keeps the last word on
each: the node running, its step of how many. It answers at
/progress/{prompt_id}. While a job's request waits, the router reads the
queue (queued, n ahead) and the gate (running, node, step, percent) and
writes it into the job, which a client reads back. Only into a job of the
same key: a tag naming another's job writes nothing.
G3, the rest of the plumbing: comfyui.nodes names custom nodes by https
repository and full commit, with the PyPI packages they need. The gate puts
each in the replica's own custom_nodes at that commit and installs the
packages beside the replica, constrained to the versions ComfyUI's
environment already has, which it never changes. The catalog refuses a
branch for a commit and anything but package specs for requirements. An
OBJ a node saves leaves the gateway as a GLB, colours kept, so every mesh
can carry its licence.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
84 lines
3.5 KiB
Go
84 lines
3.5 KiB
Go
package catalog
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func comfyModel(source string, c *ComfyUI, profiles ...Profile) Model {
|
|
if len(profiles) == 0 {
|
|
profiles = []Profile{{Name: "gpu", Engine: EngineComfyUI, GPUClass: "6000-24g", TP: 1}}
|
|
}
|
|
return Model{Name: "flux", Source: source, ComfyUI: c, Profiles: profiles}
|
|
}
|
|
|
|
func TestAComfyUIModelMakesPicturesFromItsSourcesFile(t *testing.T) {
|
|
cat, err := New([]Model{comfyModel("hf:Comfy-Org/flux1-schnell/flux1-schnell-fp8.safetensors", nil)})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
m, _ := cat.Resolve("flux")
|
|
if !m.MakesImages() || !m.ServesComfyUI() {
|
|
t.Error("a ComfyUI model does not say it makes pictures")
|
|
}
|
|
if w, ckpt := m.ComfyWorkflow(); w != "" || ckpt != "flux1-schnell-fp8.safetensors" {
|
|
t.Errorf("workflow %q, checkpoint %q", w, ckpt)
|
|
}
|
|
}
|
|
|
|
func TestAComfyUIModelIsRefusedWhatItCannotDo(t *testing.T) {
|
|
for name, c := range map[string]struct {
|
|
m Model
|
|
want string
|
|
}{
|
|
"a set without a checkpoint": {comfyModel("https://mirror/flux-dev/v1/", nil), "comfyui.checkpoint"},
|
|
"an unknown workflow": {comfyModel("/w/f.safetensors", &ComfyUI{Workflow: "sdxl"}), "no built-in workflow"},
|
|
"mixed engines": {comfyModel("/w/f.safetensors", nil,
|
|
Profile{Name: "a", Engine: EngineComfyUI, GPUClass: "6000-24g", TP: 1},
|
|
Profile{Name: "b", Engine: "vllm", GPUClass: "6000-24g", TP: 1}), "every profile"},
|
|
"comfyui settings without the engine": {Model{Name: "x", Source: "/w/x", ComfyUI: &ComfyUI{},
|
|
Profiles: []Profile{{Name: "a", Engine: "vllm", GPUClass: "6000-24g", TP: 1}}}, "engine comfyui"},
|
|
"no card": {comfyModel("/w/f.safetensors", nil, Profile{Name: "a", Engine: EngineComfyUI, TP: 1}), "gpu_class"},
|
|
} {
|
|
if _, err := New([]Model{c.m}); err == nil || !strings.Contains(err.Error(), c.want) {
|
|
t.Errorf("%s: %v", name, err)
|
|
}
|
|
}
|
|
ok := comfyModel("https://mirror/flux-dev/v1/", &ComfyUI{Workflow: "flux-dev", Checkpoint: "flux1-dev-fp8.safetensors"})
|
|
if _, err := New([]Model{ok}); err != nil {
|
|
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)
|
|
}
|
|
}
|