135 lines
3.4 KiB
Go
135 lines
3.4 KiB
Go
package ml
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"fn-registry/functions/core"
|
|
)
|
|
|
|
// SdcliProgressCallback es una funcion llamada cada vez que se parsea una linea
|
|
// de progreso del proceso sd. Puede ser nil.
|
|
type SdcliProgressCallback func(p SdcliProgress)
|
|
|
|
// SdcliGenerate ejecuta el binario sd para generar una imagen y escribe el
|
|
// resultado en outPath.
|
|
//
|
|
// Flujo:
|
|
// 1. Construye args con GenconfigToSdcliArgs(cfg) + ["-o", outPath].
|
|
// 2. Lanza el proceso via SubprocessStream.
|
|
// 3. Goroutine interna lee eventos: lineas stderr se pasan a SdcliParseProgress;
|
|
// si onProgress != nil y hay progreso reconocible, llama onProgress(p).
|
|
// 4. Espera el resultado. ExitCode != 0 => error con las ultimas 10 lineas de stderr.
|
|
// 5. Lee outPath y retorna ImageGenResult con bytes, meta y duration_ms.
|
|
//
|
|
// ctx controla el timeout y cancelacion: se pasa directamente a SubprocessStream,
|
|
// que maneja SIGTERM -> grace 2s -> SIGKILL.
|
|
func SdcliGenerate(
|
|
ctx context.Context,
|
|
bin SdcliBinary,
|
|
cfg GenerationConfig,
|
|
outPath string,
|
|
onProgress SdcliProgressCallback,
|
|
) (ImageGenResult, error) {
|
|
start := time.Now()
|
|
|
|
args := GenconfigToSdcliArgs(cfg)
|
|
args = append(args, "-o", outPath)
|
|
|
|
events, results := core.SubprocessStream(ctx, bin.Path, args, nil, nil)
|
|
|
|
// Consumir eventos en goroutine: parsear progreso y acumular stderr para
|
|
// mensajes de error utiles.
|
|
type collectedStderr struct {
|
|
lines []string
|
|
}
|
|
stderrCh := make(chan collectedStderr, 1)
|
|
|
|
go func() {
|
|
var stderrLines []string
|
|
for ev := range events {
|
|
if ev.Stream == "stderr" {
|
|
stderrLines = append(stderrLines, ev.Line)
|
|
if onProgress != nil {
|
|
if p, ok := SdcliParseProgress(ev.Line); ok {
|
|
onProgress(p)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
stderrCh <- collectedStderr{lines: stderrLines}
|
|
}()
|
|
|
|
res := <-results
|
|
collected := <-stderrCh
|
|
|
|
if res.Err != nil {
|
|
return ImageGenResult{}, fmt.Errorf("sdcli subprocess: %w", res.Err)
|
|
}
|
|
|
|
if res.ExitCode != 0 {
|
|
tail := stderrTail(collected.lines, 10)
|
|
return ImageGenResult{}, fmt.Errorf(
|
|
"sdcli exited with code %d:\n%s",
|
|
res.ExitCode, tail,
|
|
)
|
|
}
|
|
|
|
imageBytes, err := os.ReadFile(outPath)
|
|
if err != nil {
|
|
return ImageGenResult{}, fmt.Errorf("sdcli: reading output image %q: %w", outPath, err)
|
|
}
|
|
|
|
durationMs := time.Since(start).Milliseconds()
|
|
|
|
meta := map[string]any{
|
|
"backend": "sdcli",
|
|
"binary_path": bin.Path,
|
|
"model": cfg.Model.Name,
|
|
"seed": cfg.Seed,
|
|
"steps": cfg.Steps,
|
|
"cfg_scale": cfg.CfgScale,
|
|
"sampler": cfg.Sampler,
|
|
"width": cfg.Width,
|
|
"height": cfg.Height,
|
|
}
|
|
if bin.Version != "" {
|
|
meta["version"] = bin.Version
|
|
}
|
|
if cfg.Model.Path != "" {
|
|
meta["model_path"] = cfg.Model.Path
|
|
}
|
|
if cfg.Model.ModelType != "" {
|
|
meta["model_type"] = cfg.Model.ModelType
|
|
}
|
|
if cfg.Model.Quantization != "" {
|
|
meta["quantization"] = cfg.Model.Quantization
|
|
}
|
|
if len(cfg.Loras) > 0 {
|
|
loras := make([]string, len(cfg.Loras))
|
|
for i, l := range cfg.Loras {
|
|
loras[i] = l.Path + ":" + strconv.FormatFloat(l.Weight, 'f', -1, 64)
|
|
}
|
|
meta["loras"] = strings.Join(loras, ",")
|
|
}
|
|
|
|
return ImageGenResult{
|
|
ImageBytes: imageBytes,
|
|
Format: "png",
|
|
Meta: meta,
|
|
DurationMs: durationMs,
|
|
}, nil
|
|
}
|
|
|
|
// stderrTail retorna las ultimas n lineas de lines, unidas con newline.
|
|
func stderrTail(lines []string, n int) string {
|
|
if len(lines) <= n {
|
|
return strings.Join(lines, "\n")
|
|
}
|
|
return strings.Join(lines[len(lines)-n:], "\n")
|
|
}
|