llmgw/internal/comfy/comfy_test.go

package comfy

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestTheBuiltInWorkflowsFillIn(t *testing.T) {
	for _, name := range Builtins() {
		if raw, _ := builtin.ReadFile("workflows/" + name + ".json"); strings.Contains(string(raw), "{{image}}") {
			continue // made from a picture: TestTheBuiltInMeshWorkflowFillsIn
		}
		tpl, err := Load(name)
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		g, err := tpl.Fill(Params{Checkpoint: "flux.safetensors", Prompt: `a "red" fox`, Width: 768, Height: 512, N: 2, Seed: 42})
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		b, _ := json.Marshal(g)
		s := string(b)
		for _, want := range []string{`"ckpt_name":"flux.safetensors"`, `"text":"a \"red\" fox"`, `"width":768`,
			`"height":512`, `"batch_size":2`, `"seed":42`} {
			if !strings.Contains(s, want) {
				t.Errorf("%s: %s not in %s", name, want, s)
			}
		}
		if strings.Contains(s, "{{") {
			t.Errorf("%s: a placeholder was left: %s", name, s)
		}
	}
}

func TestTheBuiltInMeshWorkflowFillsIn(t *testing.T) {
	tpl, err := Load("triposr")
	if err != nil {
		t.Fatal(err)
	}
	g, err := tpl.Fill(Params{Checkpoint: "model.ckpt", Image: "llmgw-1.png"})
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(g)
	s := string(b)
	// The picture is cut out and centred before TripoSR sees it: without
	// that, its background becomes geometry.
	for _, want := range []string{`"image":"llmgw-1.png"`, `"model":"model.ckpt"`, `"geometry_resolution":256`,
		`"class_type":"LLMGWPrepareForMesh"`, `"reference_image":["5",0]`} {
		if !strings.Contains(s, want) {
			t.Errorf("%s not in %s", want, s)
		}
	}
	if strings.Contains(s, "{{") {
		t.Errorf("a placeholder was left: %s", s)
	}
}

func TestAStepCountGivenReplacesTheDefault(t *testing.T) {
	tpl, _ := Load("flux-schnell")
	g, _ := tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512})
	if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 4.0 {
		t.Errorf("default steps = %v, want 4", steps)
	}
	g, _ = tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512, Steps: 8})
	if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 8 {
		t.Errorf("steps = %v, want 8", steps)
	}
}

func TestAPlaceholderWithoutAValueOrADefaultIsSaid(t *testing.T) {
	tpl, err := Load(`{"1":{"class_type":"X","inputs":{"a":"{{checkpoint}}","b":"size {{width}}x{{height}}"}}}`)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := tpl.Fill(Params{Width: 5}); err == nil || !strings.Contains(err.Error(), "checkpoint, height") {
		t.Fatalf("err = %v", err)
	}
	g, err := tpl.Fill(Params{Checkpoint: "c", Width: 5, Height: 6})
	if err != nil {
		t.Fatal(err)
	}
	if b := g["1"].(map[string]any)["inputs"].(map[string]any)["b"]; b != "size 5x6" {
		t.Errorf("b = %v", b)
	}
}

func TestAWorkflowInTheUIsFormatIsRefused(t *testing.T) {
	if _, err := Load(`{"nodes":[{"id":1}],"links":[]}`); err == nil || !strings.Contains(err.Error(), "API format") {
		t.Fatalf("err = %v", err)
	}
	if _, err := Load("sdxl-turbo"); err == nil {
		t.Fatal("an unknown built-in was loaded")
	}
}

// fakeComfy is just enough of ComfyUI: /prompt queues, /history answers
// after a few asks, /view serves a picture.
type fakeComfy struct {
	mu          sync.Mutex
	asked       int
	fail        string
	refuse      bool
	deleted     []string
	queued      map[string]any
	readyAt     int
	srv         *httptest.Server
	interrupted bool
}

