majordomo/provider/llamaswap/faceswap.go

// faceswap.go implements imagegen.FaceSwapper against the InsightFace shim
// (buffalo_l + inswapper_128) reached through llama-swap's /upstream
// passthrough (ADR-0024):
//
//	POST /upstream/<id>/v1/faces     multipart file            -> JSON
//	POST /upstream/<id>/v1/faceswap  multipart target, source  -> PNG
package llamaswap

import (
	"context"
	"encoding/json"
	"fmt"
	"mime"
	"net/http"
	"strconv"
	"strings"

	"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// maxFaceSwapResponseBytes bounds the returned PNG. Generous: the shim echoes
// the target's dimensions, and a 4K photo round-trips as a large lossless PNG.
const maxFaceSwapResponseBytes = 64 << 20

// FaceSwapModel implements the face-transfer surface. The id selects which
// upstream llama-swap loads.
func (p *Provider) FaceSwapModel(id string) (imagegen.FaceSwapper, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	return &faceSwapModel{p: p, id: id}, nil
}

type faceSwapModel struct {
	p  *Provider
	id string
}

// facesResponse mirrors the shim's /v1/faces body.
type facesResponse struct {
	Count int `json:"count"`
	Faces []struct {
		Index  int     `json:"index"`
		Box    []int   `json:"box"`
		Score  float64 `json:"score"`
		Width  int     `json:"width"`
		Height int     `json:"height"`
	} `json:"faces"`
}

// ListFaces implements imagegen.FaceSwapper.
func (m *faceSwapModel) ListFaces(ctx context.Context, img imagegen.Image) ([]imagegen.DetectedFace, error) {
	if len(img.Data) == 0 {
		return nil, fmt.Errorf("%w: face detection requires image bytes", llm.ErrUnsupported)
	}
	path, err := upstreamPath(m.id, "/v1/faces")
	if err != nil {
		return nil, err
	}
	body, contentType, err := buildMultipart("build faces form",
		filePart{field: "file", filename: imageFilename(img.MIME, "image"), data: img.Data}, nil)
	if err != nil {
		return nil, err
	}
	raw, _, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxFaceSwapResponseBytes)
	if err != nil {
		return nil, err
	}
	var parsed facesResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
			Message: fmt.Sprintf("faces response is not JSON: %s", truncateForError(raw))}
	}
	out := make([]imagegen.DetectedFace, 0, len(parsed.Faces))
	for _, f := range parsed.Faces {
		df := imagegen.DetectedFace{Index: f.Index, Score: f.Score, Width: f.Width, Height: f.Height}
		// A short box would silently index out of range below; treat a
		// malformed entry as a protocol error rather than zero-filling it,
		// because a wrong box sends the caller at the wrong face.
		if len(f.Box) != 4 {
			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
				Message: fmt.Sprintf("face %d has a %d-element box, want 4", f.Index, len(f.Box))}
		}
		copy(df.Box[:], f.Box)
		out = append(out, df)
	}
	return out, nil
}

// FaceSwap implements imagegen.FaceSwapper. The endpoint always answers PNG.
func (m *faceSwapModel) FaceSwap(ctx context.Context, req imagegen.FaceSwapRequest, opts ...imagegen.FaceSwapOption) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Target.Data) == 0 {
		return nil, fmt.Errorf("%w: face swap requires a target image", llm.ErrUnsupported)
	}
	if len(req.Source.Data) == 0 {
		return nil, fmt.Errorf("%w: face swap requires a source image", llm.ErrUnsupported)
	}
	if req.Index != nil && *req.Index < 0 {
		return nil, fmt.Errorf("%w: face index must be >= 0, got %d", llm.ErrUnsupported, *req.Index)
	}
	path, err := upstreamPath(m.id, "/v1/faceswap")
	if err != nil {
		return nil, err
	}

	var fields []formField
	if req.All {
		fields = append(fields, formField{key: "all", value: "true", required: true})
	} else if req.Index != nil {
		// Only sent when NOT swapping all: the shim ignores index under
		// all=true, and sending both would imply a precedence the caller
		// cannot see.
		fields = append(fields, formField{key: "index", value: strconv.Itoa(*req.Index), required: true})
	}

	body, contentType, err := buildMultipartFiles("build faceswap form",
		[]filePart{
			{field: "target", filename: imageFilename(req.Target.MIME, "target"), data: req.Target.Data},
			{field: "source", filename: imageFilename(req.Source.MIME, "source"), data: req.Source.Data},
		}, fields)
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxFaceSwapResponseBytes)
	if err != nil {
		return nil, err
	}
	if len(raw) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "face swap response contained no image"}
	}
	mime := sniffImageMIME(raw)
	if respType != "" && !isImageContentType(respType) {
		// A JSON error body sniffs as text, not image — say what came back
		// rather than handing the caller bytes that are not a picture.
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
			Message: fmt.Sprintf("face swap response is not an image (Content-Type %q): %s", respType, truncateForError(raw))}
	}
	return &imagegen.Result{Images: []llm.ImagePart{{MIME: mime, Data: raw}}}, nil
}

// imageFilename picks a multipart filename for an image part. The shim reads
// bytes, not names, but a plausible extension keeps server-side sniffing and
// request logs honest. base distinguishes the parts of a multi-file form
// ("target"/"source") so a log line says which one was malformed.
func imageFilename(mimeType, base string) string {
	if base == "" {
		base = "image"
	}
	mt := strings.ToLower(strings.TrimSpace(mimeType))
	if parsed, _, err := mime.ParseMediaType(mt); err == nil {
		mt = parsed
	}
	switch mt {
	case "image/jpeg", "image/jpg":
		return base + ".jpg"
	case "image/webp":
		return base + ".webp"
	case "image/gif":
		return base + ".gif"
	case "image/bmp":
		return base + ".bmp"
	default:
		// PNG is the safe default: every caller in this repo either sends PNG
		// or sends something the decoder identifies by magic bytes anyway.
		return base + ".png"
	}
}

// isImageContentType reports whether a response Content-Type is an image.
// Used to tell a returned picture apart from a JSON error body, which would
// otherwise be handed back as "image" bytes.
func isImageContentType(contentType string) bool {
	return mimeFromContentType(contentType, "image/") != ""
}