feat: wave-3 image + document surfaces — segmentation, colorize, face restore, ocr (ADR-0023)
CI / Tidy (pull_request) Successful in 9m25s
CI / Build & Test (pull_request) Successful in 9m48s
Gadfly review (reusable) / review (pull_request) Successful in 47m4s
Adversarial Review (Gadfly) / review (pull_request) Successful in 47m4s

- imagegen.Segmenter/SegmentationProvider: prompted mask via
  POST /upstream/<id>/v1/segment (file, prompt[, threshold], output=mask);
  white = prompted region, EditRequest.Mask polarity.
- imagegen.Colorizer/ColorizeProvider: POST /upstream/<id>/v1/colorize.
- imagegen.FaceRestorer/FaceRestoreProvider:
  POST /upstream/<id>/v1/restore_faces (upscale 1|2).
- New ocr leaf package (Request/Page/Result, Recognize) + llamaswap
  OCRModel: POST /upstream/<id>/v1/ocr (file[, langs, max_pages]),
  tolerant per-page decode (join lines when page text absent), Raw
  escape hatch.
- httptest contract tests per surface; ADR-0023; ADR index backfilled
  (0020-0022 rows were missing).

Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
2026-07-16 16:54:30 -04:00
co-authored by Claude Fable 5
parent cd43009672
commit e6987f54b2
11 changed files with 1155 additions and 0 deletions
+59
View File
@@ -0,0 +1,59 @@
# ADR-0023: Wave-3 image + document surfaces (segmentation, colorize, face restore, OCR)
Status: Accepted (2026-07-16)
## Context
The llama-swap host is gaining four wave-3 capabilities (spec:
mort docs/specs/2026-07-16-llamaswap-wave3.md): promptable segmentation
(GroundingDINO + SAM 2.1), photo colorization (DDColor), face restoration
(GFPGAN), and document OCR (Surya). All ride the `/upstream/<model>/<path>`
passthrough (ADR-0020); none of their native APIs are OpenAI-shaped.
## Decision
1. **Grow `imagegen` with three optional interfaces**, ADR-0016→0022
conventions (functional options, zero value = backend default, bytes-only
I/O, provider mints model):
- `imagegen.Segmenter` / `SegmentationProvider` — `SegmentationRequest
{Image, Prompt, Threshold}` → one grayscale mask image, WHITE = the
prompted region. Same polarity as `EditRequest.Mask` (white = repaint),
so the mask feeds inpainting directly; cutouts are derived client-side
(mask-as-alpha), one host call serving both. The client always sends
`output=mask` — the shim's `cutout`/`boxes` modes are not exposed.
- `imagegen.Colorizer` / `ColorizeProvider` — `ColorizeRequest{Image}`,
no knobs (the reference backend takes none; options reserved).
- `imagegen.FaceRestorer` / `FaceRestoreProvider` —
`FaceRestoreRequest{Image, Upscale}` (1 or 2, 0 = backend default).
2. **New `ocr` leaf package** rather than a method on an existing surface:
OCR is document-shaped (multi-page, PDFs), not image-generation-shaped.
`Request{Document, MIME, Filename, Languages, MaxPages}` →
`Result{Text, Pages []Page{Number, Text}, Raw}`. `Result.Text` is the
pages joined with a blank line; the per-line bbox/confidence/layout
detail stays in `Raw` (json.RawMessage) — exact text is the contract,
geometry is the escape hatch.
3. **provider/llamaswap wire shapes** (pinned by the netherstorm image
builds; host smoke tests are the drift defence):
- `POST /upstream/<id>/v1/segment` multipart `file`,`prompt`
[,`threshold`],`output=mask` → mask PNG.
- `POST /upstream/<id>/v1/colorize` multipart `file` → PNG.
- `POST /upstream/<id>/v1/restore_faces` multipart `file`[,`upscale`] → PNG.
- `POST /upstream/<id>/v1/ocr` multipart `file`[,`langs` (comma-joined),
`max_pages`] → JSON `{pages:[{number,text,lines,layout}]}`. The decode
is tolerant: a page without aggregate `text` joins its line texts; a
missing page `number` defaults to position. Zero pages is an error
(a blank page still arrives as a page), mirroring the "no transcript"
honesty rule.
4. **Binary success bodies are validated before wrapping** (ADR-0020 rule):
the three image surfaces reuse `singleImageResult` (positive evidence of
image-ness required), OCR requires decodable JSON.
## Consequences
- imagegen grows from five to eight optional surfaces; consumers
type-assert or use the provider methods directly, as before.
- `ocr` is the seventh leaf media package (imagegen, audio, videogen,
meshgen, musicgen, embeddings, ocr); the conventions have held across all
of them.
- PDF handling lives host-side (the shim rasterizes via pypdfium2);
majordomo ships bytes and never needs a PDF dependency.
+4
View File
@@ -23,3 +23,7 @@ One decision per file, append-only; supersede rather than rewrite.
| [0017](0017-audio-interfaces.md) | audio — canonical speech synthesis + transcription interfaces | Accepted |
| [0018](0018-imagegen-editor.md) | imagegen.Editor — image-to-image as a separate optional interface | Accepted |
| [0019](0019-videogen-interface.md) | videogen — canonical video-generation surface | Accepted |
| [0020](0020-upstream-passthrough-media-surfaces.md) | Upstream-passthrough media surfaces (mask, upscale, background removal, interpolation, diarization, meshgen) | Accepted |
| [0021](0021-musicgen-interface.md) | musicgen — blocking Generate over an async job queue | Accepted |
| [0022](0022-embeddings-rerank-interface.md) | embeddings + rerank interface | Accepted |
| [0023](0023-image-doc-surfaces.md) | Wave-3 image + document surfaces (segmentation, colorize, face restore, OCR) | Accepted |
+54
View File
@@ -0,0 +1,54 @@
package imagegen
import "context"
// ColorizeRequest asks a colorization backend (DDColor style) to add color to
// a grayscale/faded photo (ADR-0023).
type ColorizeRequest struct {
// Image is the image to colorize. Required.
Image Image
}
// ColorizeOption mutates a ColorizeRequest before it is sent. Reserved for
// future request settings (the reference backend takes no parameters).
type ColorizeOption func(*ColorizeRequest)
// Apply returns a copy of the request with all options applied.
func (r ColorizeRequest) Apply(opts ...ColorizeOption) ColorizeRequest {
for _, opt := range opts {
opt(&r)
}
return r
}
// Colorizer adds color to grayscale images.
type Colorizer interface {
// Colorize returns the colorized image as a one-image Result.
Colorize(ctx context.Context, req ColorizeRequest, opts ...ColorizeOption) (*Result, error)
}
// ColorizeModelOption configures a Colorizer at construction time. Reserved
// for future per-model settings.
type ColorizeModelOption func(*ColorizeModelConfig)
// ColorizeModelConfig carries per-model construction settings.
type ColorizeModelConfig struct{}
// ApplyColorizeModelOptions folds options into a config.
func ApplyColorizeModelOptions(opts []ColorizeModelOption) ColorizeModelConfig {
var cfg ColorizeModelConfig
for _, opt := range opts {
opt(&cfg)
}
return cfg
}
// ColorizeProvider mints Colorizers bound to one backend.
type ColorizeProvider interface {
// Name is the registry identifier for the provider.
Name() string
// ColorizeModel returns a Colorizer bound to the given id (passed through
// to the backend verbatim; no catalog validation).
ColorizeModel(id string, opts ...ColorizeModelOption) (Colorizer, error)
}
+62
View File
@@ -0,0 +1,62 @@
package imagegen
import "context"
// FaceRestoreRequest asks a face-restoration backend (GFPGAN style) to repair
// degraded faces in a photo. Zero values mean "backend default" (ADR-0023).
type FaceRestoreRequest struct {
// Image is the image to restore. Required.
Image Image
// Upscale is the output enlargement factor (1 or 2 on the reference
// backend); 0 = backend default.
Upscale int
}
// FaceRestoreOption mutates a FaceRestoreRequest before it is sent.
type FaceRestoreOption func(*FaceRestoreRequest)
// WithFaceRestoreUpscale sets the output enlargement factor.
func WithFaceRestoreUpscale(n int) FaceRestoreOption {
return func(r *FaceRestoreRequest) { r.Upscale = n }
}
// Apply returns a copy of the request with all options applied.
func (r FaceRestoreRequest) Apply(opts ...FaceRestoreOption) FaceRestoreRequest {
for _, opt := range opts {
opt(&r)
}
return r
}
// FaceRestorer repairs degraded/blurry faces in photos.
type FaceRestorer interface {
// RestoreFaces returns the restored image as a one-image Result.
RestoreFaces(ctx context.Context, req FaceRestoreRequest, opts ...FaceRestoreOption) (*Result, error)
}
// FaceRestoreModelOption configures a FaceRestorer at construction time.
// Reserved for future per-model settings.
type FaceRestoreModelOption func(*FaceRestoreModelConfig)
// FaceRestoreModelConfig carries per-model construction settings.
type FaceRestoreModelConfig struct{}
// ApplyFaceRestoreModelOptions folds options into a config.
func ApplyFaceRestoreModelOptions(opts []FaceRestoreModelOption) FaceRestoreModelConfig {
var cfg FaceRestoreModelConfig
for _, opt := range opts {
opt(&cfg)
}
return cfg
}
// FaceRestoreProvider mints FaceRestorers bound to one backend.
type FaceRestoreProvider interface {
// Name is the registry identifier for the provider.
Name() string
// FaceRestoreModel returns a FaceRestorer bound to the given id (passed
// through to the backend verbatim; no catalog validation).
FaceRestoreModel(id string, opts ...FaceRestoreModelOption) (FaceRestorer, error)
}
+70
View File
@@ -0,0 +1,70 @@
package imagegen
import "context"
// SegmentationRequest asks a promptable-segmentation backend (GroundingDINO +
// SAM style) for the mask of a text-described region. Zero values mean
// "backend default" (ADR-0023).
type SegmentationRequest struct {
// Image is the image to segment. Required.
Image Image
// Prompt is the text description of the region to segment (e.g. "the red
// car"). Required — promptless segmentation is a different capability.
Prompt string
// Threshold is the detection confidence threshold in (0,1]; 0 = backend
// default.
Threshold float64
}
// SegmentationOption mutates a SegmentationRequest before it is sent.
type SegmentationOption func(*SegmentationRequest)
// WithSegmentationThreshold sets the detection confidence threshold.
func WithSegmentationThreshold(t float64) SegmentationOption {
return func(r *SegmentationRequest) { r.Threshold = t }
}
// Apply returns a copy of the request with all options applied.
func (r SegmentationRequest) Apply(opts ...SegmentationOption) SegmentationRequest {
for _, opt := range opts {
opt(&r)
}
return r
}
// Segmenter extracts a prompted region's mask from an image. The result is a
// single grayscale mask image where WHITE marks the prompted region — the
// same polarity as EditRequest.Mask (white = repaint), so it feeds inpainting
// directly; derive a cutout client-side by applying the mask as alpha.
type Segmenter interface {
// Segment returns the region mask as a one-image Result.
Segment(ctx context.Context, req SegmentationRequest, opts ...SegmentationOption) (*Result, error)
}
// SegmentationModelOption configures a Segmenter at construction time.
// Reserved for future per-model settings (mirrors ModelOption).
type SegmentationModelOption func(*SegmentationModelConfig)
// SegmentationModelConfig carries per-model construction settings.
type SegmentationModelConfig struct{}
// ApplySegmentationModelOptions folds options into a config.
func ApplySegmentationModelOptions(opts []SegmentationModelOption) SegmentationModelConfig {
var cfg SegmentationModelConfig
for _, opt := range opts {
opt(&cfg)
}
return cfg
}
// SegmentationProvider mints Segmenters bound to one backend.
type SegmentationProvider interface {
// Name is the registry identifier for the provider.
Name() string
// SegmentationModel returns a Segmenter bound to the given id (passed
// through to the backend verbatim; no catalog validation).
SegmentationModel(id string, opts ...SegmentationModelOption) (Segmenter, error)
}
+114
View File
@@ -0,0 +1,114 @@
// Package ocr is majordomo's canonical document text-recognition surface.
// Like imagegen/audio/videogen/musicgen, it is a deliberately separate leaf
// contract from the llm package (ADR-0023, following the ADR-0016→0022
// lineage: functional options, zero values = backend default, bytes-only
// I/O, Raw escape hatch). OCR is not chat-vision: it targets dedicated
// detection+recognition models (Surya style) that return exact per-page,
// per-line text rather than a model's paraphrase.
//
// The first implementation is provider/llamaswap, which posts the document
// to a Surya shim through the /upstream passthrough; the shim rasterizes
// PDFs itself, so Document may be an image (png/jpg/webp) or a PDF.
package ocr
import "context"
// Request is a text-recognition request. Zero values mean "backend default".
type Request struct {
// Document is the encoded document to recognize: an image (png/jpg/webp)
// or a PDF. Required. Carried as bytes, never a URL.
Document []byte
// MIME is the document MIME type (e.g. "image/png", "application/pdf");
// "" = let the backend sniff it.
MIME string
// Filename is the multipart filename hint some backends key their format
// detection on; "" derives one from MIME ("document.pdf") or falls back
// to "document".
Filename string
// Languages are ISO-639 hints for the recognizer (e.g. "en", "de");
// nil = backend default (auto/multilingual).
Languages []string
// MaxPages caps how many pages of a multi-page document are recognized;
// 0 = backend default (all pages, up to the backend's own ceiling).
MaxPages int
}
// Page is the recognized text of one page.
type Page struct {
// Number is the 1-based page number.
Number int
// Text is the page's recognized text.
Text string
}
// Result is the canonical text-recognition result.
type Result struct {
// Text is the full recognized text, pages joined in order.
Text string
// Pages are the per-page results in page order.
Pages []Page
// Raw is the provider-native response object (e.g. the per-line
// bbox/confidence/layout detail), an escape hatch for provider-specific
// fields. May be nil; never required for normal use.
Raw any
}
// Option mutates a Request before it is sent. Options passed to Recognize are
// applied to a copy of the request, so a Request value can be reused.
type Option func(*Request)
// WithLanguages sets the recognizer language hints.
func WithLanguages(langs ...string) Option {
return func(r *Request) { r.Languages = langs }
}
// WithMaxPages caps how many pages are recognized.
func WithMaxPages(n int) Option { return func(r *Request) { r.MaxPages = n } }
// Apply returns a copy of the request with all options applied. Providers
// call this once at the top of Recognize.
func (r Request) Apply(opts ...Option) Request {
for _, opt := range opts {
opt(&r)
}
return r
}
// Model recognizes text in documents.
type Model interface {
// Recognize extracts the document's text, page by page.
Recognize(ctx context.Context, req Request, opts ...Option) (*Result, error)
}
// ModelOption configures a Model at construction time. Reserved for future
// per-model settings.
type ModelOption func(*ModelConfig)
// ModelConfig carries per-model construction settings.
type ModelConfig struct{}
// ApplyModelOptions folds options into a config.
func ApplyModelOptions(opts []ModelOption) ModelConfig {
var cfg ModelConfig
for _, opt := range opts {
opt(&cfg)
}
return cfg
}
// Provider mints OCR models bound to one backend.
type Provider interface {
// Name is the registry identifier for the provider.
Name() string
// OCRModel returns a Model bound to the given id (passed through to the
// backend verbatim; no catalog validation).
OCRModel(id string, opts ...ModelOption) (Model, error)
}
+140
View File
@@ -0,0 +1,140 @@
// ocr.go implements ocr.Provider against a Surya-style shim reached through
// llama-swap's /upstream passthrough (ADR-0023):
//
// POST /upstream/<id>/v1/ocr multipart file[,langs,max_pages]
//
// The document may be an image or a PDF (the shim rasterizes PDFs itself).
// The response is per-page JSON: {pages:[{number,text,lines,layout}]}. The
// decode is tolerant — when a page carries no aggregate `text`, its line
// texts are joined instead — and the full payload survives in Result.Raw for
// callers that want bboxes/confidence/layout.
package llamaswap
import (
"context"
"encoding/json"
"fmt"
"mime"
"net/http"
"strconv"
"strings"
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
"gitea.stevedudenhoeffer.com/steve/majordomo/ocr"
)
// OCRModel implements ocr.Provider. The id selects which upstream llama-swap
// loads.
func (p *Provider) OCRModel(id string, opts ...ocr.ModelOption) (ocr.Model, error) {
if err := p.requireBaseURL(); err != nil {
return nil, err
}
_ = ocr.ApplyModelOptions(opts)
return &ocrModel{p: p, id: id}, nil
}
type ocrModel struct {
p *Provider
id string
}
// ocrResponse is the Surya shim's /v1/ocr shape (the subset this package
// relies on; per-line bbox/confidence and layout stay in Raw).
type ocrResponse struct {
Pages []struct {
Number int `json:"number"`
Text string `json:"text"`
Lines []struct {
Text string `json:"text"`
} `json:"lines"`
} `json:"pages"`
}
// Recognize implements ocr.Model.
func (m *ocrModel) Recognize(ctx context.Context, req ocr.Request, opts ...ocr.Option) (*ocr.Result, error) {
req = req.Apply(opts...)
if len(req.Document) == 0 {
return nil, fmt.Errorf("%w: ocr requires document bytes", llm.ErrUnsupported)
}
if req.MaxPages < 0 {
return nil, fmt.Errorf("%w: ocr max pages must be >= 0, got %d", llm.ErrUnsupported, req.MaxPages)
}
path, err := upstreamPath(m.id, "/v1/ocr")
if err != nil {
return nil, err
}
maxPages := ""
if req.MaxPages != 0 {
maxPages = strconv.Itoa(req.MaxPages)
}
body, contentType, err := buildMultipart("build ocr form",
filePart{field: "file", filename: documentFilename(req.Filename, req.MIME), data: req.Document},
[]formField{
{"langs", strings.Join(req.Languages, ","), false},
{"max_pages", maxPages, false},
})
if err != nil {
return nil, err
}
raw, _, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxResponseBytes)
if err != nil {
return nil, err
}
var out ocrResponse
if err := json.Unmarshal(raw, &out); err != nil {
return nil, fmt.Errorf("llama-swap: decode ocr response: %w", err)
}
if len(out.Pages) == 0 {
// A blank page still comes back as a page with empty text; zero pages
// is API drift or a soft error, never "the document was empty".
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
Message: "ocr response contained no pages: " + truncateForError(raw)}
}
res := &ocr.Result{Raw: json.RawMessage(raw)}
texts := make([]string, 0, len(out.Pages))
for i, pg := range out.Pages {
text := pg.Text
if text == "" && len(pg.Lines) > 0 {
lines := make([]string, 0, len(pg.Lines))
for _, ln := range pg.Lines {
lines = append(lines, ln.Text)
}
text = strings.Join(lines, "\n")
}
number := pg.Number
if number == 0 {
number = i + 1
}
res.Pages = append(res.Pages, ocr.Page{Number: number, Text: text})
texts = append(texts, text)
}
res.Text = strings.Join(texts, "\n\n")
return res, nil
}
// documentFilename picks the multipart filename hint for an OCR document: the
// caller's (sanitized), else one derived from the MIME subtype
// ("document.pdf"), else "document". MIME parameters are stripped before
// matching, mirroring transcriptionFilename.
func documentFilename(filename, mimeType string) string {
if name := sanitizeFilename(filename); name != "" {
return name
}
mt := strings.ToLower(strings.TrimSpace(mimeType))
if parsed, _, err := mime.ParseMediaType(mt); err == nil {
mt = parsed
}
switch mt {
case "application/pdf":
return "document.pdf"
case "image/png":
return "document.png"
case "image/jpeg", "image/jpg":
return "document.jpg"
case "image/webp":
return "document.webp"
default:
return "document"
}
}
+200
View File
@@ -0,0 +1,200 @@
package llamaswap
import (
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
"gitea.stevedudenhoeffer.com/steve/majordomo/ocr"
)
func TestRecognize(t *testing.T) {
var gotPath, gotLangs, gotMaxPages, gotFilename string
var gotFile []byte
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart: %v", err)
}
gotLangs = r.FormValue("langs")
gotMaxPages = r.FormValue("max_pages")
f, hdr, err := r.FormFile("file")
if err != nil {
t.Fatalf("form file: %v", err)
}
defer f.Close()
gotFile, _ = io.ReadAll(f)
gotFilename = hdr.Filename
_, _ = w.Write([]byte(`{"pages":[
{"number":1,"text":"page one","lines":[{"text":"page","bbox":[0,0,1,1],"confidence":0.9},{"text":"one"}],"layout":{}},
{"number":2,"text":"page two"}
]}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
om, err := p.OCRModel("ocr-surya")
if err != nil {
t.Fatalf("OCRModel: %v", err)
}
res, err := om.Recognize(context.Background(),
ocr.Request{Document: []byte("%PDF"), MIME: "application/pdf"},
ocr.WithLanguages("en", "de"), ocr.WithMaxPages(5))
if err != nil {
t.Fatalf("Recognize: %v", err)
}
if gotPath != "/upstream/ocr-surya/v1/ocr" {
t.Errorf("path = %q", gotPath)
}
if gotLangs != "en,de" || gotMaxPages != "5" {
t.Errorf("langs/max_pages = %q/%q", gotLangs, gotMaxPages)
}
if string(gotFile) != "%PDF" || gotFilename != "document.pdf" {
t.Errorf("file = %q name = %q", gotFile, gotFilename)
}
if len(res.Pages) != 2 || res.Pages[0].Number != 1 || res.Pages[0].Text != "page one" ||
res.Pages[1].Number != 2 || res.Pages[1].Text != "page two" {
t.Errorf("pages = %+v", res.Pages)
}
if res.Text != "page one\n\npage two" {
t.Errorf("text = %q", res.Text)
}
if res.Raw == nil {
t.Error("Raw = nil, want raw payload")
}
}
func TestRecognizeJoinsLinesWhenPageTextAbsent(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"pages":[{"lines":[{"text":"first line"},{"text":"second line"}]}]}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
om, _ := p.OCRModel("ocr-surya")
res, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")})
if err != nil {
t.Fatalf("Recognize: %v", err)
}
if len(res.Pages) != 1 || res.Pages[0].Text != "first line\nsecond line" {
t.Errorf("pages = %+v", res.Pages)
}
// Missing page number defaults to position.
if res.Pages[0].Number != 1 {
t.Errorf("number = %d, want 1", res.Pages[0].Number)
}
if res.Text != "first line\nsecond line" {
t.Errorf("text = %q", res.Text)
}
}
func TestRecognizeOmitsUnsetFields(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart: %v", err)
}
for _, k := range []string{"langs", "max_pages"} {
if v, ok := r.MultipartForm.Value[k]; ok {
t.Errorf("unset request sent %q = %v, want omitted", k, v)
}
}
if _, hdr, err := r.FormFile("file"); err == nil {
if hdr.Filename != "document" {
t.Errorf("filename = %q, want document fallback", hdr.Filename)
}
}
_, _ = w.Write([]byte(`{"pages":[{"number":1,"text":"x"}]}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
om, _ := p.OCRModel("ocr-surya")
if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err != nil {
t.Fatalf("Recognize: %v", err)
}
}
func TestRecognizeRejectsBadArgs(t *testing.T) {
p := New(WithBaseURL("http://unused"))
om, _ := p.OCRModel("ocr-surya")
if _, err := om.Recognize(context.Background(), ocr.Request{}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no document: err = %v, want ErrUnsupported", err)
}
if _, err := om.Recognize(context.Background(),
ocr.Request{Document: []byte{1}, MaxPages: -1}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("negative max pages: err = %v, want ErrUnsupported", err)
}
}
func TestRecognizeRejectsZeroPages(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"pages":[]}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
om, _ := p.OCRModel("ocr-surya")
_, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %v, want APIError for zero pages", err)
}
}
func TestRecognizeRejectsNonJSONResponse(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
_, _ = w.Write([]byte("<html>proxy error page</html>"))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
om, _ := p.OCRModel("ocr-surya")
if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err == nil ||
!strings.Contains(err.Error(), "decode ocr response") {
t.Fatalf("err = %v, want decode error for non-JSON body", err)
}
}
func TestRecognizeSurfacesAPIError(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":{"message":"unsupported file type"}}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
om, _ := p.OCRModel("ocr-surya")
_, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("bad")})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %T %v, want *llm.APIError", err, err)
}
if apiErr.Status != http.StatusBadRequest || apiErr.Message != "unsupported file type" || apiErr.Model != "ocr-surya" {
t.Errorf("apiErr = %+v", apiErr)
}
}
func TestDocumentFilename(t *testing.T) {
cases := []struct {
filename, mime, want string
}{
{"scan.pdf", "", "scan.pdf"},
{"evil\r\nname.pdf", "", "evilname.pdf"},
{"", "application/pdf", "document.pdf"},
{"", "image/png", "document.png"},
{"", "image/jpeg", "document.jpg"},
{"", "image/webp; charset=binary", "document.webp"},
{"", "", "document"},
}
for _, tc := range cases {
if got := documentFilename(tc.filename, tc.mime); got != tc.want {
t.Errorf("documentFilename(%q, %q) = %q, want %q", tc.filename, tc.mime, got, tc.want)
}
}
}
+106
View File
@@ -0,0 +1,106 @@
// restore.go implements the imagegen.ColorizeProvider and
// imagegen.FaceRestoreProvider surfaces against the mediautils shim reached
// through llama-swap's /upstream passthrough (ADR-0023):
//
// colorize POST /upstream/<id>/v1/colorize (DDColor)
// restore_faces POST /upstream/<id>/v1/restore_faces (GFPGAN)
//
// Both are one-file multipart in, one PNG out, mirroring mediautil.go.
package llamaswap
import (
"context"
"fmt"
"net/http"
"strconv"
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)
// --- colorize ---
// ColorizeModel implements imagegen.ColorizeProvider against the mediautils
// shim's POST /v1/colorize. The id selects which upstream llama-swap loads.
func (p *Provider) ColorizeModel(id string, opts ...imagegen.ColorizeModelOption) (imagegen.Colorizer, error) {
if err := p.requireBaseURL(); err != nil {
return nil, err
}
_ = imagegen.ApplyColorizeModelOptions(opts)
return &colorizeModel{p: p, id: id}, nil
}
type colorizeModel struct {
p *Provider
id string
}
// Colorize implements imagegen.Colorizer.
func (m *colorizeModel) Colorize(ctx context.Context, req imagegen.ColorizeRequest, opts ...imagegen.ColorizeOption) (*imagegen.Result, error) {
req = req.Apply(opts...)
if len(req.Image.Data) == 0 {
return nil, fmt.Errorf("%w: colorization requires an image", llm.ErrUnsupported)
}
path, err := upstreamPath(m.id, "/v1/colorize")
if err != nil {
return nil, err
}
body, contentType, err := buildMultipart("build colorize form",
filePart{field: "file", filename: "image.png", data: req.Image.Data},
nil)
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, "colorize", raw, respType)
}
// --- face restoration ---
// FaceRestoreModel implements imagegen.FaceRestoreProvider against the
// mediautils shim's POST /v1/restore_faces.
func (p *Provider) FaceRestoreModel(id string, opts ...imagegen.FaceRestoreModelOption) (imagegen.FaceRestorer, error) {
if err := p.requireBaseURL(); err != nil {
return nil, err
}
_ = imagegen.ApplyFaceRestoreModelOptions(opts)
return &faceRestoreModel{p: p, id: id}, nil
}
type faceRestoreModel struct {
p *Provider
id string
}
// RestoreFaces implements imagegen.FaceRestorer.
func (m *faceRestoreModel) RestoreFaces(ctx context.Context, req imagegen.FaceRestoreRequest, opts ...imagegen.FaceRestoreOption) (*imagegen.Result, error) {
req = req.Apply(opts...)
if len(req.Image.Data) == 0 {
return nil, fmt.Errorf("%w: face restoration requires an image", llm.ErrUnsupported)
}
if req.Upscale != 0 && req.Upscale != 1 && req.Upscale != 2 {
return nil, fmt.Errorf("%w: face-restore upscale must be 1 or 2, got %d", llm.ErrUnsupported, req.Upscale)
}
path, err := upstreamPath(m.id, "/v1/restore_faces")
if err != nil {
return nil, err
}
upscale := ""
if req.Upscale != 0 {
upscale = strconv.Itoa(req.Upscale)
}
body, contentType, err := buildMultipart("build restore-faces form",
filePart{field: "file", filename: "image.png", data: req.Image.Data},
[]formField{{"upscale", upscale, 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, "face restoration", raw, respType)
}
+76
View File
@@ -0,0 +1,76 @@
// segment.go implements imagegen.SegmentationProvider against a
// GroundingDINO+SAM shim (segment-langsam) reached through llama-swap's
// /upstream passthrough (ADR-0023):
//
// POST /upstream/<id>/v1/segment multipart file,prompt[,threshold],output=mask
//
// The response is a single grayscale mask PNG where WHITE marks the prompted
// region — directly usable as imagegen.EditRequest.Mask (white = repaint).
// The shim also offers output=cutout|boxes; this client always requests the
// mask, because a cutout is derivable client-side from mask+original with no
// second GPU call.
package llamaswap
import (
"context"
"fmt"
"net/http"
"strconv"
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)
// SegmentationModel implements imagegen.SegmentationProvider. The id selects
// which upstream llama-swap loads.
func (p *Provider) SegmentationModel(id string, opts ...imagegen.SegmentationModelOption) (imagegen.Segmenter, error) {
if err := p.requireBaseURL(); err != nil {
return nil, err
}
_ = imagegen.ApplySegmentationModelOptions(opts)
return &segmentationModel{p: p, id: id}, nil
}
type segmentationModel struct {
p *Provider
id string
}
// Segment implements imagegen.Segmenter.
func (m *segmentationModel) Segment(ctx context.Context, req imagegen.SegmentationRequest, opts ...imagegen.SegmentationOption) (*imagegen.Result, error) {
req = req.Apply(opts...)
if len(req.Image.Data) == 0 {
return nil, fmt.Errorf("%w: segmentation requires an image", llm.ErrUnsupported)
}
if req.Prompt == "" {
return nil, fmt.Errorf("%w: segmentation requires a prompt", llm.ErrUnsupported)
}
if req.Threshold < 0 || req.Threshold > 1 {
return nil, fmt.Errorf("%w: segmentation threshold must be in [0,1], got %g", llm.ErrUnsupported, req.Threshold)
}
path, err := upstreamPath(m.id, "/v1/segment")
if err != nil {
return nil, err
}
threshold := ""
if req.Threshold != 0 {
threshold = strconv.FormatFloat(req.Threshold, 'g', -1, 64)
}
body, contentType, err := buildMultipart("build segment form",
filePart{field: "file", filename: "image.png", data: req.Image.Data},
[]formField{
{"prompt", req.Prompt, true},
{"threshold", threshold, false},
// Always the mask: white = prompted region, EditRequest.Mask
// polarity. Cutouts are derived client-side.
{"output", "mask", true},
})
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, "segmentation", raw, respType)
}
+270
View File
@@ -0,0 +1,270 @@
package llamaswap
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)
func TestSegment(t *testing.T) {
png := pngFixture(t)
var gotPath, gotPrompt, gotThreshold, gotOutput, gotFilename string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
gotPrompt = r.FormValue("prompt")
gotThreshold = r.FormValue("threshold")
gotOutput = r.FormValue("output")
if f, hdr, err := r.FormFile("file"); err == nil {
gotFilename = hdr.Filename
f.Close()
} else {
t.Errorf("file part: %v", err)
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
sg, err := p.SegmentationModel("segment-langsam")
if err != nil {
t.Fatalf("SegmentationModel: %v", err)
}
res, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{MIME: "image/png", Data: png}, Prompt: "the red car"},
imagegen.WithSegmentationThreshold(0.35))
if err != nil {
t.Fatalf("Segment: %v", err)
}
if gotPath != "/upstream/segment-langsam/v1/segment" {
t.Errorf("path = %q", gotPath)
}
if gotPrompt != "the red car" || gotThreshold != "0.35" || gotOutput != "mask" || gotFilename != "image.png" {
t.Errorf("prompt/threshold/output/filename = %q/%q/%q/%q", gotPrompt, gotThreshold, gotOutput, gotFilename)
}
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
t.Fatalf("images = %+v", res.Images)
}
}
func TestSegmentOmitsDefaultThreshold(t *testing.T) {
png := pngFixture(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
if _, ok := r.MultipartForm.Value["threshold"]; ok {
t.Error("threshold field sent for default; want omitted")
}
if got := r.FormValue("output"); got != "mask" {
t.Errorf("output = %q, want mask (always sent)", got)
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
sg, _ := p.SegmentationModel("segment-langsam")
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: png}, Prompt: "dog"}); err != nil {
t.Fatalf("Segment: %v", err)
}
}
func TestSegmentRejectsBadArgs(t *testing.T) {
p := New(WithBaseURL("http://unused"))
sg, _ := p.SegmentationModel("segment-langsam")
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Prompt: "dog"}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no image: err = %v, want ErrUnsupported", err)
}
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no prompt: err = %v, want ErrUnsupported", err)
}
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: 1.5}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("threshold 1.5: err = %v, want ErrUnsupported", err)
}
}
func TestSegmentRejectsNonImageResponse(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
_, _ = w.Write([]byte("<html>proxy error page</html>"))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
sg, _ := p.SegmentationModel("segment-langsam")
_, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: pngFixture(t)}, Prompt: "dog"})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %v, want APIError for non-image body", err)
}
}
func TestColorize(t *testing.T) {
png := pngFixture(t)
var gotPath, gotFilename string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
if f, hdr, err := r.FormFile("file"); err == nil {
gotFilename = hdr.Filename
f.Close()
} else {
t.Errorf("file part: %v", err)
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
cl, err := p.ColorizeModel("mediautils")
if err != nil {
t.Fatalf("ColorizeModel: %v", err)
}
res, err := cl.Colorize(context.Background(),
imagegen.ColorizeRequest{Image: imagegen.Image{Data: png}})
if err != nil {
t.Fatalf("Colorize: %v", err)
}
if gotPath != "/upstream/mediautils/v1/colorize" {
t.Errorf("path = %q", gotPath)
}
if gotFilename != "image.png" {
t.Errorf("filename = %q", gotFilename)
}
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
t.Fatalf("images = %+v", res.Images)
}
}
func TestColorizeRejectsEmptyImage(t *testing.T) {
p := New(WithBaseURL("http://unused"))
cl, _ := p.ColorizeModel("mediautils")
if _, err := cl.Colorize(context.Background(), imagegen.ColorizeRequest{}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("err = %v, want ErrUnsupported", err)
}
}
func TestColorizeSurfacesAPIError(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusServiceUnavailable)
_, _ = w.Write([]byte(`{"error":{"message":"model is loading"}}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
cl, _ := p.ColorizeModel("mediautils")
_, err := cl.Colorize(context.Background(),
imagegen.ColorizeRequest{Image: imagegen.Image{Data: pngFixture(t)}})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %T %v, want *llm.APIError", err, err)
}
if apiErr.Status != http.StatusServiceUnavailable || apiErr.Message != "model is loading" || apiErr.Model != "mediautils" {
t.Errorf("apiErr = %+v", apiErr)
}
}
func TestRestoreFaces(t *testing.T) {
png := pngFixture(t)
var gotPath, gotUpscale string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
gotUpscale = r.FormValue("upscale")
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
fr, err := p.FaceRestoreModel("mediautils")
if err != nil {
t.Fatalf("FaceRestoreModel: %v", err)
}
res, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}},
imagegen.WithFaceRestoreUpscale(2))
if err != nil {
t.Fatalf("RestoreFaces: %v", err)
}
if gotPath != "/upstream/mediautils/v1/restore_faces" {
t.Errorf("path = %q", gotPath)
}
if gotUpscale != "2" {
t.Errorf("upscale = %q", gotUpscale)
}
if len(res.Images) != 1 {
t.Fatalf("images = %+v", res.Images)
}
}
func TestRestoreFacesOmitsDefaultUpscale(t *testing.T) {
png := pngFixture(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
if _, ok := r.MultipartForm.Value["upscale"]; ok {
t.Error("upscale field sent for default; want omitted")
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
fr, _ := p.FaceRestoreModel("mediautils")
if _, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}}); err != nil {
t.Fatalf("RestoreFaces: %v", err)
}
}
func TestRestoreFacesRejectsBadArgs(t *testing.T) {
p := New(WithBaseURL("http://unused"))
fr, _ := p.FaceRestoreModel("mediautils")
if _, err := fr.RestoreFaces(context.Background(), imagegen.FaceRestoreRequest{}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no image: err = %v, want ErrUnsupported", err)
}
if _, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: []byte{1}}, Upscale: 3}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("upscale 3: err = %v, want ErrUnsupported", err)
}
}
func TestRestoreFacesRejectsNonImageResponse(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header()["Content-Type"] = nil // NO Content-Type at all
_, _ = w.Write([]byte("502 bad gateway"))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
fr, _ := p.FaceRestoreModel("mediautils")
_, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: pngFixture(t)}})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %v, want APIError for headerless non-image body", err)
}
}