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}
		// 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)
	}
	// Only when it will actually be sent: under All the index is documented
	// as ignored, so rejecting a negative one there would fail a request that
	// is perfectly well formed.
	if !req.All && 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, respHdr, err := m.p.doRawHeaders(ctx, http.MethodPost, path, m.id, contentType, body, maxFaceSwapResponseBytes)
	if err != nil {
		return nil, err
	}
	respType := respHdr.Get("Content-Type")
	if len(raw) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "face swap response contained no image"}
	}
	// Validate the BYTES, not the header. sniffImageMIME falls back to
	// image/png when detection is inconclusive, so trusting it here would
	// label a JSON error body as a PNG and return it as a successful image —
	// and a header check alone misses the case where the response carries no
	// Content-Type at all. The shim answers JSON on a semantic miss (no face
	// found), which is exactly the body that would sail through.
	detected := http.DetectContentType(raw)
	if !strings.HasPrefix(detected, "image/") {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
			Message: fmt.Sprintf("face swap response is not an image (sniffed %q, Content-Type %q): %s",
				detected, respType, truncateForError(raw))}
	}
	// Prefer the server's own label when it is an image type (it knows
	// subtypes the sniffer does not), else what the bytes actually are.
	mimeType := detected
	if hdr := mimeFromContentType(respType, "image/"); hdr != "" {
		mimeType = hdr
	}
	return &imagegen.Result{
		Images:       []llm.ImagePart{{MIME: mimeType, Data: raw}},
		SwappedFaces: parseSwapReport(respHdr.Get("X-Swap-Report")),
	}, nil
}

// swapReport mirrors the shim's X-Swap-Report header.
type swapReport struct {
	Image []int `json:"image"`
	Faces []struct {
		Index              int      `json:"index"`
		Size               []int    `json:"size"`
		Yaw                *float64 `json:"yaw"`
		IdentitySimilarity *float64 `json:"identity_similarity"`
	} `json:"faces"`
}

// parseSwapReport decodes the measured outcome. A missing or malformed header
// yields nil rather than an error: an older shim does not send it, and a swap
// that produced a good image must not fail because the diagnostics alongside
// it were unreadable.
func parseSwapReport(header string) []imagegen.SwappedFace {
	header = strings.TrimSpace(header)
	if header == "" {
		return nil
	}
	var rep swapReport
	if err := json.Unmarshal([]byte(header), &rep); err != nil {
		return nil
	}
	out := make([]imagegen.SwappedFace, 0, len(rep.Faces))
	for _, f := range rep.Faces {
		sf := imagegen.SwappedFace{
			Index:              f.Index,
			Yaw:                f.Yaw,
			IdentitySimilarity: f.IdentitySimilarity,
		}
		if len(f.Size) == 2 {
			sf.Width, sf.Height = f.Size[0], f.Size[1]
		}
		if len(rep.Image) == 2 {
			sf.ImageWidth, sf.ImageHeight = rep.Image[0], rep.Image[1]
		}
		out = append(out, sf)
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

// 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.
//
// initImageFilename (video.go) is this function with base fixed to "frame"
// and delegates here — two copies of one extension table is how they drift.
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"
	}
}