mAi: #10 - multi-model backend expansion (workflow templates + compare harness)
Path 1 architecture: one comfyui adapter, workflows as data.
- workflow_template.go: embed.FS + token substitution with type-preserving
whole-value placeholders. ${prompt} → string, ${seed} → int64,
${cfg} → float64 — no JSON round-tripping. Partial matches ignored.
- comfyui.go: refactored to load workflow from embedded FS or filesystem
path. Back-compat preserved: workflow: defaults to flux1-schnell.
- workflows/{flux1-schnell,flux2-klein,sd35-medium}.json — bundled
templates. flux1-schnell migrated from hardcoded with identical node IDs.
- compare.go: new `imagen compare` subcommand. Sequential N-backend run
(one GPU on mRock — parallel would OOM), per-backend PNG, sidecar JSON
with per-model metadata + errors, composite contact sheet via Go image
package (no ImageMagick dep).
- Sample config gains flux2-klein-local + sd35-medium-local instances.
- docs/backends.md: architecture rationale + per-model HF download paths
+ how to add a new bundled workflow + compare-harness reference.
Live smoke verified: compare mock + flux-schnell-local at 768×768 →
both PNGs written, sidecar JSON has workflow="flux1-schnell" + full
metadata, contact sheet renders. Worker contract (Request → Generate)
unchanged, so flexsiebels /imagine UI API surface preserved.
Tests: 11 existing comfyui + 6 new workflow_template + 5 new compare
tests, all green.
Adding a new model is now yaml + JSON, never Go.
This commit is contained in:
153
internal/backend/workflow_template_test.go
Normal file
153
internal/backend/workflow_template_test.go
Normal file
@@ -0,0 +1,153 @@
|
||||
package backend
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBundledWorkflowsParseable(t *testing.T) {
|
||||
names := BundledWorkflowNames()
|
||||
if len(names) == 0 {
|
||||
t.Fatal("expected at least one bundled workflow")
|
||||
}
|
||||
mustHave := []string{"flux1-schnell", "flux2-klein", "sd35-medium"}
|
||||
for _, want := range mustHave {
|
||||
if !slices.Contains(names, want) {
|
||||
t.Errorf("bundled workflows missing %q (have: %v)", want, names)
|
||||
}
|
||||
}
|
||||
// Every bundled template must parse and contain at least one node.
|
||||
for _, n := range names {
|
||||
wf, err := LoadWorkflowTemplate(n)
|
||||
if err != nil {
|
||||
t.Errorf("LoadWorkflowTemplate(%q): %v", n, err)
|
||||
continue
|
||||
}
|
||||
if len(wf) == 0 {
|
||||
t.Errorf("workflow %q has zero nodes", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadWorkflowFromFilesystem(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "custom.json")
|
||||
body := `{"1":{"class_type":"X","inputs":{"v":"${prompt}"}}}`
|
||||
if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
|
||||
t.Fatalf("write tmp workflow: %v", err)
|
||||
}
|
||||
wf, err := LoadWorkflowTemplate(path)
|
||||
if err != nil {
|
||||
t.Fatalf("load from path: %v", err)
|
||||
}
|
||||
if _, ok := wf["1"]; !ok {
|
||||
t.Errorf("custom workflow missing node 1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadWorkflowUnknownNameErrors(t *testing.T) {
|
||||
_, err := LoadWorkflowTemplate("definitely-not-a-real-workflow")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown workflow name")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not found") {
|
||||
t.Errorf("error should say not found, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubstituteWorkflowPreservesTypes(t *testing.T) {
|
||||
wf := map[string]any{
|
||||
"31": map[string]any{
|
||||
"class_type": "KSampler",
|
||||
"inputs": map[string]any{
|
||||
"seed": "${seed}",
|
||||
"steps": "${steps}",
|
||||
"text": "${prompt}",
|
||||
"cfg": "${cfg}",
|
||||
},
|
||||
},
|
||||
}
|
||||
subs := map[string]any{
|
||||
"seed": int64(42),
|
||||
"steps": 11,
|
||||
"prompt": "a cat",
|
||||
"cfg": 4.5,
|
||||
}
|
||||
used, err := SubstituteWorkflow(wf, subs)
|
||||
if err != nil {
|
||||
t.Fatalf("Substitute: %v", err)
|
||||
}
|
||||
if len(used) != 4 {
|
||||
t.Errorf("used = %v, want all four", used)
|
||||
}
|
||||
inputs := wf["31"].(map[string]any)["inputs"].(map[string]any)
|
||||
if seed, ok := inputs["seed"].(int64); !ok || seed != 42 {
|
||||
t.Errorf("seed = %T %v, want int64 42", inputs["seed"], inputs["seed"])
|
||||
}
|
||||
if steps, ok := inputs["steps"].(int); !ok || steps != 11 {
|
||||
t.Errorf("steps = %T %v, want int 11", inputs["steps"], inputs["steps"])
|
||||
}
|
||||
if text, ok := inputs["text"].(string); !ok || text != "a cat" {
|
||||
t.Errorf("text = %T %v, want string", inputs["text"], inputs["text"])
|
||||
}
|
||||
if cfg, ok := inputs["cfg"].(float64); !ok || cfg != 4.5 {
|
||||
t.Errorf("cfg = %T %v, want float64 4.5", inputs["cfg"], inputs["cfg"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubstituteWorkflowMissingPlaceholderErrors(t *testing.T) {
|
||||
wf := map[string]any{
|
||||
"1": map[string]any{"inputs": map[string]any{"v": "${missing}"}},
|
||||
}
|
||||
_, err := SubstituteWorkflow(wf, map[string]any{})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing placeholder")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "${missing}") {
|
||||
t.Errorf("error should name the placeholder, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubstituteWorkflowOnlyWholeTokens(t *testing.T) {
|
||||
// Partial-match strings ("prefix ${prompt} suffix") are NOT substituted —
|
||||
// the placeholder must be the whole value so we can preserve types.
|
||||
wf := map[string]any{
|
||||
"1": map[string]any{"inputs": map[string]any{
|
||||
"keep_string": "stuff with ${prompt} inside",
|
||||
"replace_full": "${prompt}",
|
||||
}},
|
||||
}
|
||||
used, err := SubstituteWorkflow(wf, map[string]any{"prompt": "x"})
|
||||
if err != nil {
|
||||
t.Fatalf("Substitute: %v", err)
|
||||
}
|
||||
inputs := wf["1"].(map[string]any)["inputs"].(map[string]any)
|
||||
if inputs["keep_string"].(string) != "stuff with ${prompt} inside" {
|
||||
t.Errorf("partial match should be left alone, got %q", inputs["keep_string"])
|
||||
}
|
||||
if inputs["replace_full"].(string) != "x" {
|
||||
t.Errorf("full-value match should substitute, got %q", inputs["replace_full"])
|
||||
}
|
||||
if _, ok := used["prompt"]; !ok {
|
||||
t.Errorf("used should track keys that fired")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFlux1SchnellTemplateMatchesLegacyShape(t *testing.T) {
|
||||
// Regression guard against the historical hardcoded workflow: every
|
||||
// node ID the old Comfy.buildWorkflow used must still exist in the
|
||||
// migrated template.
|
||||
wf, err := LoadWorkflowTemplate("flux1-schnell")
|
||||
if err != nil {
|
||||
t.Fatalf("load flux1-schnell: %v", err)
|
||||
}
|
||||
legacyNodes := []string{"6", "8", "9", "10", "11", "12", "13", "27", "30", "31"}
|
||||
for _, id := range legacyNodes {
|
||||
if _, ok := wf[id]; !ok {
|
||||
t.Errorf("flux1-schnell template missing node %q (legacy parity)", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user