majordomo/provider/llamaswap/mediautil.go

// mediautil.go implements the imagegen.UpscaleProvider,
// imagegen.BackgroundRemovalProvider, and videogen.InterpolationProvider
// surfaces against upstreams reached through llama-swap's /upstream
// passthrough (ADR-0020):
//
//	upscale      POST /upstream/<id>/v1/upscale       (mediautils shim)
//	background   POST /upstream/<id>/api/remove       (rembg server)
//	interpolate  POST /upstream/<id>/v1/interpolate   (mediautils shim)
//
// All three are one-file multipart in, raw bytes out.
package llamaswap

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"strings"

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

// maxImageResponseBytes caps the raw-image bodies from the upscale and
// background-removal endpoints. The 64MB maxResponseBytes is a JSON cap;
// a 4x-upscaled PNG legitimately exceeds it. Bounded (not video-sized)
// because a single still image past this is an upstream bug, not data.
const maxImageResponseBytes = 256 << 20

// --- upscale ---

// UpscaleModel implements imagegen.UpscaleProvider against the mediautils
// shim's POST /v1/upscale. The id selects which upstream llama-swap loads.
func (p *Provider) UpscaleModel(id string, opts ...imagegen.UpscaleModelOption) (imagegen.Upscaler, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = imagegen.ApplyUpscaleModelOptions(opts)
	return &upscaleModel{p: p, id: id}, nil
}

type upscaleModel struct {
	p  *Provider
	id string
}

// Upscale implements imagegen.Upscaler.
func (m *upscaleModel) Upscale(ctx context.Context, req imagegen.UpscaleRequest, opts ...imagegen.UpscaleOption) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Image.Data) == 0 {
		return nil, fmt.Errorf("%w: upscale requires an image", llm.ErrUnsupported)
	}
	if req.Scale != 0 && req.Scale != 2 && req.Scale != 4 {
		return nil, fmt.Errorf("%w: upscale scale must be 2 or 4, got %d", llm.ErrUnsupported, req.Scale)
	}
	path, err := upstreamPath(m.id, "/v1/upscale")
	if err != nil {
		return nil, err
	}
	scale := ""
	if req.Scale != 0 {
		scale = strconv.Itoa(req.Scale)
	}
	body, contentType, err := buildMultipart("build upscale form",
		filePart{field: "file", filename: "image.png", data: req.Image.Data},
		[]formField{{"scale", scale, false}})
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
	if err != nil {
		return nil, err
	}
	return singleImageResult(m.p.name, m.id, "upscale", raw, respType)
}

// --- background removal ---

// BackgroundRemovalModel implements imagegen.BackgroundRemovalProvider
// against a rembg server's POST /api/remove.
func (p *Provider) BackgroundRemovalModel(id string, opts ...imagegen.BackgroundRemovalModelOption) (imagegen.BackgroundRemover, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = imagegen.ApplyBackgroundRemovalModelOptions(opts)
	return &backgroundModel{p: p, id: id}, nil
}

type backgroundModel struct {
	p  *Provider
	id string
}

// RemoveBackground implements imagegen.BackgroundRemover. rembg's `model`
// form field is the request's Net (its internal network); the llama-swap
// model id only picks the upstream. `om=true` returns the black/white
// foreground mask instead of the RGBA cutout.
func (m *backgroundModel) RemoveBackground(ctx context.Context, req imagegen.BackgroundRemovalRequest, opts ...imagegen.BackgroundRemovalOption) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Image.Data) == 0 {
		return nil, fmt.Errorf("%w: background removal requires an image", llm.ErrUnsupported)
	}
	path, err := upstreamPath(m.id, "/api/remove")
	if err != nil {
		return nil, err
	}
	om := ""
	if req.OnlyMask {
		om = "true"
	}
	body, contentType, err := buildMultipart("build background-removal form",
		filePart{field: "file", filename: "image.png", data: req.Image.Data},
		[]formField{
			{"model", req.Net, false},
			{"om", om, false},
		})
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
	if err != nil {
		return nil, err
	}
	return singleImageResult(m.p.name, m.id, "background removal", raw, respType)
}

// singleImageResult wraps one raw image body into an imagegen.Result.
// Acceptance requires positive evidence of image-ness — sniffed magic
// bytes or a declared image/* Content-Type — so a proxy error page with
// an empty Content-Type can never become "the image" (sniffImageMIME's
// PNG default is a labelling fallback, not a validator).
func singleImageResult(provider, model, verb string, raw []byte, contentType string) (*imagegen.Result, error) {
	if len(raw) == 0 {
		return nil, &llm.APIError{Provider: provider, Model: model, Message: verb + " response contained no image"}
	}
	sniffed := http.DetectContentType(raw)
	declared := strings.TrimSpace(contentType)
	if !strings.HasPrefix(sniffed, "image/") && !strings.HasPrefix(declared, "image/") {
		return nil, &llm.APIError{Provider: provider, Model: model,
			Message: fmt.Sprintf("%s response is not an image (Content-Type %q, sniffed %q)", verb, declared, sniffed)}
	}
	mimeType := sniffed
	if !strings.HasPrefix(mimeType, "image/") {
		mimeType = declared
	}
	return &imagegen.Result{Images: []llm.ImagePart{{MIME: mimeType, Data: raw}}}, nil
}

// --- frame interpolation ---

// InterpolatorModel implements videogen.InterpolationProvider against the
// mediautils shim's POST /v1/interpolate.
func (p *Provider) InterpolatorModel(id string, opts ...videogen.InterpolatorModelOption) (videogen.Interpolator, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = videogen.ApplyInterpolatorModelOptions(opts)
	return &interpolatorModel{p: p, id: id}, nil
}

type interpolatorModel struct {
	p  *Provider
	id string
}

// Interpolate implements videogen.Interpolator. The response body IS the
// encoded clip (same contract as /v1/videos/sync), hence the video-sized
// response cap and the same MIME resolution rules.
func (m *interpolatorModel) Interpolate(ctx context.Context, req videogen.InterpolateRequest, opts ...videogen.InterpolateOption) (*videogen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Video.Data) == 0 {
		return nil, fmt.Errorf("%w: interpolation requires a video", llm.ErrUnsupported)
	}
	if req.Multi != 0 && req.Multi != 2 && req.Multi != 4 {
		return nil, fmt.Errorf("%w: interpolation multi must be 2 or 4, got %d", llm.ErrUnsupported, req.Multi)
	}
	if req.TargetFPS < 0 {
		return nil, fmt.Errorf("%w: target fps must be >= 0, got %d", llm.ErrUnsupported, req.TargetFPS)
	}
	path, err := upstreamPath(m.id, "/v1/interpolate")
	if err != nil {
		return nil, err
	}
	multi := ""
	if req.Multi != 0 {
		multi = strconv.Itoa(req.Multi)
	}
	mode := ""
	targetFPS := ""
	if req.SlowMo {
		mode = "slowmo"
	} else if req.TargetFPS > 0 {
		targetFPS = strconv.Itoa(req.TargetFPS)
	}
	body, contentType, err := buildMultipart("build interpolate form",
		filePart{field: "file", filename: "video.mp4", data: req.Video.Data},
		[]formField{
			{"multi", multi, false},
			{"mode", mode, false},
			{"target_fps", targetFPS, false},
		})
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxVideoResponseBytes)
	if err != nil {
		return nil, err
	}
	return singleVideoResult(m.p.name, m.id, "interpolate", raw, respType)
}