feat: wave-3 image + document surfaces — segmentation, colorize, face restore, ocr (ADR-0023)
- 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:
@@ -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.
|
||||
@@ -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 |
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user