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"`
	// OwnerUserID is m's auth.users.id on msupabase. The cloud-sync writer
	// uses it to populate imagen.images.owner_user_id (NOT NULL, owns RLS).
	// Empty disables DB inserts even when cloud_sync is on.
	OwnerUserID    string                 `yaml:"owner_user_id"`
	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"`
	// CloudSync controls whether successful generations also upload to
	// Supabase Storage and insert into imagen.images. Tri-state mirroring
	// Preview: "auto" (default — on when SUPABASE_URL + SUPABASE_SERVICE_KEY
	// are set), "on" (errors if env unset), "off". --no-cloud overrides.
	CloudSync string `yaml:"cloud_sync"`
}

// 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)
	}
	switch c.Output.CloudSync {
	case "", "auto", "on", "off":
	default:
		return fmt.Errorf("output.cloud_sync = %q (must be auto|on|off)", c.Output.CloudSync)
	}
	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

# Owner UUID for the cloud-sync row in imagen.images. Look up via:
#   SELECT id FROM auth.users WHERE email = '<your-supabase-email>';
# Empty disables imagen.images inserts even when cloud_sync is on.
owner_user_id: ""

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
  # Sync the PNG to Supabase Storage (bucket: imagen-generated) and insert
  # a row into imagen.images. Reads SUPABASE_URL + SUPABASE_SERVICE_KEY
  # from env (same as mai.imagen_usage cost-tracking).
  # auto (default): on iff env is configured AND owner_user_id is set.
  # on: always upload (errors if env or owner_user_id is missing).
  # off: never upload. --no-cloud also forces off per-call.
  cloud_sync: auto

backends:
  # FLUX.1-schnell on the local ComfyUI server. The "workflow" key picks the
  # bundled template under internal/backend/workflows/; omit it for back-compat
  # (defaults to flux1-schnell). See docs/backends.md for the per-model setup.
  flux-schnell-local:
    type: comfyui
    base_url: http://mrock:8188
    workflow: flux1-schnell
    # Filename of the unet checkpoint inside the ComfyUI server's
    # models/unet/ directory.
    model: flux1-schnell.safetensors
    vae: ae.safetensors
    clip_l: clip_l.safetensors
    clip_t5: t5xxl_fp8_e4m3fn.safetensors
    dtype: fp8_e4m3fn
    default_steps: 4
    default_sampler: euler
    default_scheduler: simple
    default_cfg: 1.0

  # FLUX.2 [klein] 4B distilled — sub-second on RTX 4070 Ti SUPER.
  # Weights: BFL non-commercial; flux-2-klein-base-4b-fp8 in models/unet/,
  # qwen_3_4b in models/text_encoders/, flux2-vae in models/vae/.
  flux2-klein-local:
    type: comfyui
    base_url: http://mrock:8188
    workflow: flux2-klein
    model: flux-2-klein-base-4b-fp8.safetensors
    vae: flux2-vae.safetensors
    clip: qwen_3_4b.safetensors
    dtype: fp8_e4m3fn
    default_steps: 4
    default_sampler: euler
    default_scheduler: simple
    default_cfg: 1.0
    guidance: 4.0

  # SD3.5 medium — single-checkpoint variant that bundles the three text
  # encoders inside the .safetensors. Drop into models/checkpoints/.
  sd35-medium-local:
    type: comfyui
    base_url: http://mrock:8188
    workflow: sd35-medium
    model: sd3.5_medium_incl_clips_t5xxlfp8scaled.safetensors
    default_steps: 28
    default_sampler: dpmpp_2m
    default_scheduler: sgm_uniform
    default_cfg: 4.5
    shift: 3.0

  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
}