ImaGen/cmd/imagen/generate.go

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"

	"mgit.msbls.de/m/ImaGen/internal/backend"
	"mgit.msbls.de/m/ImaGen/internal/config"
	"mgit.msbls.de/m/ImaGen/internal/output"
	"mgit.msbls.de/m/ImaGen/internal/preview"
	"mgit.msbls.de/m/ImaGen/internal/prompt"
	"mgit.msbls.de/m/ImaGen/internal/usage"
)

func runGenerate(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("generate", flag.ContinueOnError)
	var (
		backendName  string
		size         string
		outPath      string
		seed         int64
		steps        int
		style        string
		negative     string
		configPath   string
		noSidecar    bool
		previewOn    bool
		previewOff   bool
	)
	fs.StringVar(&backendName, "backend", "", "backend instance name (default: config.default_backend)")
	fs.StringVar(&size, "size", "1024x1024", "WxH, e.g. 1024x1024")
	fs.StringVar(&outPath, "output", "", "explicit output path (overrides config naming template)")
	fs.Int64Var(&seed, "seed", 0, "deterministic seed (0 = backend default)")
	fs.IntVar(&steps, "steps", 0, "diffusion steps (0 = backend default)")
	fs.StringVar(&style, "style", "", "style preset name (see imagen config init for the list)")
	fs.StringVar(&negative, "negative", "", "negative prompt (ignored by backends that don't support it)")
	fs.StringVar(&configPath, "config", "", "config file path (default: ~/.config/imagen.yaml)")
	fs.BoolVar(&noSidecar, "no-sidecar", false, "skip the JSON sidecar even if config enables it")
	fs.BoolVar(&previewOn, "preview", false, "force tmux preview window on (errors outside $TMUX)")
	fs.BoolVar(&previewOff, "no-preview", false, "skip the tmux preview window")
	fs.Usage = func() {
		fmt.Fprintln(fs.Output(), `Usage: imagen generate "<prompt>" [flags]`)
		fs.PrintDefaults()
	}
	// stdlib flag stops parsing at the first non-flag arg, so split the
	// prompt (leading positional args) from the flags ourselves before parsing.
	leadingPositional, flagArgs := splitLeadingPositional(args)
	if err := fs.Parse(flagArgs); err != nil {
		return err
	}
	positional := append(leadingPositional, fs.Args()...)
	if len(positional) == 0 {
		fs.Usage()
		return userErr("missing prompt")
	}
	rawPrompt := strings.Join(positional, " ")

	w, h, err := parseSize(size)
	if err != nil {
		return userErr("bad --size: %v", err)
	}

	cfg, cfgErr := config.Load(configPath)
	if cfgErr != nil && !os.IsNotExist(cfgErr) {
		return cfgErr
	}

	if backendName == "" {
		if cfg != nil {
			backendName = cfg.DefaultBackend
		}
	}
	if backendName == "" {
		return userErr("no --backend given and no default_backend in config")
	}

	be, err := buildBackend(cfg, backendName)
	if err != nil {
		return err
	}
	attachUsageSink(be)

	finalPrompt, err := prompt.Apply(rawPrompt, style)
	if err != nil {
		return userErr("%v", err)
	}

	req := backend.Request{
		Prompt:         finalPrompt,
		NegativePrompt: negative,
		Width:          w,
		Height:         h,
		Steps:          steps,
		Seed:           seed,
		Style:          style,
	}
	res, err := be.Generate(ctx, req)
	if err != nil {
		return fmt.Errorf("backend %q: %w", backendName, err)
	}
	defer res.ImageReader.Close()

	writer := buildWriter(cfg, noSidecar)
	in := output.Inputs{
		Prompt:   rawPrompt,
		Backend:  be.Name(),
		Seed:     seedFromMetadata(res.Metadata, seed),
		Ext:      extFromMime(res.MimeType),
		Metadata: res.Metadata,
	}
	var paths *output.Outputs
	if outPath != "" {
		paths, err = writer.WriteToPath(res.ImageReader, outPath, in)
	} else {
		paths, err = writer.Write(res.ImageReader, in)
	}
	if err != nil {
		return err
	}
	fmt.Println(paths.ImagePath)
	if paths.SidecarPath != "" {
		fmt.Fprintln(os.Stderr, "sidecar:", paths.SidecarPath)
	}

	if err := maybePreview(cfg, previewOn, previewOff, paths.ImagePath, rawPrompt); err != nil {
		// preview failures are warnings — the image already wrote.
		fmt.Fprintln(os.Stderr, "imagen: preview:", err)
	}
	return nil
}

