ImaGen/internal/config/config.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,
// and whether `imagen generate` opens a tmux preview window.
type OutputConfig struct {
	Directory         string `yaml:"directory"`
	Naming            string `yaml:"naming"`
	WriteMetadataJSON bool   `yaml:"write_metadata_json"`
	// Preview is the tri-state preview mode: "auto" (default), "on", "off".
	// Empty / unset is treated as "auto". $IMAGEN_PREVIEW and the
	// --preview/--no-preview flags override this in turn.
	Preview string `yaml:"preview"`
}

// 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)
		}
	}
	switch c.Output.Preview {
	case "", "auto", "on", "off":
	default:
		return fmt.Errorf("output.preview = %q (must be auto|on|off)", c.Output.Preview)
	}
	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
  # Open a tmux window with tmux-img after a successful generation.
  # auto (default): preview iff stdout is a TTY and $TMUX is set.
  # on: always preview (errors outside a tmux session).
  # off: never preview (use this for batch / CI callers).
  preview: auto

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-schnell-replicate:
    type: replicate
    api_token_env: REPLICATE_API_TOKEN
    model: black-forest-labs/flux-schnell
    default_steps: 4
    default_aspect_ratio: "1:1"

  flux-dev-replicate:
    type: replicate
    api_token_env: REPLICATE_API_TOKEN
    model: black-forest-labs/flux-dev
    default_steps: 28
    default_aspect_ratio: "1:1"

  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
}