llmgw/internal/control/catalog/workflow.go

package catalog

import (
	"fmt"
	"regexp"
	"slices"
	"strings"
)

// What a ComfyUI workflow takes and what it makes, said in the catalog so a
// model is data rather than code: a request is checked against its inputs
// before a job is written, a client can build a form from /v1/models, and
// the answer is picked from the files the workflow saved by its output kind.

// The kinds of thing a ComfyUI model makes, beside pictures and meshes.
const (
	OutputVideo = "video"
	OutputAudio = "audio"
	// OutputFiles is whatever the workflow saves, of any type.
	OutputFiles = "files"
)

// Outputs are every output kind a ComfyUI model may name.
var Outputs = []string{OutputPictures, OutputMeshes, OutputVideo, OutputAudio, OutputFiles}

// The types of a workflow's input.
const (
	InputText   = "text"
	InputInt    = "int"
	InputFloat  = "float"
	InputBool   = "bool"
	InputChoice = "choice"
	InputImage  = "image"
	InputMask   = "mask"
	InputAudio  = "audio"
	InputVideo  = "video"
	InputFile   = "file"
)

// InputTypes are every type an input may have.
var InputTypes = []string{InputText, InputInt, InputFloat, InputBool, InputChoice,
	InputImage, InputMask, InputAudio, InputVideo, InputFile}

// Input is one thing a workflow takes from a request: the placeholder
// {{name}} in its graph, filled with the request's inputs.name.
type Input struct {
	// Name is the placeholder's, and the request's key.
	Name string `json:"name" koanf:"name"`
	// Type is one of text, int, float, bool, choice, image, mask, audio,
	// video, file. The last five are files, sent as base64, a data URL, or
	// the id of a file uploaded to /v1/files, and put in ComfyUI's input
	// directory; the placeholder gets the name they were given there.
	Type string `json:"type" koanf:"type"`
	// Required inputs must be in every request; the others take Default,
	// or the workflow's own {{name|default}}.
	Required bool `json:"required,omitempty" koanf:"required"`
	Default  any  `json:"default,omitempty" koanf:"default"`
	// Min and Max bound a number; for text, Max bounds its length.
	Min *float64 `json:"min,omitempty" koanf:"min"`
	Max *float64 `json:"max,omitempty" koanf:"max"`
	// Choices are what a choice may be.
	Choices []string `json:"choices,omitempty" koanf:"choices"`
	// MaxMB bounds a file input; zero is the type's own bound.
	MaxMB int `json:"max_mb,omitempty" koanf:"max_mb"`
	// Doc says what the input is for, to whoever fills it in.
	Doc string `json:"doc,omitempty" koanf:"doc"`
}

// IsFile reports whether the input is a file put in ComfyUI's input directory.
func (in Input) IsFile() bool {
	switch in.Type {
	case InputImage, InputMask, InputAudio, InputVideo, InputFile:
		return true
	}
	return false
}

// MaxBytes is the most a file input may be.
func (in Input) MaxBytes() int64 {
	if in.MaxMB > 0 {
		return int64(in.MaxMB) << 20
	}
	switch in.Type {
	case InputImage, InputMask:
		return 16 << 20
	case InputVideo:
		return 256 << 20
	}
	return 64 << 20
}

var inputName = regexp.MustCompile(`^[a-z][a-z0-9_]{0,63}$`)

func (in Input) validate(model string) error {
	if !inputName.MatchString(in.Name) {
		return fmt.Errorf("model %q: input %q is not a name of lowercase letters, digits and '_'", model, in.Name)
	}
	if !slices.Contains(InputTypes, in.Type) {
		return fmt.Errorf("model %q: input %q has type %q, not one of %s", model, in.Name, in.Type, strings.Join(InputTypes, ", "))
	}
	if in.Type == InputChoice && len(in.Choices) == 0 {
		return fmt.Errorf("model %q: input %q is a choice with no choices", model, in.Name)
	}
	if in.Type != InputChoice && len(in.Choices) > 0 {
		return fmt.Errorf("model %q: input %q has choices but is not a choice", model, in.Name)
	}
	if in.Min != nil && in.Max != nil && *in.Min > *in.Max {
		return fmt.Errorf("model %q: input %q has min above max", model, in.Name)
	}
	if in.Required && in.Default != nil {
		return fmt.Errorf("model %q: input %q is required and has a default: it is one or the other", model, in.Name)
	}
	if in.IsFile() && in.Default != nil {
		return fmt.Errorf("model %q: input %q is a file and cannot have a default", model, in.Name)
	}
	if in.MaxMB < 0 {
		return fmt.Errorf("model %q: input %q has a negative max_mb", model, in.Name)
	}
	return nil
}

// ComfyOutput is what a ComfyUI model makes: pictures unless it says.
func (m Model) ComfyOutput() string { return m.comfyOutput() }

// ComfyInputs are what a ComfyUI model's workflow takes from a request.
func (m Model) ComfyInputs() []Input {
	if m.ComfyUI == nil {
		return nil
	}
	return m.ComfyUI.Inputs
}

// RunsWorkflows reports whether a model answers /v1/workflows/runs: any the
// cluster's ComfyUI or a ComfyUI outside it serves.
func (m Model) RunsWorkflows() bool { return m.ServesComfyUI() }

func validOutput(model, out string) error {
	if !slices.Contains(Outputs, out) {
		return fmt.Errorf("model %q: comfyui.output %q is not one of %s", model, out, strings.Join(Outputs, ", "))
	}
	return nil
}

func validateInputs(model string, inputs []Input) error {
	seen := map[string]bool{}
	for _, in := range inputs {
		if err := in.validate(model); err != nil {
			return err
		}
		if seen[in.Name] {
			return fmt.Errorf("model %q: input %q is declared twice", model, in.Name)
		}
		seen[in.Name] = true
	}
	return nil
}