From fde1501e8a46f0975cc4f95ed53c743ab10c238f Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Sat, 10 Oct 2026 00:39:10 +0200 Subject: [PATCH] ComfyUI as an engine of the cluster's own MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A model whose profiles run engine comfyui is placed on a node's card like any other, and makes pictures there. ComfyUI asks for no key, so the agent starts it behind llmgw comfyui-gate: ComfyUI on the loopback, the gate on the replica's port demanding the replica's key and letting through only what the router asks — queueing, waiting, fetching, cancelling. The gate answers the agent's /health and /v1/models, shows ComfyUI the fetched weights by their own names, and gives each replica its own base directory. ComfyUI runs as the gate's child, so its card memory is the replica's. A node offers the engine when agent.yaml names its ComfyUI (engines.comfyui.dir). The model's comfyui: block names the workflow and, for a source that is a set of files, the checkpoint. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/admin-api.openapi.json | 18 ++ docs/reference.md | 56 +++++ internal/agent/engines/comfyui.go | 124 ++++++++++ internal/agent/engines/comfyui_test.go | 66 +++++ internal/agent/engines/engine.go | 3 + internal/agent/engines/engines_test.go | 2 +- internal/agent/engines/unsafeenv.go | 13 + internal/cli/comfyuigate.go | 49 ++++ internal/cli/root.go | 1 + internal/comfy/gate.go | 291 +++++++++++++++++++++++ internal/comfy/gate_test.go | 201 ++++++++++++++++ internal/control/catalog/catalog.go | 98 +++++++- internal/control/catalog/comfyui_test.go | 51 ++++ internal/control/control.go | 41 +++- internal/docgen/reference_gen.go | 10 + internal/docgen/render.go | 1 + internal/router/images.go | 3 +- 17 files changed, 1008 insertions(+), 20 deletions(-) create mode 100644 internal/agent/engines/comfyui.go create mode 100644 internal/agent/engines/comfyui_test.go create mode 100644 internal/cli/comfyuigate.go create mode 100644 internal/comfy/gate.go create mode 100644 internal/comfy/gate_test.go create mode 100644 internal/control/catalog/comfyui_test.go diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index a20a546..c3b8f90 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -8576,6 +8576,20 @@ } } }, + "catalog.ComfyUI": { + "type": "object", + "description": "ComfyUI is a ComfyUI model's workflow and checkpoint.", + "properties": { + "workflow": { + "type": "string", + "description": "Workflow is a built-in's name (flux-schnell, the default, or flux-dev) or a JSON file on the control plane, in ComfyUI's API format." + }, + "checkpoint": { + "type": "string", + "description": "Checkpoint is the file the workflow's {{checkpoint}} names: the source's own name unless set, which a source that is a set of files must." + } + } + }, "catalog.EvalGate": { "type": "object", "description": "EvalGate is a quality suite a model's new weights must pass.", @@ -8683,6 +8697,10 @@ "$ref": "#/components/schemas/catalog.External", "description": "External is a model served by someone else — an OpenAI-compatible API outside the cluster — rather than by replicas placed here. Such a model has no profiles and nothing is placed for it; the router sends its requests to the provider. Most useful named in another model's fallbacks: used only when the cluster has no capacity left. Its prompts leave the cluster, so which keys may name it is worth an allowlist." }, + "comfyui": { + "$ref": "#/components/schemas/catalog.ComfyUI", + "description": "ComfyUI is how a model whose profiles run engine comfyui makes its pictures: the workflow, and the file it loads." + }, "adapters": { "type": "array", "items": { diff --git a/docs/reference.md b/docs/reference.md index 2421227..84b3d78 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -61,6 +61,10 @@ rather than repeated. | `vision` | `bool` | [Capabilities](#capabilities) | | `embeddings` | `bool` | [Capabilities](#capabilities) | | `images` | `bool` | [Capabilities](#capabilities) | +| `workflow` | `string` | [ComfyUI](#comfyui) | +| `checkpoint` | `string` | [ComfyUI](#comfyui) | +| `dir` | `string` | [ComfyUIConfig](#comfyuiconfig) | +| `python` | `string` | [ComfyUIConfig](#comfyuiconfig) | | `interval` | `time.Duration` | [Control](#control) | | `base_backoff` | `time.Duration` | [Control](#control) | | `max_backoff` | `time.Duration` | [Control](#control) | @@ -139,6 +143,7 @@ rather than repeated. | `class` | `string` | [Model](#model) | | `owner` | `string` | [Model](#model) | | `external` | `*External` | [Model](#model) | +| `comfyui` | `*ComfyUI` | [Model](#model) | | `adapters` | `[]Adapter` | [Model](#model) | | `max_loaded_adapters` | `int` | [Model](#model) | | `fallbacks` | `[]string` | [Model](#model) | @@ -389,6 +394,7 @@ rather than repeated. | `llamacpp` | `LlamaCppConfig` | [engines.Config](#enginesconfig) | | `trtllm` | `TRTLLMConfig` | [engines.Config](#enginesconfig) | | `sim` | `SimConfig` | [engines.Config](#enginesconfig) | +| `comfyui` | `ComfyUIConfig` | [engines.Config](#enginesconfig) | | `enabled` | `*bool` | [guard.Config](#guardconfig) | | `default` | `[]string` | [guard.Config](#guardconfig) | | `policies` | `map[string]Policy` | [guard.Config](#guardconfig) | @@ -730,6 +736,39 @@ Images is a model that makes pictures, asked for at /v1/images/generations. Implied by an external model whose API is ComfyUI's. +## ComfyUI + +### `workflow` + +`string` + +Workflow is a built-in's name (flux-schnell, the default, or +flux-dev) or a JSON file on the control plane, in ComfyUI's API +format. + +### `checkpoint` + +`string` + +Checkpoint is the file the workflow's {{checkpoint}} names: the +source's own name unless set, which a source that is a set of files +must. + +## ComfyUIConfig + +### `dir` + +`string` + +Dir is ComfyUI's directory, where main.py is. Readable by the user +engines run as. + +### `python` + +`string` + +Python is the interpreter of its environment; Dir's .venv when unset. + ## Control ### `interval` @@ -1344,6 +1383,15 @@ allowlist. Takes a [External](#external) block. +### `comfyui` + +`*ComfyUI` + +ComfyUI is how a model whose profiles run engine comfyui makes its +pictures: the workflow, and the file it loads. + +Takes a [ComfyUI](#comfyui) block. + ### `adapters` `[]Adapter` @@ -3238,6 +3286,14 @@ Takes a [TRTLLMConfig](#trtllmconfig) block. Takes a [SimConfig](#simconfig) block. +### `comfyui` + +`ComfyUIConfig` + +ComfyUI makes pictures; offered by a node whose ComfyUI is named here. + +Takes a [ComfyUIConfig](#comfyuiconfig) block. + ## guard.Config ### `enabled` diff --git a/internal/agent/engines/comfyui.go b/internal/agent/engines/comfyui.go new file mode 100644 index 0000000..c0b04b7 --- /dev/null +++ b/internal/agent/engines/comfyui.go @@ -0,0 +1,124 @@ +package engines + +import ( + "context" + "fmt" + "net" + "os" + "path/filepath" + "regexp" + "strconv" + "sync/atomic" +) + +// ComfyUIConfig locates a ComfyUI installation. +type ComfyUIConfig struct { + // Dir is ComfyUI's directory, where main.py is. Readable by the user + // engines run as. + Dir string `koanf:"dir"` + // Python is the interpreter of its environment; Dir's .venv when unset. + Python string `koanf:"python"` +} + +// ComfyUI runs ComfyUI as an engine, behind llmgw's gate (llmgw +// comfyui-gate): ComfyUI asks for no key, so it listens on the loopback and +// the gate, which does, on the replica's port. Its replicas make pictures +// for /v1/images/generations. +type ComfyUI struct { + cfg ComfyUIConfig + installed atomic.Pointer[string] +} + +// Name identifies the engine in the catalog. +func (*ComfyUI) Name() string { return "comfyui" } + +// HealthPath is the gate's, which answers once ComfyUI does. +func (*ComfyUI) HealthPath() string { return "/health" } + +// MetricsPath is empty: ComfyUI publishes no metrics, and the router reads +// none from it. +func (*ComfyUI) MetricsPath() string { return "" } + +// ServesHTTP is true: a ComfyUI replica is one process on one node. +func (*ComfyUI) ServesHTTP(ShardSpec) bool { return true } + +// UseInstalled points the engine at a ComfyUI directory the agent installed. +func (c *ComfyUI) UseInstalled(dir string) { c.installed.Store(&dir) } + +// Configured reports whether the configuration names a ComfyUI. +func (c *ComfyUI) Configured() bool { return c.cfg.Dir != "" } + +func (c *ComfyUI) dir() string { + if p := c.installed.Load(); p != nil { + return *p + } + return c.cfg.Dir +} + +func (c *ComfyUI) python() string { + if c.cfg.Python != "" && c.installed.Load() == nil { + return c.cfg.Python + } + return filepath.Join(c.dir(), ".venv", "bin", "python") +} + +var comfyVersion = regexp.MustCompile(`__version__\s*=\s*"([^"]+)"`) + +// Probe reports ComfyUI's version, from the file its build writes it to: +// importing ComfyUI to ask would load torch and look for cards. +func (c *ComfyUI) Probe(context.Context) (string, error) { + dir := c.dir() + if dir == "" { + return "", fmt.Errorf("no ComfyUI directory is configured") + } + if _, err := os.Stat(filepath.Join(dir, "main.py")); err != nil { + return "", fmt.Errorf("%s is not ComfyUI: %w", dir, err) + } + b, err := os.ReadFile(filepath.Join(dir, "comfyui_version.py")) + if err != nil { + return Unknown, nil //nolint:nilerr // installed, of a version it does not say + } + if m := comfyVersion.FindSubmatch(b); m != nil { + return string(m[1]), nil + } + return Unknown, nil +} + +// Validate checks the profile: one card, on one node. ComfyUI runs a +// workflow on one device and has no tensor parallelism to spread it by. +func (*ComfyUI) Validate(p Profile, gpus []GPUInfo) error { + if p.Nodes > 1 || p.TP > 1 || p.PP > 1 || p.GPUsPerNode > 1 { + return fmt.Errorf("profile %q: ComfyUI runs on one GPU of one node", p.Name) + } + return validateCommon(p, gpus) +} + +// Commands is the gate, which starts ComfyUI itself. +func (c *ComfyUI) Commands(s ShardSpec) ([]Command, error) { + dir := c.dir() + if dir == "" { + return nil, fmt.Errorf("comfyui: no ComfyUI directory is configured (engines.comfyui.dir)") + } + bin, err := os.Executable() + if err != nil { + return nil, err + } + p := s.Profile + args := []string{"comfyui-gate", + "--listen", net.JoinHostPort(s.Host(), strconv.Itoa(s.HTTPPort)), + "--model", p.ServedModelName, + "--weights", p.ModelPath, + // Relative to where engines run — their user's home — so it is + // theirs to write. + "--state", filepath.Join("comfyui", s.ReplicaID), + "--python", c.python(), + "--comfyui", dir, + "--", + } + args = append(args, p.Args...) + env := baseEnv(s) + if s.APIKey != "" { + env["LLMGW_ENGINE_KEY"] = s.APIKey + } + return []Command{{Name: "comfyui", Path: bin, Args: args, Env: materializeEnv(env)}}, nil +} diff --git a/internal/agent/engines/comfyui_test.go b/internal/agent/engines/comfyui_test.go new file mode 100644 index 0000000..eb54143 --- /dev/null +++ b/internal/agent/engines/comfyui_test.go @@ -0,0 +1,66 @@ +package engines + +import ( + "os" + "path/filepath" + "slices" + "strings" + "testing" +) + +func TestComfyUIRunsBehindTheGateWithItsKeyOutOfSight(t *testing.T) { + c := &ComfyUI{cfg: ComfyUIConfig{Dir: "/opt/comfyui"}} + cmds, err := c.Commands(ShardSpec{ + ReplicaID: "r-flux-1", HTTPPort: 8100, BindAddr: "10.0.0.147", APIKey: "sekret", GPUIndices: []int{0}, + Profile: Profile{Name: "p", Engine: "comfyui", ModelPath: "/w/flux.safetensors", ServedModelName: "flux", + Args: []string{"--disable-smart-memory"}}, + }) + if err != nil || len(cmds) != 1 { + t.Fatalf("%v %v", cmds, err) + } + args := strings.Join(cmds[0].Args, " ") + for _, want := range []string{"comfyui-gate", "--listen 10.0.0.147:8100", "--model flux", "--weights /w/flux.safetensors", + "--state comfyui/r-flux-1", "--python /opt/comfyui/.venv/bin/python", "--comfyui /opt/comfyui", "-- --disable-smart-memory"} { + if !strings.Contains(args, want) { + t.Errorf("%q not in %q", want, args) + } + } + if strings.Contains(args, "sekret") { + t.Error("the key is on the command line, where every user can read it") + } + if !slices.Contains(cmds[0].Env, "LLMGW_ENGINE_KEY=sekret") || !slices.Contains(cmds[0].Env, "CUDA_VISIBLE_DEVICES=0") { + t.Error("the key or the card is not in the environment") + } +} + +func TestComfyUIRunsOnOneCard(t *testing.T) { + c := &ComfyUI{} + if err := c.Validate(Profile{Name: "p", ModelPath: "/w", TP: 2}, h100(2)); err == nil { + t.Error("a two-card ComfyUI profile was accepted") + } + if err := c.Validate(Profile{Name: "p", ModelPath: "/w", TP: 1}, h100(1)); err != nil { + t.Error(err) + } + if _, err := c.Commands(ShardSpec{Profile: Profile{ModelPath: "/w"}}); err == nil { + t.Error("commands without a ComfyUI directory") + } +} + +func TestComfyUIIsOfferedByTheVersionItsBuildWrote(t *testing.T) { + dir := t.TempDir() + c := &ComfyUI{cfg: ComfyUIConfig{Dir: dir}} + if _, err := c.Probe(t.Context()); err == nil { + t.Error("a directory without main.py passed for ComfyUI") + } + _ = os.WriteFile(filepath.Join(dir, "main.py"), nil, 0o644) + _ = os.WriteFile(filepath.Join(dir, "comfyui_version.py"), []byte("__version__ = \"0.39.2\"\n"), 0o644) + if v, err := c.Probe(t.Context()); err != nil || v != "0.39.2" { + t.Errorf("%q %v", v, err) + } + if got := NewRegistry(Config{ComfyUI: ComfyUIConfig{Dir: dir}}).ProbeConfigured(t.Context()); got["comfyui"] != "0.39.2" { + t.Errorf("a configured ComfyUI is reported as %v", got) + } + if err := UnsafeArgs("comfyui", []string{"--enable-manager"}); err == nil { + t.Error("a profile may install ComfyUI nodes") + } +} diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index 9b9d0fc..35952a8 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -214,6 +214,7 @@ func NewRegistry(cfg Config) *Registry { r.Add(&LlamaCpp{cfg: cfg.LlamaCpp}) r.Add(&TRTLLM{cfg: cfg.TRTLLM}) r.Add(&Sim{cfg: cfg.Sim}) + r.Add(&ComfyUI{cfg: cfg.ComfyUI}) return r } @@ -262,6 +263,8 @@ type Config struct { // offers it only when its interpreter is named here. TRTLLM TRTLLMConfig `koanf:"trtllm"` Sim SimConfig `koanf:"sim"` + // ComfyUI makes pictures; offered by a node whose ComfyUI is named here. + ComfyUI ComfyUIConfig `koanf:"comfyui"` } // --- shared helpers --- diff --git a/internal/agent/engines/engines_test.go b/internal/agent/engines/engines_test.go index 7d0fb9d..2f58000 100644 --- a/internal/agent/engines/engines_test.go +++ b/internal/agent/engines/engines_test.go @@ -477,7 +477,7 @@ func TestValidateRequiresModelPath(t *testing.T) { func TestRegistry(t *testing.T) { r := NewRegistry(Config{}) - if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "trtllm", "vllm"}) { + if got := r.Names(); !slices.Equal(got, []string{"comfyui", "llamacpp", "sglang", "sim", "trtllm", "vllm"}) { t.Errorf("registry contents: %v", got) } if _, err := r.Get("vllm"); err != nil { diff --git a/internal/agent/engines/unsafeenv.go b/internal/agent/engines/unsafeenv.go index 26c8cc8..c1d267b 100644 --- a/internal/agent/engines/unsafeenv.go +++ b/internal/agent/engines/unsafeenv.go @@ -119,6 +119,19 @@ var flagsNever = map[string]map[string]string{ "--model-path": "is the model's source, set by the agent", "--served-model-name": "is set by the agent", "--lora-paths": "loads files", "--download-dir": "writes files", }, + // ComfyUI's main.py, behind the gate that sets where it listens and + // where its files are. + "comfyui": { + "--listen": "is set by the agent", "--base-directory": "is set by the agent", + "--extra-model-paths-config": "is set by the agent", + "--output-directory": "writes files", "--temp-directory": "writes files", + "--input-directory": "reads files", "--user-directory": "writes files", + "--database-url": "writes a database elsewhere", "--front-end-root": "serves a directory", + "--enable-cors-header": "lets other sites' pages call it", + "--enable-manager": "installs nodes, which run Python", + "--tls-keyfile": "reads a file", "--tls-certfile": "reads a file", + "--comfy-api-base": "sends work to another service", + }, // trtllm-serve spells its flags with underscores, so the shared list's // hyphenated names do not reach them. "trtllm": { diff --git a/internal/cli/comfyuigate.go b/internal/cli/comfyuigate.go new file mode 100644 index 0000000..817c7c5 --- /dev/null +++ b/internal/cli/comfyuigate.go @@ -0,0 +1,49 @@ +package cli + +import ( + "os" + "os/signal" + "syscall" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" +) + +// comfyKeyEnv carries the replica's key to the gate: the environment, which +// only the engine's user and root can read, rather than the command line, +// which every user on the node can. +const comfyKeyEnv = "LLMGW_ENGINE_KEY" + +func newComfyUIGateCmd() *cobra.Command { + var cfg comfy.GateConfig + cmd := &cobra.Command{ + Use: "comfyui-gate [flags] [-- main.py flags]", + Short: "Run ComfyUI as an engine: behind a key, on the loopback (started by the agent)", + Hidden: true, + Long: "Started by the node agent for a replica whose engine is comfyui. ComfyUI asks\n" + + "nobody for a key, so it listens on the loopback only and this gate on the\n" + + "replica's port, letting through, with the replica's key, only what the router\n" + + "asks of it. It also lays the fetched weights out where ComfyUI looks for them.", + RunE: func(cmd *cobra.Command, args []string) error { + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + cfg.Key = os.Getenv(comfyKeyEnv) + _ = os.Unsetenv(comfyKeyEnv) + cfg.Args = args + cfg.Log = logger() + return comfy.RunGate(ctx, cfg) + }, + } + f := cmd.Flags() + f.StringVar(&cfg.Listen, "listen", "", "the replica's address") + f.StringVar(&cfg.Model, "model", "", "the model's name, for /v1/models") + f.StringVar(&cfg.Weights, "weights", "", "the fetched weights: a file or a directory") + f.StringVar(&cfg.State, "state", "", "the replica's own directory") + f.StringVar(&cfg.Python, "python", "python3", "the interpreter of ComfyUI's environment") + f.StringVar(&cfg.Dir, "comfyui", "", "ComfyUI's directory, where main.py is") + for _, n := range []string{"listen", "weights", "state", "comfyui"} { + _ = cmd.MarkFlagRequired(n) + } + return cmd +} diff --git a/internal/cli/root.go b/internal/cli/root.go index 7f08832..48c4ee7 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -54,6 +54,7 @@ func Root() *cobra.Command { newMirrorCmd(), newBundleCmd(), newSimEngineCmd(), + newComfyUIGateCmd(), newDoctorCmd(), newNodeCmd(), newSelfUpgradeCmd(), diff --git a/internal/comfy/gate.go b/internal/comfy/gate.go new file mode 100644 index 0000000..a602e68 --- /dev/null +++ b/internal/comfy/gate.go @@ -0,0 +1,291 @@ +package comfy + +import ( + "context" + "crypto/subtle" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/http/httputil" + "net/url" + "os" + "os/exec" + "path/filepath" + "regexp" + "strconv" + "strings" + "syscall" + "time" +) + +// The gate is how a node runs ComfyUI as one of llmgw's engines. +// +// ComfyUI asks nobody for a key: whoever reaches its port can queue work, +// upload files and install nodes. Every other engine here demands the key +// the control plane made for its replica, so ComfyUI is put behind a gate +// that does: ComfyUI listens on the loopback only, and the gate on the +// replica's port, letting through only the requests the router makes. The +// gate also answers what the agent asks every engine — /health once +// ComfyUI is up, and /v1/models naming the model — and lays the weights out +// where ComfyUI looks for them. +// +// ComfyUI runs as the gate's child, in its process group, so stopping the +// replica stops both, the card memory ComfyUI holds is counted as the +// replica's, and either one exiting ends the other. + +// GateConfig is one replica's gate. +type GateConfig struct { + // Listen is the replica's address, the one the router is given. + Listen string + // Model is the name the agent asks /v1/models for. + Model string + // Key is what every request must carry as a bearer token; empty lets + // every request through, as an engine started without one does. + Key string + // Weights is the model as the agent fetched it: a file, or a directory + // of files. + Weights string + // State is the replica's own directory: ComfyUI's base directory, with + // its output, temp and user directories, and the models it is shown. + State string + // Python runs ComfyUI's main.py, found in Dir. + Python, Dir string + // Args are passed on to main.py, after the gate's own. + Args []string + // Output is where ComfyUI's output goes. + Output io.Writer + Log *slog.Logger +} + +// The requests the router makes of ComfyUI: queueing, waiting, fetching and +// cancelling. Nothing that uploads, installs or changes ComfyUI. +var gatePaths = []string{"/prompt", "/history/", "/view", "/queue", "/interrupt", "/system_stats", "/object_info"} + +// modelFolders are the kinds of file ComfyUI's loaders look for. Each is +// shown the same directory: a workflow names its files, so a file being +// listed where it is not wanted costs nothing. +var modelFolders = []string{"checkpoints", "diffusion_models", "text_encoders", "vae", "loras", + "clip_vision", "controlnet", "upscale_models", "embeddings", "style_models"} + +// RunGate lays out the replica, starts ComfyUI and serves the gate until +// ctx ends or ComfyUI exits. +func RunGate(ctx context.Context, cfg GateConfig) error { + log := cfg.Log + if log == nil { + log = slog.Default() + } + state, err := filepath.Abs(cfg.State) + if err != nil { + return err + } + models := filepath.Join(state, "models") + if err := LinkWeights(cfg.Weights, models); err != nil { + return 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 + } + + inner, err := freePort() + if err != nil { + return err + } + args := append([]string{filepath.Join(cfg.Dir, "main.py"), + "--listen", "127.0.0.1", "--port", strconv.Itoa(inner), + "--base-directory", state, "--extra-model-paths-config", extra, + "--disable-auto-launch"}, cfg.Args...) + 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 + } + child.Stdout, child.Stderr = out, out + if err := child.Start(); err != nil { + return fmt.Errorf("starting ComfyUI: %w", err) + } + exited := make(chan error, 1) + go func() { exited <- child.Wait() }() + + upstream := &url.URL{Scheme: "http", Host: net.JoinHostPort("127.0.0.1", strconv.Itoa(inner))} + srv := &http.Server{ + Addr: cfg.Listen, + Handler: NewGate(upstream, cfg.Key, cfg.Model), + ReadHeaderTimeout: 10 * time.Second, + } + served := make(chan error, 1) + go func() { served <- srv.ListenAndServe() }() + log.Info("ComfyUI gate up", "listen", cfg.Listen, "comfyui", upstream.Host, "weights", cfg.Weights) + + stop := func() { + _ = child.Process.Signal(syscall.SIGTERM) + select { + case <-exited: + case <-time.After(20 * time.Second): + _ = child.Process.Kill() + <-exited + } + } + select { + case <-ctx.Done(): + shutdown(srv) + stop() + return nil + case err := <-exited: + shutdown(srv) + if err == nil { + err = errors.New("exited") + } + return fmt.Errorf("ComfyUI %w", err) + case err := <-served: + stop() + return fmt.Errorf("the gate stopped serving: %w", err) + } +} + +func shutdown(srv *http.Server) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = srv.Shutdown(ctx) +} + +// NewGate is the gate's handler, in front of the ComfyUI at upstream. +func NewGate(upstream *url.URL, key, model string) http.Handler { + proxy := httputil.NewSingleHostReverseProxy(upstream) + base := proxy.Director + proxy.Director = func(r *http.Request) { + base(r) + r.Header.Del("Authorization") + // ComfyUI refuses a request whose Origin is not its own host; the + // router sends none, and nothing else gets this far. + r.Header.Del("Origin") + r.Host = upstream.Host + } + client := &http.Client{Timeout: 3 * time.Second} + up := func(ctx context.Context) error { + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, upstream.String()+"/system_stats", nil) + resp, err := client.Do(req) + if err != nil { + return err + } + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("ComfyUI answered %s", resp.Status) + } + return nil + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if key != "" { + got, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ") + if !ok || subtle.ConstantTimeCompare([]byte(got), []byte(key)) != 1 { + http.Error(w, "this engine wants its key", http.StatusUnauthorized) + return + } + } + switch r.URL.Path { + case "/health": + if err := up(r.Context()); err != nil { + http.Error(w, "ComfyUI is not up yet: "+err.Error(), http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + return + case "/v1/models": + if err := up(r.Context()); err != nil { + http.Error(w, "ComfyUI is not up yet: "+err.Error(), http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"object":"list","data":[{"id":%q,"object":"model","owned_by":"comfyui"}]}`, model) + return + } + for _, p := range gatePaths { + if r.URL.Path == p || strings.HasSuffix(p, "/") && strings.HasPrefix(r.URL.Path, p) { + proxy.ServeHTTP(w, r) + return + } + } + http.NotFound(w, r) + }) +} + +// cachedName is a file the agent fetched by itself: its cache names it by +// the source's hash, then the name it had. +var cachedName = regexp.MustCompile(`^[0-9a-f]{16}-`) + +// FileName is what a weights file is called where ComfyUI sees it: its own +// name, without the prefix the agent's cache gives it. +func FileName(path string) string { + return cachedName.ReplaceAllString(filepath.Base(path), "") +} + +// LinkWeights shows ComfyUI the weights: a link in dir to the file, or to +// every file of the directory. Links made before are taken away first, so +// a replica started again on other weights sees only those. +func LinkWeights(weights, dir string) error { + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + old, err := os.ReadDir(dir) + if err != nil { + return err + } + for _, e := range old { + if e.Type()&os.ModeSymlink != 0 { + _ = os.Remove(filepath.Join(dir, e.Name())) + } + } + weights, err = filepath.Abs(weights) + if err != nil { + return err + } + info, err := os.Stat(weights) + if err != nil { + return fmt.Errorf("the weights: %w", 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.Type().IsRegular() && !strings.HasPrefix(e.Name(), ".") { + files = append(files, filepath.Join(weights, e.Name())) + } + } + if len(files) == 0 { + return fmt.Errorf("the weights directory %s has no files", weights) + } + } + for _, f := range files { + if err := os.Symlink(f, filepath.Join(dir, FileName(f))); err != nil { + return err + } + } + return nil +} + +func extraModelPaths(models string) string { + var b strings.Builder + fmt.Fprintf(&b, "llmgw:\n base_path: %q\n is_default: true\n", models) + for _, f := range modelFolders { + fmt.Fprintf(&b, " %s: .\n", f) + } + return b.String() +} + +func freePort() (int, error) { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return 0, err + } + defer l.Close() + return l.Addr().(*net.TCPAddr).Port, nil +} diff --git a/internal/comfy/gate_test.go b/internal/comfy/gate_test.go new file mode 100644 index 0000000..679d523 --- /dev/null +++ b/internal/comfy/gate_test.go @@ -0,0 +1,201 @@ +package comfy + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestTheGateWantsTheKeyAndLetsThroughOnlyTheRoutersRequests(t *testing.T) { + var got []string + comfy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = append(got, r.Method+" "+r.URL.Path+" auth="+r.Header.Get("Authorization")) + w.WriteHeader(http.StatusOK) + })) + defer comfy.Close() + u, _ := url.Parse(comfy.URL) + gate := httptest.NewServer(NewGate(u, "k", "flux")) + defer gate.Close() + + do := func(method, path, key string) (int, string) { + req, _ := http.NewRequest(method, gate.URL+path, strings.NewReader("{}")) + if key != "" { + req.Header.Set("Authorization", "Bearer "+key) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(b) + } + if code, _ := do("POST", "/prompt", ""); code != http.StatusUnauthorized { + t.Errorf("no key: %d", code) + } + if code, _ := do("POST", "/prompt", "wrong"); code != http.StatusUnauthorized { + t.Errorf("wrong key: %d", code) + } + for _, p := range []string{"/upload/image", "/api/manager/queue/install", "/userdata/x", "/internal/logs"} { + if code, _ := do("POST", p, "k"); code != http.StatusNotFound { + t.Errorf("%s: %d, want it refused", p, code) + } + } + for _, p := range []string{"/prompt", "/history/p1", "/view?filename=a.png", "/queue", "/interrupt"} { + if code, _ := do("POST", p, "k"); code != http.StatusOK { + t.Errorf("%s: %d", p, code) + } + } + if code, body := do("GET", "/v1/models", "k"); code != 200 || !strings.Contains(body, `"id":"flux"`) { + t.Errorf("/v1/models: %d %s", code, body) + } + if code, _ := do("GET", "/health", "k"); code != 200 { + t.Errorf("/health: %d", code) + } + for _, g := range got { + if strings.Contains(g, "auth=Bearer") { + t.Errorf("the key reached ComfyUI: %s", g) + } + } + comfy.Close() + if code, _ := do("GET", "/health", "k"); code != http.StatusServiceUnavailable { + t.Errorf("/health with ComfyUI down: %d", code) + } +} + +func TestTheWeightsAreShownByTheirOwnNames(t *testing.T) { + cache := t.TempDir() + one := filepath.Join(cache, "0123456789abcdef-flux1-schnell-fp8.safetensors") + _ = os.WriteFile(one, []byte("w"), 0o644) + models := filepath.Join(t.TempDir(), "models") + if err := LinkWeights(one, models); err != nil { + t.Fatal(err) + } + if target, err := os.Readlink(filepath.Join(models, "flux1-schnell-fp8.safetensors")); err != nil || target != one { + t.Fatalf("%q %v", target, err) + } + + set := t.TempDir() + for _, n := range []string{"flux1-dev.safetensors", "ae.safetensors", "t5xxl_fp8.safetensors", ".complete"} { + _ = os.WriteFile(filepath.Join(set, n), []byte("w"), 0o644) + } + if err := LinkWeights(set, models); err != nil { + t.Fatal(err) + } + entries, _ := os.ReadDir(models) + var names []string + for _, e := range entries { + names = append(names, e.Name()) + } + if strings.Join(names, ",") != "ae.safetensors,flux1-dev.safetensors,t5xxl_fp8.safetensors" { + t.Errorf("shown %v: the old weights' link should be gone and the marker not shown", names) + } + if !strings.Contains(extraModelPaths(models), "checkpoints: .") { + t.Error(extraModelPaths(models)) + } +} + +// TestMain lets the test binary stand in for ComfyUI: started by the gate +// as `python main.py --listen … --port …`, it answers /system_stats and +// says what it was given. +func TestMain(m *testing.M) { + if os.Getenv("LLMGW_FAKE_COMFYUI") == "1" { + fakeComfyUI() + return + } + os.Exit(m.Run()) +} + +func fakeComfyUI() { + var port, base string + for i, a := range os.Args { + switch a { + case "--port": + port = os.Args[i+1] + case "--base-directory": + base = os.Args[i+1] + } + } + _ = os.WriteFile(filepath.Join(base, "args"), []byte(strings.Join(os.Args[1:], " ")+"\nkey="+os.Getenv("LLMGW_ENGINE_KEY")), 0o644) + http.HandleFunc("/system_stats", func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, "{}") }) + http.HandleFunc("/quit", func(w http.ResponseWriter, _ *http.Request) { os.Exit(3) }) + if os.Getenv("LLMGW_FAKE_COMFYUI_EXIT") == "1" { + time.AfterFunc(200*time.Millisecond, func() { os.Exit(3) }) + } + _ = http.ListenAndServe("127.0.0.1:"+port, nil) +} + +func TestTheGateRunsComfyUIAndEndsWithIt(t *testing.T) { + t.Setenv("LLMGW_FAKE_COMFYUI", "1") + weights := filepath.Join(t.TempDir(), "flux.safetensors") + _ = os.WriteFile(weights, []byte("w"), 0o644) + state := t.TempDir() + l, _ := net.Listen("tcp", "127.0.0.1:0") + listen := l.Addr().String() + l.Close() + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { + done <- RunGate(ctx, GateConfig{Listen: listen, Model: "flux", Weights: weights, State: state, + Python: os.Args[0], Dir: t.TempDir(), Args: []string{"--disable-smart-memory"}, Output: io.Discard}) + }() + get := func(path string) int { + resp, err := http.Get("http://" + listen + path) + if err != nil { + return 0 + } + resp.Body.Close() + return resp.StatusCode + } + deadline := time.Now().Add(10 * time.Second) + for get("/health") != http.StatusOK { + if time.Now().After(deadline) { + t.Fatal("the gate never said ComfyUI was up") + } + time.Sleep(20 * time.Millisecond) + } + b, _ := os.ReadFile(filepath.Join(state, "args")) + for _, want := range []string{"main.py --listen 127.0.0.1 --port", "--base-directory " + state, + "--extra-model-paths-config " + filepath.Join(state, "extra_model_paths.yaml"), "--disable-smart-memory"} { + if !strings.Contains(string(b), want) { + t.Errorf("%q not in what ComfyUI was given: %s", want, b) + } + } + if _, err := os.Lstat(filepath.Join(state, "models", "flux.safetensors")); err != nil { + t.Error("the weights were not shown to ComfyUI:", err) + } + if get("/quit") != http.StatusNotFound { + t.Error("a path the router never asks for was let through") + } + cancel() + select { + case err := <-done: + if err != nil { + t.Errorf("stopped, the gate said %v", err) + } + case <-time.After(25 * time.Second): + t.Fatal("the gate did not stop") + } +} + +func TestComfyUIExitingEndsTheGate(t *testing.T) { + t.Setenv("LLMGW_FAKE_COMFYUI", "1") + t.Setenv("LLMGW_FAKE_COMFYUI_EXIT", "1") + weights := filepath.Join(t.TempDir(), "w.safetensors") + _ = os.WriteFile(weights, []byte("w"), 0o644) + err := RunGate(t.Context(), GateConfig{Listen: "127.0.0.1:0", Weights: weights, State: t.TempDir(), + Python: os.Args[0], Dir: t.TempDir(), Output: io.Discard}) + if err == nil || !strings.Contains(err.Error(), "exit status 3") { + t.Fatalf("err = %v", err) + } +} diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 0f52fed..adf3c15 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -8,6 +8,7 @@ import ( "encoding/json" "fmt" "net/url" + "path" "slices" "strings" "time" @@ -322,6 +323,9 @@ type Model struct { // prompts leave the cluster, so which keys may name it is worth an // allowlist. External *External `json:"external,omitempty" koanf:"external"` + // ComfyUI is how a model whose profiles run engine comfyui makes its + // pictures: the workflow, and the file it loads. + ComfyUI *ComfyUI `json:"comfyui,omitempty" koanf:"comfyui"` // Adapters are LoRA adapters served on this model's replicas, each asked // for by its own name as though it were a model. Every replica loads // every adapter, applying one only to a request that names it — unless @@ -466,7 +470,90 @@ type Capabilities struct { // 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 + return m.Capabilities.Images || m.ServesComfyUI() +} + +// EngineComfyUI is the engine that runs ComfyUI on a node. +const EngineComfyUI = "comfyui" + +// ServesComfyUI reports whether a model is made by ComfyUI, on the +// cluster's own nodes or on a server outside it. +func (m Model) ServesComfyUI() bool { + if m.External != nil { + return m.External.API == APIComfyUI + } + return len(m.Profiles) > 0 && m.Profiles[0].Engine == EngineComfyUI +} + +// ComfyUI is a ComfyUI model's workflow and checkpoint. +type ComfyUI struct { + // Workflow is a built-in's name (flux-schnell, the default, or + // flux-dev) or a JSON file on the control plane, in ComfyUI's API + // format. + Workflow string `json:"workflow,omitempty" koanf:"workflow"` + // Checkpoint is the file the workflow's {{checkpoint}} names: the + // source's own name unless set, which a source that is a set of files + // must. + Checkpoint string `json:"checkpoint,omitempty" koanf:"checkpoint"` +} + +// ComfyWorkflow is what a ComfyUI model's replicas are sent: the workflow +// as the catalog names it, and the checkpoint it loads. +func (m Model) ComfyWorkflow() (workflow, checkpoint string) { + if m.External != nil { + return m.External.Workflow, m.External.Model + } + if m.ComfyUI != nil { + workflow, checkpoint = m.ComfyUI.Workflow, m.ComfyUI.Checkpoint + } + if checkpoint == "" && !strings.HasSuffix(m.Source, "/") { + checkpoint = path.Base(strings.TrimPrefix(m.Source, "hf:")) + } + return workflow, checkpoint +} + +// checkWorkflow refuses a workflow that is neither a built-in nor a file. +func checkWorkflow(model, w string) error { + // A file is read where the router runs; a name is checked here. + if 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(), ", ")) + } + return nil +} + +func (m Model) validateComfyUI() error { + comfyProfiles := 0 + for _, p := range m.Profiles { + if p.Engine == EngineComfyUI { + comfyProfiles++ + } + } + if comfyProfiles == 0 { + if m.ComfyUI != nil { + return fmt.Errorf("model %q: comfyui: is for a model whose profiles run engine comfyui", m.Name) + } + return nil + } + if comfyProfiles != len(m.Profiles) { + return fmt.Errorf("model %q: every profile runs engine comfyui or none does: "+ + "a request for pictures cannot be answered by an engine that writes text", m.Name) + } + if len(m.Adapters) > 0 { + return fmt.Errorf("model %q: ComfyUI serves no adapters", m.Name) + } + var w string + if m.ComfyUI != nil { + w = m.ComfyUI.Workflow + } + if err := checkWorkflow(m.Name, w); err != nil { + return err + } + if _, ckpt := m.ComfyWorkflow(); ckpt == "" { + return fmt.Errorf("model %q: its source is a set of files, so comfyui.checkpoint must name the one the workflow loads", + m.Name) + } + return nil } // Catalog is the set of models the cluster should serve. @@ -707,6 +794,9 @@ func (m Model) validate() error { } else if len(m.Profiles) == 0 { return fmt.Errorf("model %q: at least one profile is required", m.Name) } + if err := m.validateComfyUI(); err != nil { + return err + } if m.Owner != "" && !validOwner(m.Owner) { return fmt.Errorf("model %q: owner %q is not org:ID, team:ID or project:ID", m.Name, m.Owner) } @@ -1148,10 +1238,8 @@ func (e *External) validate(model string, profiles int) error { 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(), ", ")) + if err := checkWorkflow(model, e.Workflow); err != nil { + return err } default: return fmt.Errorf("model %q: external api %q is not openai or comfyui", model, e.API) diff --git a/internal/control/catalog/comfyui_test.go b/internal/control/catalog/comfyui_test.go new file mode 100644 index 0000000..dfde589 --- /dev/null +++ b/internal/control/catalog/comfyui_test.go @@ -0,0 +1,51 @@ +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) + } +} diff --git a/internal/control/control.go b/internal/control/control.go index 2d8260c..ff9b12f 100644 --- a/internal/control/control.go +++ b/internal/control/control.go @@ -361,7 +361,7 @@ func (c *Controller) endpoints() []endpoint.Info { names := append([]string{m.Name}, m.Aliases...) prof, _ := m.ProfileByName(rep.Profile) - out = append(out, endpoint.Info{ + info := endpoint.Info{ ID: rep.ID, URL: strings.TrimRight(rep.Endpoint, "/"), Models: names, @@ -370,7 +370,18 @@ func (c *Controller) endpoints() []endpoint.Info { Pool: rep.Pool, Weight: rep.Weight, Priority: prof.Priority, - }) + } + if prof.Engine == catalog.EngineComfyUI { + workflow, ckpt := m.ComfyWorkflow() + w, err := c.workflow(workflow) + if err != nil { + c.log.Warn("a ComfyUI model's workflow could not be read; its replica is not offered", + "model", m.Name, "replica", rep.ID, "err", err) + continue + } + info.Workflow, info.UpstreamModel = w, ckpt + } + out = append(out, info) } return append(out, c.externalEndpoints()...) } @@ -415,17 +426,12 @@ func (c *Controller) externalEndpoints() []endpoint.Info { } 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) + w, err := c.workflow(m.External.Workflow) + if err != nil { + c.log.Warn("a ComfyUI model's workflow could not be read; not offered", "model", m.Name, "err", err) + continue } + engine, workflow = endpoint.EngineComfyUI, w } out = append(out, endpoint.Info{ ID: id, URL: strings.TrimRight(m.External.URL, "/"), @@ -439,6 +445,17 @@ func (c *Controller) externalEndpoints() []endpoint.Info { return out } +// workflow is a ComfyUI workflow as the router is sent it: a built-in's +// name as it is, a file read here, on the control plane the catalog says it +// is on, and sent as the workflow itself. +func (c *Controller) workflow(w string) (string, error) { + if !strings.ContainsAny(w, "/.") { + return w, nil + } + b, err := os.ReadFile(w) + return string(b), err +} + // fingerprint summarizes a list so an unchanged one can be skipped. func fingerprint(infos []endpoint.Info) string { var b strings.Builder diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 73d0a96..ec728e3 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -66,6 +66,14 @@ var Reference = []Block{ {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: "catalog", Name: "ComfyUI", Title: "ComfyUI", Doc: "", Fields: []Field{ + {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: ""}, + }}, + {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: ""}, + {Key: "python", Type: "string", Doc: "Python is the interpreter of its environment; Dir's .venv when unset.", 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: ""}, {Key: "base_backoff", Type: "time.Duration", Doc: "BaseBackoff and MaxBackoff bound how fast a failing model is retried.", Nested: ""}, @@ -175,6 +183,7 @@ var Reference = []Block{ {Key: "class", Type: "string", Doc: "Class is how much this model's capacity is worth against another's,\nwhen there is not enough for both: critical, standard or opportunistic.\nEmpty is standard, or opportunistic for a model loaded on demand. See\nEffectiveClass.", Nested: ""}, {Key: "owner", Type: "string", Doc: "Owner is the part of the tenancy tree whose capacity quota this model\ncounts against, written org:ID, team:ID or project:ID. Empty is shared:\nheld by the cluster, counted against nobody.", Nested: ""}, {Key: "external", Type: "*External", Doc: "External is a model served by someone else — an OpenAI-compatible API\noutside the cluster — rather than by replicas placed here. Such a model\nhas no profiles and nothing is placed for it; the router sends its\nrequests to the provider. Most useful named in another model's\nfallbacks: used only when the cluster has no capacity left. Its\nprompts leave the cluster, so which keys may name it is worth an\nallowlist.", Nested: "catalog.External"}, + {Key: "comfyui", Type: "*ComfyUI", Doc: "ComfyUI is how a model whose profiles run engine comfyui makes its\npictures: the workflow, and the file it loads.", Nested: "catalog.ComfyUI"}, {Key: "adapters", Type: "[]Adapter", Doc: "Adapters are LoRA adapters served on this model's replicas, each asked\nfor by its own name as though it were a model. Every replica loads\nevery adapter, applying one only to a request that names it — unless\nMaxLoadedAdapters is set.", Nested: ""}, {Key: "max_loaded_adapters", Type: "int", Doc: "MaxLoadedAdapters, above zero, has each replica hold at most this many\nadapters at once, loaded when a request first asks for one and unloaded\nleast recently used first to make room. The router sends a request for\nan adapter to a replica that already holds it while that one has room\nto take it. vLLM profiles only: llama-server loads adapters at start\nand cannot load one afterwards. Zero loads every adapter on every\nreplica at start.", Nested: ""}, {Key: "fallbacks", Type: "[]string", Doc: "", Nested: ""}, @@ -503,6 +512,7 @@ var Reference = []Block{ {Key: "llamacpp", Type: "LlamaCppConfig", Doc: "", Nested: "engines.LlamaCppConfig"}, {Key: "trtllm", Type: "TRTLLMConfig", Doc: "TRTLLM is TensorRT-LLM, which the agent does not install: a node\noffers it only when its interpreter is named here.", Nested: "engines.TRTLLMConfig"}, {Key: "sim", Type: "SimConfig", Doc: "", Nested: "engines.SimConfig"}, + {Key: "comfyui", Type: "ComfyUIConfig", Doc: "ComfyUI makes pictures; offered by a node whose ComfyUI is named here.", Nested: "engines.ComfyUIConfig"}, }}, {Package: "guard", Name: "Config", Title: "guard.Config", Doc: "", Fields: []Field{ {Key: "enabled", Type: "*bool", Doc: "Enabled turns every policy on or off at once: false checks nothing,\nwhatever is defined or attached. Unset is on.", Nested: ""}, diff --git a/internal/docgen/render.go b/internal/docgen/render.go index 515b7be..9577755 100644 --- a/internal/docgen/render.go +++ b/internal/docgen/render.go @@ -22,6 +22,7 @@ var Sources = []string{ "internal/agent/engines/sglang.go", "internal/agent/engines/trtllm.go", "internal/agent/engines/sim.go", + "internal/agent/engines/comfyui.go", "internal/agent/supervisor/supervisor.go", "internal/tlsconf/tlsconf.go", "internal/control/catalog/catalog.go", diff --git a/internal/router/images.go b/internal/router/images.go index dab5bc4..71634c9 100644 --- a/internal/router/images.go +++ b/internal/router/images.go @@ -15,7 +15,6 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/auth" "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" "gitea.mynodes.xyz/ilker/llmgw/internal/store" @@ -121,7 +120,7 @@ func (r *Router) imageRouteRefusal(path, model string, body []byte) string { } return "" } - if known && m.External != nil && m.External.API == catalog.APIComfyUI { + if known && m.ServesComfyUI() { return fmt.Sprintf("%s makes pictures: ask for them at %s", m.Name, imagesPath) } return ""