func newFakeComfy(t *testing.T) *fakeComfy {
	f := &fakeComfy{readyAt: 2}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		f.mu.Lock()
		defer f.mu.Unlock()
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/prompt":
			if f.refuse {
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprint(w, `{"error":{"message":"Prompt outputs failed validation","details":""},"node_errors":{"1":{"class_type":"CheckpointLoaderSimple","errors":[{"message":"Value not in list","details":"ckpt_name: 'x' not in []"}]}}}`)
				return
			}
			var b struct {
				Prompt map[string]any `json:"prompt"`
			}
			_ = json.NewDecoder(r.Body).Decode(&b)
			f.queued = b.Prompt
			fmt.Fprint(w, `{"prompt_id":"p1","number":1,"node_errors":{}}`)
		case r.URL.Path == "/history/p1":
			f.asked++
			if f.asked < f.readyAt {
				fmt.Fprint(w, `{}`)
				return
			}
			if f.fail != "" {
				fmt.Fprintf(w, `{"p1":{"status":{"status_str":"error","completed":false,"messages":[["execution_start",{}],["execution_error",{"node_type":"KSampler","exception_message":%q}]]},"outputs":{}}}`, f.fail)
				return
			}
			fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true,"messages":[]},"outputs":{"7":{"images":[{"filename":"a.png","subfolder":"","type":"temp"},{"filename":"b.png","subfolder":"","type":"temp"}]}}}}`)
		case r.URL.Path == "/view":
			w.Header().Set("Content-Type", "image/png")
			fmt.Fprint(w, "PNG:"+r.URL.Query().Get("filename"))
		case r.URL.Path == "/queue" && r.Method == http.MethodPost:
			var b struct {
				Delete []string `json:"delete"`
			}
			_ = json.NewDecoder(r.Body).Decode(&b)
			f.deleted = append(f.deleted, b.Delete...)
		case r.URL.Path == "/queue":
			fmt.Fprint(w, `{"queue_running":[[1,"p1",{},{},[]]],"queue_pending":[]}`)
		case r.URL.Path == "/interrupt":
			f.interrupted = true
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(f.srv.Close)
	return f
}

func TestGenerateWaitsForTheWorkflowAndFetchesItsPictures(t *testing.T) {
	f := newFakeComfy(t)
	c := &Client{Poll: time.Millisecond}
	images, err := c.Generate(t.Context(), f.srv.URL, "", map[string]any{"1": map[string]any{"class_type": "X"}})
	if err != nil {
		t.Fatal(err)
	}
	if len(images) != 2 || string(images[0].Data) != "PNG:a.png" || images[1].ContentType != "image/png" {
		t.Fatalf("images = %+v", images)
	}
	if f.queued["1"] == nil {
		t.Error("the workflow was not what was queued")
	}
}

func TestAWorkflowThatFailsSaysWhy(t *testing.T) {
	f := newFakeComfy(t)
	f.fail = "CUDA out of memory"
	_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
	if err == nil || !strings.Contains(err.Error(), "KSampler: CUDA out of memory") {
		t.Fatalf("err = %v", err)
	}
}

func TestAWorkflowComfyUIRefusesIsNotRetried(t *testing.T) {
	f := newFakeComfy(t)
	f.refuse = true
	_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
	ce, ok := err.(*Error)
	if !ok || ce.Retryable || !strings.Contains(ce.Msg, "Value not in list") {
		t.Fatalf("err = %#v", err)
	}
}

func TestAComfyUIThatCannotBeReachedIsRetryable(t *testing.T) {
	f := newFakeComfy(t)
	f.srv.Close()
	_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
	if ce, ok := err.(*Error); !ok || !ce.Retryable {
		t.Fatalf("err = %#v", err)
	}
}

func TestAClientThatLeavesTakesItsWorkflowOffTheCard(t *testing.T) {
	f := newFakeComfy(t)
	f.readyAt = 1 << 30
	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond)
	defer cancel()
	_, err := (&Client{Poll: time.Millisecond}).Generate(ctx, f.srv.URL, "", map[string]any{})
	if err == nil {
		t.Fatal("no error")
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.deleted) != 1 || f.deleted[0] != "p1" || !f.interrupted {
		t.Errorf("deleted %v, interrupted %v", f.deleted, f.interrupted)
	}
}

func TestAPictureLosesTheWorkflowComfyUIWroteIntoIt(t *testing.T) {
	chunk := func(kind, data string) []byte {
		b := make([]byte, 0, 12+len(data))
		b = binary.BigEndian.AppendUint32(b, uint32(len(data)))
		b = append(b, kind...)
		b = append(b, data...)
		return append(b, 0, 0, 0, 0) // a CRC nobody here checks
	}
	png := []byte("\x89PNG\r\n\x1a\n")
	png = append(png, chunk("IHDR", "0123456789abc")...)
	png = append(png, chunk("tEXt", "prompt\x00{\"1\":...}")...)
	png = append(png, chunk("IDAT", "pixels")...)
	png = append(png, chunk("iTXt", "workflow\x00...")...)
	png = append(png, chunk("IEND", "")...)
	got := string(withoutPNGText(png))
	if strings.Contains(got, "prompt") || strings.Contains(got, "workflow") {
		t.Errorf("text left in: %q", got)
	}
	for _, kind := range []string{"IHDR", "IDAT", "pixels", "IEND"} {
		if !strings.Contains(got, kind) {
			t.Errorf("%s gone: %q", kind, got)
		}
	}
	if broken := png[:30]; string(withoutPNGText(broken)) != string(broken) {
		t.Error("a truncated PNG was changed")
	}
}

func TestATrueOrFalseDefaultIsABoolean(t *testing.T) {
	// A BOOLEAN input refuses the word "false": sf3d's albedo_match.
	tpl, err := Load((`{"5":{"class_type":"LLMGWSaveMesh","inputs":{"albedo_match":"{{albedo_match|false}}","count":"{{count|2}}","s":"{{s|x}}"}}}`))
	if err != nil {
		t.Fatal(err)
	}
	g, err := tpl.Fill(Params{})
	if err != nil {
		t.Fatal(err)
	}
	in := g["5"].(map[string]any)["inputs"].(map[string]any)
	if in["albedo_match"] != false || in["count"] != 2.0 || in["s"] != "x" {
		t.Errorf("inputs %#v", in)
	}
}