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/prompt"
)

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
	)
	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.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
	}

	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)
	}
	return nil
}

// 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
}

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"
}