Files
ImaGen/internal/config/config.go
mAi 127bbf3ed5 mAi: #2 - phase 2 ComfyUI Go adapter, tests, config sample
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.
2026-05-08 16:59:21 +02:00

148 lines
3.8 KiB
Go

// Package config loads ~/.config/imagen.yaml. The framework knows the global
// shape (default backend + output settings + a per-backend block); each
// adapter owns the schema of its own block.
package config
import (
"errors"
"fmt"
"os"
"path/filepath"
"gopkg.in/yaml.v3"
)
// Config is the top-level shape of imagen.yaml.
type Config struct {
DefaultBackend string `yaml:"default_backend"`
Output OutputConfig `yaml:"output"`
Backends map[string]BackendSpec `yaml:"backends"`
}
// OutputConfig controls where generated images and metadata sidecars land.
type OutputConfig struct {
Directory string `yaml:"directory"`
Naming string `yaml:"naming"`
WriteMetadataJSON bool `yaml:"write_metadata_json"`
}
// BackendSpec is one entry under `backends:`. Type identifies the adapter;
// the rest is opaque to the framework and handed to the adapter as-is.
type BackendSpec struct {
Type string `yaml:"type"`
Raw map[string]any `yaml:",inline"`
}
// DefaultPath returns ~/.config/imagen.yaml, honouring XDG_CONFIG_HOME.
func DefaultPath() (string, error) {
if x := os.Getenv("XDG_CONFIG_HOME"); x != "" {
return filepath.Join(x, "imagen.yaml"), nil
}
home, err := os.UserHomeDir()
if err != nil {
return "", err
}
return filepath.Join(home, ".config", "imagen.yaml"), nil
}
// Load reads and validates the config at path. If path is empty the default
// path is used. A missing file returns os.ErrNotExist so callers can decide
// whether to fall back to defaults.
func Load(path string) (*Config, error) {
if path == "" {
p, err := DefaultPath()
if err != nil {
return nil, err
}
path = p
}
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
cfg := &Config{}
if err := yaml.Unmarshal(data, cfg); err != nil {
return nil, fmt.Errorf("parse %s: %w", path, err)
}
if err := cfg.Validate(); err != nil {
return nil, fmt.Errorf("validate %s: %w", path, err)
}
return cfg, nil
}
// Validate enforces the framework-level invariants. Adapter-level validation
// happens when the adapter constructor runs.
func (c *Config) Validate() error {
if c.DefaultBackend != "" {
if _, ok := c.Backends[c.DefaultBackend]; !ok {
return fmt.Errorf("default_backend %q is not defined under backends:", c.DefaultBackend)
}
}
for name, spec := range c.Backends {
if name == "" {
return errors.New("empty backend name")
}
if spec.Type == "" {
return fmt.Errorf("backend %q is missing a type:", name)
}
}
return nil
}
// Sample is the canonical example written by `imagen config init`.
const Sample = `# imagen.yaml — config for the imagen CLI.
# Adapters get only their own sub-block at construction. Add a new backend by
# implementing the Backend interface, registering its type name, and listing
# an instance here.
default_backend: flux-schnell-local
output:
directory: ~/Pictures/imagen
naming: "{date}-{slug}-{seed}.png"
write_metadata_json: true
backends:
flux-schnell-local:
type: comfyui
base_url: http://mrock:8188
# Filename of the unet checkpoint inside the ComfyUI server's
# models/unet/ directory. See docs/setup-comfyui-mrock.md.
model: flux1-schnell.safetensors
default_steps: 4
default_sampler: euler
default_scheduler: simple
mock:
type: mock
flux-dev-replicate:
type: replicate
api_token_env: REPLICATE_API_TOKEN
model: black-forest-labs/flux-dev
default_steps: 28
dalle3:
type: openai
api_key_env: OPENAI_API_KEY
model: dall-e-3
`
// ExpandPath resolves leading ~ to the user's home directory.
func ExpandPath(p string) string {
if p == "" || p[0] != '~' {
return p
}
home, err := os.UserHomeDir()
if err != nil {
return p
}
if len(p) == 1 {
return home
}
if p[1] == '/' {
return filepath.Join(home, p[2:])
}
return p
}