feat: wave-3 image + document surfaces — segmentation, colorize, face restore, OCR (ADR-0023) #17

Merged
steve merged 2 commits from feat/wave3-image-doc-surfaces into main 2026-07-16 23:24:33 +00:00
11 changed files with 1155 additions and 0 deletions
Showing only changes of commit e6987f54b2 - Show all commits
+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
Review

🟡 Threshold range comment (0,1] contradicts validation [0,1] and 0 sentinel usage

correctness · flagged by 1 model

Finding 1: imagegen/segment.go:16-18 — Threshold range comment

🪰 Gadfly · advisory

🟡 **Threshold range comment (0,1] contradicts validation [0,1] and 0 sentinel usage** _correctness · flagged by 1 model_ **Finding 1: `imagegen/segment.go:16-18` — Threshold range comment** <sub>🪰 Gadfly · advisory</sub>
// 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")
}
Review

🟡 documentFilename duplicates transcriptionFilename scaffolding; extract a shared MIME->filename helper parameterized by the extension table

maintainability · flagged by 1 model

  • provider/llamaswap/ocr.go:104documentFilename is a near-verbatim copy of transcriptionFilename (audio.go:151). Both do the same four steps: sanitizeFilenamestrings.ToLower/TrimSpacemime.ParseMediaTypeswitch on the parsed type → typed fallback. Only the MIME→extension table and the default string differ. The comment on documentFilename even calls out the mirroring ("mirroring transcriptionFilename"), which is exactly the signal that the shared scaffolding wan…

🪰 Gadfly · advisory

🟡 **documentFilename duplicates transcriptionFilename scaffolding; extract a shared MIME->filename helper parameterized by the extension table** _maintainability · flagged by 1 model_ - **`provider/llamaswap/ocr.go:104` — `documentFilename` is a near-verbatim copy of `transcriptionFilename` (`audio.go:151`).** Both do the same four steps: `sanitizeFilename` → `strings.ToLower`/`TrimSpace` → `mime.ParseMediaType` → `switch` on the parsed type → typed fallback. Only the MIME→extension table and the default string differ. The comment on `documentFilename` even calls out the mirroring ("mirroring transcriptionFilename"), which is exactly the signal that the shared scaffolding wan… <sub>🪰 Gadfly · advisory</sub>
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 {
Review

documentFilename structurally duplicates transcriptionFilename pattern

maintainability · flagged by 1 model

  • provider/llamaswap/ocr.go:120 and provider/llamaswap/audio.go:151documentFilename structurally duplicates transcriptionFilename: Both functions follow the identical pattern: sanitize → mime.ParseMediaType → switch on MIME → return default. They're small (~20 lines each) and domain-specific, so extraction may not be worth the indirection, but any future change to the sanitize/parse/fallback logic now has two copy-paste sites to update. *(Verified by reading both functions side-…

🪰 Gadfly · advisory

⚪ **documentFilename structurally duplicates transcriptionFilename pattern** _maintainability · flagged by 1 model_ - **`provider/llamaswap/ocr.go:120` and `provider/llamaswap/audio.go:151` — `documentFilename` structurally duplicates `transcriptionFilename`**: Both functions follow the identical pattern: sanitize → `mime.ParseMediaType` → switch on MIME → return default. They're small (~20 lines each) and domain-specific, so extraction may not be worth the indirection, but any future change to the sanitize/parse/fallback logic now has two copy-paste sites to update. *(Verified by reading both functions side-… <sub>🪰 Gadfly · advisory</sub>
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
Review

Filename restore.go doesn't reflect it also implements ColorizeProvider

maintainability · flagged by 1 model

  • provider/llamaswap/restore.go:1 — filename doesn't reflect dual responsibility: The file is named restore.go but implements both ColorizeProvider and FaceRestoreProvider. The file comment is clear, but the name only hints at face restoration. The existing pattern (mediautil.go groups upscale + background removal + interpolation) shows grouping is fine, but the name should match. Consider colorize_restore.go or image_enhance.go. *(Verified by reading the file header and comp…

🪰 Gadfly · advisory

⚪ **Filename restore.go doesn't reflect it also implements ColorizeProvider** _maintainability · flagged by 1 model_ - **`provider/llamaswap/restore.go:1` — filename doesn't reflect dual responsibility**: The file is named `restore.go` but implements *both* `ColorizeProvider` and `FaceRestoreProvider`. The file comment is clear, but the name only hints at face restoration. The existing pattern (`mediautil.go` groups upscale + background removal + interpolation) shows grouping is fine, but the name should match. Consider `colorize_restore.go` or `image_enhance.go`. *(Verified by reading the file header and comp… <sub>🪰 Gadfly · advisory</sub>
// 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 {
Review

🟠 NaN passes segmentation threshold validation and is sent to backend

correctness, error-handling, security · flagged by 4 models

  • provider/llamaswap/segment.go:48 — Segmentation threshold validation lets NaN slip through. In Go, NaN < 0 and NaN > 1 are both false, so a Threshold of NaN passes the [0,1] guard and is serialized as "NaN" to the backend. It should be rejected alongside negatives and values greater than 1. Add a math.IsNaN(req.Threshold) check (or a < 0 || > 1 || != req.Threshold NaN test) before the path is built.

🪰 Gadfly · advisory

🟠 **NaN passes segmentation threshold validation and is sent to backend** _correctness, error-handling, security · flagged by 4 models_ - `provider/llamaswap/segment.go:48` — Segmentation threshold validation lets `NaN` slip through. In Go, `NaN < 0` and `NaN > 1` are both `false`, so a `Threshold` of `NaN` passes the `[0,1]` guard and is serialized as `"NaN"` to the backend. It should be rejected alongside negatives and values greater than 1. Add a `math.IsNaN(req.Threshold)` check (or a `< 0 || > 1 || != req.Threshold` NaN test) before the path is built. <sub>🪰 Gadfly · advisory</sub>
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
Review

Test filename omits 'colorize' despite containing colorize tests

maintainability · flagged by 1 model

  • provider/llamaswap/segment_restore_test.go:1 — test filename omits "colorize": The file contains TestColorize, TestColorizeRejectsEmptyImage, and TestColorizeSurfacesAPIError but the filename only says segment_restore. Someone searching for colorize tests by filename would miss them. Rename to segment_colorize_restore_test.go or split colorize tests into their own file. (Verified by reading the test file contents.)

🪰 Gadfly · advisory

⚪ **Test filename omits 'colorize' despite containing colorize tests** _maintainability · flagged by 1 model_ - **`provider/llamaswap/segment_restore_test.go:1` — test filename omits "colorize"**: The file contains `TestColorize`, `TestColorizeRejectsEmptyImage`, and `TestColorizeSurfacesAPIError` but the filename only says `segment_restore`. Someone searching for colorize tests by filename would miss them. Rename to `segment_colorize_restore_test.go` or split colorize tests into their own file. *(Verified by reading the test file contents.)* <sub>🪰 Gadfly · advisory</sub>
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)
}
}