trainer

type AIToolkit struct{}

type Capabilities struct {
	Device     string // mps | cpu | cuda
	DeviceName string
	MemoryGB   int
	Version    string // backend version (best effort)
}

type Config struct {
	Home   string
	Python string
}

type GenerateProgress struct {
	Raw    string
	Loaded bool // true once the model has finished loading and rendering begins
}

type GenerateRequest struct {
	Name           string
	Base           string
	BaseCheckpoint string   // name_or_path (the base model dir, as in training)
	LoRAPath       string   // trained .safetensors to fuse
	Prompts        []string // inline prompts (ignored when PromptFile is set)
	PromptFile     string   // path to a newline-delimited prompts file
	Negative       string
	Width          int
	Height         int
	Steps          int     // diffusion steps (sample_steps)
	Guidance       float64 // guidance scale
	Seed           int     // -1 = random per image
	Count          int     // images per prompt (num_repeats)
	Sampler        string  // "" → flowmatch for FLUX, else ddpm
	Precision      string  // dtype, e.g. bf16
	Quantize       bool
	Device         string
	OutputDir      string
}

type GenerateResult struct {
	Images []string
}

type Metrics struct {
	FinalLoss       float64  `json:"final_loss"`
	StepsCompleted  int      `json:"steps_completed"`
	DurationSeconds int      `json:"duration_seconds"`
	PeakMemoryBytes int64    `json:"peak_memory_bytes"`
	Checkpoints     []string `json:"checkpoints"`
	Samples         []string `json:"samples"`
}

type PerfPlan struct {
	CacheLatents        bool // VAE-encode images once to disk, then free the VAE
	CacheTextEmbeddings bool // encode captions once, then unload the text encoder
	DisableSampling     bool // skip sampling, including the pre-train baseline sample
}

type Plan struct {
	Req        Request
	Trainer    string
	Device     string
	DeviceName string
	MemoryGB   int
	Steps      int
	ConfigPath string
	OutputPath string // final collected .safetensors path
	Quantize   bool
}

type Progress struct {
	Step       int
	TotalSteps int
	Loss       float64
	Raw        string
}

type Request struct {
	Name           string
	Base           string
	BaseCheckpoint string
	Trigger        string
	DatasetDir     string
	CaptionMode    string
	// CaptionsHaveTrigger: the per-image .txt captions already include the
	// trigger (we generated them), so the trainer must not inject it again.
	CaptionsHaveTrigger bool
	Profile             profile.Profile
	Device              string
	CacheDir            string // .loradex/cache/<run-id>
	OutputDir           string // version dir to collect results into
	OutputFile          string // final .safetensors filename
	RawConfig           string // --config escape hatch (path); skips generation
	Samples             int
	RunID               string
}

type Result struct {
	WeightsPath string
	NetworkRank int
	NetworkDim  int
	Metrics     Metrics
	Stopped     bool // ctrl-c → resumable
}

type Trainer interface {
	Name() string
	Detect(cfg Config) (Capabilities, error)
	Plan(req Request) (Plan, error)
	Train(ctx context.Context, plan Plan, onProgress func(Progress)) (Result, error)
}