// resolvePreviewMode applies the precedence chain config -> env -> flag.
// Flags win, env beats config, config beats the implicit auto default.
func resolvePreviewMode(cfg *config.Config, flagOn, flagOff bool, env string) (preview.Mode, error) {
	mode := preview.ModeAuto
	if cfg != nil && cfg.Output.Preview != "" {
		m, err := preview.ParseMode(cfg.Output.Preview)
		if err != nil {
			return "", fmt.Errorf("config output.preview: %w", err)
		}
		mode = m
	}
	if env != "" {
		m, err := preview.ParseMode(env)
		if err != nil {
			return "", fmt.Errorf("$IMAGEN_PREVIEW: %w", err)
		}
		mode = m
	}
	if flagOn && flagOff {
		return "", userErr("--preview and --no-preview are mutually exclusive")
	}
	if flagOn {
		mode = preview.ModeOn
	}
	if flagOff {
		mode = preview.ModeOff
	}
	return mode, nil
}

// maybePreview resolves the effective preview mode and, if it says yes,
// spawns a tmux window via tmux-img. Always non-fatal.
func maybePreview(cfg *config.Config, flagOn, flagOff bool, imagePath, rawPrompt string) error {
	mode, err := resolvePreviewMode(cfg, flagOn, flagOff, os.Getenv("IMAGEN_PREVIEW"))
	if err != nil {
		return err
	}
	decision, err := preview.Resolve(mode, os.Getenv("TMUX") != "", stdoutIsTTY())
	if err != nil {
		return err
	}
	if !decision.ShouldPreview {
		return nil
	}
	spawner := &preview.Spawner{}
	return spawner.Spawn(imagePath, output.Slug(rawPrompt))
}

func stdoutIsTTY() bool {
	fi, err := os.Stdout.Stat()
	if err != nil {
		return false
	}
	return fi.Mode()&os.ModeCharDevice != 0
}

// splitLeadingPositional separates the positional args at the start of args
// from the rest (which begins with the first flag). A literal "--" terminator
// pushes everything after it into the positional list and out of flag parsing.
func splitLeadingPositional(args []string) (positional, flags []string) {
	for i, a := range args {
		if a == "--" {
			return append(positional, args[i+1:]...), flags
		}
		if strings.HasPrefix(a, "-") {
			return positional, args[i:]
		}
		positional = append(positional, a)
	}
	return positional, flags
}

func parseSize(s string) (int, int, error) {
	parts := strings.SplitN(s, "x", 2)
	if len(parts) != 2 {
		return 0, 0, fmt.Errorf("expected WxH, got %q", s)
	}
	w, err := strconv.Atoi(parts[0])
	if err != nil {
		return 0, 0, err
	}
	h, err := strconv.Atoi(parts[1])
	if err != nil {
		return 0, 0, err
	}
	return w, h, nil
}

// attachUsageSink wires a Supabase cost-tracking sink into the backend
// when it accepts one and the env is configured. Adapters that record
// usage expose a public Sink field of type backend.UsageSink.
func attachUsageSink(be backend.Backend) {
	r, ok := be.(*backend.Replicate)
	if !ok {
		return
	}
	sink, ok := usage.NewSupabaseSinkFromEnv()
	if !ok {
		return
	}
	r.Sink = sink
}

func buildBackend(cfg *config.Config, name string) (backend.Backend, error) {
	if cfg != nil {
		spec, ok := cfg.Backends[name]
		if ok {
			return backend.Default.Build(spec.Type, name, spec.Raw)
		}
	}
	if backend.Default.Has(name) {
		return backend.Default.Build(name, name, nil)
	}
	return nil, userErr("backend %q not found in config and not a registered type (registered types: %v)",
		name, backend.Default.Types())
}

func buildWriter(cfg *config.Config, noSidecar bool) *output.Writer {
	w := &output.Writer{}
	if cfg != nil {
		w.Directory = config.ExpandPath(cfg.Output.Directory)
		w.NameTemplate = cfg.Output.Naming
		w.WriteSidecar = cfg.Output.WriteMetadataJSON
	}
	if w.Directory == "" {
		w.Directory = "."
	}
	if noSidecar {
		w.WriteSidecar = false
	}
	return w
}

func seedFromMetadata(meta map[string]any, fallback int64) int64 {
	if v, ok := meta["seed"]; ok {
		switch n := v.(type) {
		case int64:
			return n
		case int:
			return int64(n)
		case float64:
			return int64(n)
		}
	}
	return fallback
}

func extFromMime(mime string) string {
	switch mime {
	case "image/png", "":
		return "png"
	case "image/jpeg":
		return "jpg"
	case "image/webp":
		return "webp"
	}
	return "bin"
}