mAi
127bbf3ed5
internal/backend/comfyui.go implements the Backend interface against ComfyUI's /prompt + /history + /view HTTP API. Workflow is the canonical FLUX.1 schnell shape — UNETLoader + DualCLIPLoader (clip_l + t5xxl fp8) + VAELoader + ModelSamplingFlux + KSampler — assembled as a Go map per request so Width / Height / Seed / Steps / sampler / scheduler all flow into the right node inputs. Resilience: one retry on /prompt 5xx and transient network errors, no retry on 4xx. Connection-refused / timeouts surface a 'boot-whitetower mrock' hint. node_errors mentioning a missing unet point users at docs/setup-comfyui-mrock.md (matches both the 4xx and 200-with-errors shapes ComfyUI uses across versions). Result.Metadata carries model, seed_used, latency_ms, steps, sampler, scheduler, width, height, prompt_id, client_id, plus best-effort vram_used_mib pulled from /system_stats post-gen. Tests use httptest with poll interval squashed to 1ms — no real mRock dependency. Coverage: happy path, defaults, retry-once on 5xx, give-up after two 5xx, no-retry on 4xx, missing-model hint (both 4xx and 200+node_errors paths), history-error surfaced, /view 4xx, unreachable host, ctx cancel during poll, workflow-shape assertion, registration. Config sample: flux-schnell-local is now default_backend; the user-facing block names the unet file by basename (the mapping into models/unet/ is the server's convention, captured in docs/setup-comfyui-mrock.md from phase 1). Smoke verified end-to-end: imagen generate ... --backend flux-schnell-local --size 1024x1024 --output /tmp/cat-via-cli.png on mRock returned a 1024x1024 PNG of a cat in a fishbowl in 10.3s with a sidecar carrying seed + latency_ms + the rest of the metadata.
77 lines
1.8 KiB
Go
77 lines
1.8 KiB
Go
package config
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
)
|
|
|
|
func TestLoadAndValidate(t *testing.T) {
|
|
dir := t.TempDir()
|
|
path := filepath.Join(dir, "imagen.yaml")
|
|
if err := os.WriteFile(path, []byte(Sample), 0o644); err != nil {
|
|
t.Fatalf("write sample: %v", err)
|
|
}
|
|
cfg, err := Load(path)
|
|
if err != nil {
|
|
t.Fatalf("Load: %v", err)
|
|
}
|
|
if cfg.DefaultBackend != "flux-schnell-local" {
|
|
t.Errorf("default = %q", cfg.DefaultBackend)
|
|
}
|
|
mock, ok := cfg.Backends["mock"]
|
|
if !ok {
|
|
t.Fatalf("mock backend missing")
|
|
}
|
|
if mock.Type != "mock" {
|
|
t.Errorf("mock type = %q", mock.Type)
|
|
}
|
|
flux, ok := cfg.Backends["flux-schnell-local"]
|
|
if !ok {
|
|
t.Fatalf("flux backend missing")
|
|
}
|
|
if flux.Type != "comfyui" {
|
|
t.Errorf("flux type = %q", flux.Type)
|
|
}
|
|
if flux.Raw["base_url"] != "http://mrock:8188" {
|
|
t.Errorf("flux base_url = %v", flux.Raw["base_url"])
|
|
}
|
|
if flux.Raw["model"] != "flux1-schnell.safetensors" {
|
|
t.Errorf("flux model = %v", flux.Raw["model"])
|
|
}
|
|
}
|
|
|
|
func TestValidateRejectsUnknownDefault(t *testing.T) {
|
|
c := &Config{
|
|
DefaultBackend: "ghost",
|
|
Backends: map[string]BackendSpec{"real": {Type: "mock"}},
|
|
}
|
|
if err := c.Validate(); err == nil {
|
|
t.Errorf("expected error for unknown default_backend")
|
|
}
|
|
}
|
|
|
|
func TestValidateRejectsMissingType(t *testing.T) {
|
|
c := &Config{
|
|
Backends: map[string]BackendSpec{"x": {}},
|
|
}
|
|
if err := c.Validate(); err == nil {
|
|
t.Errorf("expected error for missing type")
|
|
}
|
|
}
|
|
|
|
func TestExpandPath(t *testing.T) {
|
|
home, _ := os.UserHomeDir()
|
|
cases := map[string]string{
|
|
"": "",
|
|
"/abs/path": "/abs/path",
|
|
"~": home,
|
|
"~/foo/bar": filepath.Join(home, "foo/bar"),
|
|
}
|
|
for in, want := range cases {
|
|
if got := ExpandPath(in); got != want {
|
|
t.Errorf("ExpandPath(%q) = %q, want %q", in, got, want)
|
|
}
|
|
}
|
|
}
|