mAi
b282325663
Implements the Replicate API backend (FLUX schnell / FLUX dev) per ImaGen issue #3: - internal/backend/replicate.go — Backend adapter. Supports model refs as "owner/name" (uses /v1/models/{owner}/{name}/predictions) and "owner/name:hash" (uses /v1/predictions with explicit version). Polls /v1/predictions/{id} every 500ms with model-aware timeout (60s schnell, 120s dev). Resilience: 401 names api_token_env, 429 with exp backoff up to 3 retries (honours Retry-After), 5xx retries once, image download retries once on transient failure. - internal/backend/replicate_pricing.go — hardcoded per-image USD rates for known FLUX models, snapshotted from replicate.com/pricing with a refresh TODO. - internal/backend/replicate_test.go — mocked-HTTP unit tests covering happy path (model + version-pinned), 401, 429 retry policy, failed prediction, poll timeout, image-download retry, ctx cancel, BackendOpts passthrough, default_steps, aspect-ratio reduction, sha256 prompt hash. - internal/usage/usage.go — Supabase REST sink + read-side query for mai.imagen_usage. Adapter writes are best-effort: failures warn but the image still lands. - cmd/imagen/usage.go — `imagen usage [--since DATE] [--raw]` reads the table and prints a tab-aligned grouped or raw table with totals. - cmd/imagen/backends.go — instances of type=replicate now report "ok" or "not configured (set REPLICATE_API_TOKEN)" depending on env. - internal/config/config.go — sample adds flux-schnell-replicate + flux-dev-replicate; default_backend stays flux-schnell-local. - Supabase migration mai.imagen_usage (id, created_at, backend, model, seed, prompt_hash, latency_ms, cost_usd_estimate, caller) + indexes on (created_at DESC) and (caller). The raw prompt is never stored. Caller identity resolves from MAI_FROM_ID, then the tmux pane's @mai-name option, mirroring the maimcp identity logic. Prompt hash is sha256 of the user-facing prompt; raw prompt never reaches the table.
81 lines
2.1 KiB
Go
81 lines
2.1 KiB
Go
// Command imagen is the model-agnostic image-generation CLI. It dispatches
|
|
// `generate`, `backends`, and `config` subcommands against backends that
|
|
// register themselves at package init time.
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
|
|
_ "mgit.msbls.de/m/ImaGen/internal/backend"
|
|
)
|
|
|
|
const helpText = `imagen — model-agnostic image generation
|
|
|
|
Usage:
|
|
imagen generate <prompt> [flags] generate one image
|
|
imagen backends list registered backend types
|
|
imagen config init print a sample imagen.yaml on stdout
|
|
imagen config validate validate the active config
|
|
imagen serve [--addr :8080] (stub) start the HTTP server
|
|
imagen usage [--since DATE] show cost-tracking rows
|
|
imagen version print version
|
|
imagen help show this help
|
|
|
|
Run "imagen <subcommand> --help" for subcommand-specific flags.
|
|
`
|
|
|
|
// Version is overridable at link time via -ldflags '-X main.Version=...'.
|
|
var Version = "dev"
|
|
|
|
func main() {
|
|
if len(os.Args) < 2 {
|
|
fmt.Fprint(os.Stderr, helpText)
|
|
os.Exit(2)
|
|
}
|
|
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
|
defer cancel()
|
|
|
|
args := os.Args[2:]
|
|
var err error
|
|
switch os.Args[1] {
|
|
case "generate":
|
|
err = runGenerate(ctx, args)
|
|
case "backends":
|
|
err = runBackends(args)
|
|
case "config":
|
|
err = runConfig(args)
|
|
case "serve":
|
|
err = runServe(args)
|
|
case "usage":
|
|
err = runUsage(ctx, args)
|
|
case "version", "-v", "--version":
|
|
fmt.Println(Version)
|
|
case "help", "-h", "--help":
|
|
fmt.Print(helpText)
|
|
default:
|
|
fmt.Fprintf(os.Stderr, "imagen: unknown subcommand %q\n\n%s", os.Args[1], helpText)
|
|
os.Exit(2)
|
|
}
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "imagen:", err)
|
|
var u *userError
|
|
if errors.As(err, &u) {
|
|
os.Exit(2)
|
|
}
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
// userError signals "user did the wrong thing" so we exit 2 rather than 1.
|
|
type userError struct{ msg string }
|
|
|
|
func (u *userError) Error() string { return u.msg }
|
|
func userErr(format string, a ...any) error {
|
|
return &userError{msg: fmt.Sprintf(format, a...)}
|
|
}
|