llmgw/internal/agent/engines/comfyui.go

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 the environment the agent installed:
// its interpreter, with ComfyUI's code beside its bin, in comfyui/.
func (c *ComfyUI) UseInstalled(python string) { c.installed.Store(&python) }

// 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 filepath.Join(filepath.Dir(filepath.Dir(*p)), "comfyui")
	}
	return c.cfg.Dir
}

func (c *ComfyUI) python() string {
	if p := c.installed.Load(); p != nil {
		return *p
	}
	if c.cfg.Python != "" {
		return c.cfg.Python
	}
	return filepath.Join(c.cfg.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
}