644 lines
17 KiB
Go
644 lines
17 KiB
Go
package openrouter
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
stderrors "errors"
|
|
"fmt"
|
|
"image"
|
|
"image/color"
|
|
_ "image/gif"
|
|
_ "image/jpeg"
|
|
"image/png"
|
|
"math"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"golang.org/x/image/draw"
|
|
_ "golang.org/x/image/webp"
|
|
)
|
|
|
|
const Gemini3ProImagePreviewModel = "google/gemini-3-pro-image-preview"
|
|
|
|
const maskImageInstruction = "下一张图片是遮罩图。非黑色像素表示允许编辑的区域；黑色像素表示必须保持不变的保护区域。"
|
|
|
|
const referenceImageInstruction = "下一张图片是地板材质参考图。生成新地面时必须严格匹配它的颜色、纹理、图案、光泽和整体材质外观；生成的地面应看起来就是这个材质。"
|
|
|
|
var supportedAspectRatios = []struct {
|
|
value string
|
|
width int
|
|
height int
|
|
}{
|
|
{value: "1:1", width: 1, height: 1},
|
|
{value: "2:3", width: 2, height: 3},
|
|
{value: "3:2", width: 3, height: 2},
|
|
{value: "3:4", width: 3, height: 4},
|
|
{value: "4:3", width: 4, height: 3},
|
|
{value: "4:5", width: 4, height: 5},
|
|
{value: "5:4", width: 5, height: 4},
|
|
{value: "9:16", width: 9, height: 16},
|
|
{value: "16:9", width: 16, height: 9},
|
|
{value: "21:9", width: 21, height: 9},
|
|
}
|
|
|
|
type ImageGenerationRequest struct {
|
|
Prompt string
|
|
OutputPath string
|
|
Model string
|
|
Modalities []string
|
|
Provider *Provider
|
|
InputImagePath string
|
|
InputMimeType string
|
|
MaskImagePath string
|
|
MaskMimeType string
|
|
ReferenceImagePath string // optional: reference image for floor material style
|
|
ReferenceMimeType string
|
|
AspectRatio string
|
|
ImageSize string
|
|
TopP *float64
|
|
TargetWidth int
|
|
TargetHeight int
|
|
MatchSizeImagePath string
|
|
}
|
|
|
|
type ImageGenerationResult struct {
|
|
Path string
|
|
MimeType string
|
|
Model string
|
|
TextContent string
|
|
Width int
|
|
Height int
|
|
}
|
|
|
|
type ImageRecognitionRequest struct {
|
|
Prompt string
|
|
OutputPath string
|
|
Model string
|
|
Modalities []string
|
|
Provider *Provider
|
|
InputImagePath string
|
|
InputMimeType string
|
|
AspectRatio string
|
|
ImageSize string
|
|
TopP *float64
|
|
TargetWidth int
|
|
TargetHeight int
|
|
MatchSizeImagePath string
|
|
ResizeOutput bool
|
|
}
|
|
|
|
type ImageRecognitionResult = ImageGenerationResult
|
|
|
|
// RecognizeImage calls an OpenRouter image-capable chat model and saves the
|
|
// first returned image as a PNG file while returning the text content for JSON parsing.
|
|
func (c *Client) RecognizeImage(ctx context.Context, req ImageRecognitionRequest) (*ImageRecognitionResult, error) {
|
|
if strings.TrimSpace(req.Prompt) == "" {
|
|
return nil, stderrors.New("openrouter: recognition prompt is required")
|
|
}
|
|
if strings.TrimSpace(req.InputImagePath) == "" {
|
|
return nil, stderrors.New("openrouter: recognition input image path is required")
|
|
}
|
|
if strings.TrimSpace(req.OutputPath) == "" {
|
|
return nil, stderrors.New("openrouter: recognition output path is required")
|
|
}
|
|
|
|
model := req.Model
|
|
if model == "" {
|
|
model = Gemini3ProImagePreviewModel
|
|
}
|
|
targetWidth, targetHeight := 0, 0
|
|
if req.ResizeOutput || req.AspectRatio == "" {
|
|
if req.MatchSizeImagePath == "" {
|
|
req.MatchSizeImagePath = req.InputImagePath
|
|
}
|
|
var err error
|
|
targetWidth, targetHeight, err = resolveTargetSize(ImageGenerationRequest{
|
|
TargetWidth: req.TargetWidth,
|
|
TargetHeight: req.TargetHeight,
|
|
MatchSizeImagePath: req.MatchSizeImagePath,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
content, err := buildImageRecognitionContent(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
chatReq := ChatRequest{
|
|
Model: model,
|
|
Modalities: imageOutputModalities(model, req.Modalities),
|
|
Provider: req.Provider,
|
|
TopP: req.TopP,
|
|
Messages: []Message{
|
|
{
|
|
Role: "user",
|
|
Content: content,
|
|
},
|
|
},
|
|
}
|
|
if req.AspectRatio == "" && targetWidth > 0 && targetHeight > 0 {
|
|
req.AspectRatio = aspectRatioString(targetWidth, targetHeight)
|
|
}
|
|
if req.AspectRatio != "" || req.ImageSize != "" {
|
|
chatReq.ImageConfig = &ImageConfig{
|
|
AspectRatio: req.AspectRatio,
|
|
ImageSize: normalizeImageSize(req.ImageSize),
|
|
}
|
|
}
|
|
|
|
resp, err := c.chatCompletion(ctx, chatReq)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
outputWidth, outputHeight := 0, 0
|
|
if req.ResizeOutput {
|
|
outputWidth, outputHeight = targetWidth, targetHeight
|
|
}
|
|
return saveFirstResponseImage(resp.Choices, req.OutputPath, model, outputWidth, outputHeight)
|
|
}
|
|
|
|
// GenerateImageToFile calls an OpenRouter image-capable chat model and saves the
|
|
// first returned image as a PNG file.
|
|
func (c *Client) GenerateImageToFile(ctx context.Context, req ImageGenerationRequest) (*ImageGenerationResult, error) {
|
|
if strings.TrimSpace(req.Prompt) == "" {
|
|
return nil, stderrors.New("openrouter: image prompt is required")
|
|
}
|
|
if strings.TrimSpace(req.OutputPath) == "" {
|
|
return nil, stderrors.New("openrouter: image output path is required")
|
|
}
|
|
|
|
model := req.Model
|
|
if model == "" {
|
|
model = Gemini3ProImagePreviewModel
|
|
}
|
|
if req.MatchSizeImagePath == "" {
|
|
req.MatchSizeImagePath = req.InputImagePath
|
|
}
|
|
|
|
targetWidth, targetHeight, err := resolveTargetSize(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
content, err := buildImageGenerationContent(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
chatReq := ChatRequest{
|
|
Model: model,
|
|
Modalities: imageOutputModalities(model, req.Modalities),
|
|
Provider: req.Provider,
|
|
TopP: req.TopP,
|
|
Messages: []Message{
|
|
{
|
|
Role: "user",
|
|
Content: content,
|
|
},
|
|
},
|
|
}
|
|
if req.AspectRatio == "" && targetWidth > 0 && targetHeight > 0 {
|
|
req.AspectRatio = aspectRatioString(targetWidth, targetHeight)
|
|
}
|
|
if req.AspectRatio != "" || req.ImageSize != "" {
|
|
chatReq.ImageConfig = &ImageConfig{
|
|
AspectRatio: req.AspectRatio,
|
|
ImageSize: normalizeImageSize(req.ImageSize),
|
|
}
|
|
}
|
|
|
|
resp, err := c.chatCompletion(ctx, chatReq)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return saveFirstResponseImage(resp.Choices, req.OutputPath, model, targetWidth, targetHeight)
|
|
}
|
|
|
|
func saveFirstResponseImage(choices []Choice, outputPath string, model string, targetWidth, targetHeight int) (*ImageGenerationResult, error) {
|
|
outputPath = normalizePNGOutputPath(outputPath)
|
|
if len(choices) == 0 {
|
|
return nil, stderrors.New("openrouter: image response has no choices")
|
|
}
|
|
|
|
message := choices[0].Message
|
|
if len(message.Images) == 0 {
|
|
return nil, fmt.Errorf("openrouter: image response has no images: %s", message.Content)
|
|
}
|
|
|
|
_, data, err := decodeDataURL(message.Images[0].ImageURL.URL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
data, width, height, err := encodePNGImageData(data, targetWidth, targetHeight)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
|
|
return nil, fmt.Errorf("openrouter: create output directory: %w", err)
|
|
}
|
|
if err := os.WriteFile(outputPath, data, 0644); err != nil {
|
|
return nil, fmt.Errorf("openrouter: write image file: %w", err)
|
|
}
|
|
|
|
return &ImageGenerationResult{
|
|
Path: outputPath,
|
|
MimeType: "image/png",
|
|
Model: model,
|
|
TextContent: message.Content,
|
|
Width: width,
|
|
Height: height,
|
|
}, nil
|
|
}
|
|
|
|
func resolveTargetSize(req ImageGenerationRequest) (int, int, error) {
|
|
if req.TargetWidth < 0 || req.TargetHeight < 0 {
|
|
return 0, 0, stderrors.New("openrouter: target image size must be positive")
|
|
}
|
|
if req.TargetWidth > 0 || req.TargetHeight > 0 {
|
|
if req.TargetWidth == 0 || req.TargetHeight == 0 {
|
|
return 0, 0, stderrors.New("openrouter: target width and height must be set together")
|
|
}
|
|
return req.TargetWidth, req.TargetHeight, nil
|
|
}
|
|
if req.MatchSizeImagePath == "" {
|
|
return 0, 0, nil
|
|
}
|
|
|
|
file, err := os.Open(req.MatchSizeImagePath)
|
|
if err != nil {
|
|
return 0, 0, fmt.Errorf("openrouter: open match size image: %w", err)
|
|
}
|
|
defer file.Close()
|
|
|
|
config, _, err := image.DecodeConfig(file)
|
|
if err != nil {
|
|
return 0, 0, fmt.Errorf("openrouter: decode match size image config: %w", err)
|
|
}
|
|
return config.Width, config.Height, nil
|
|
}
|
|
|
|
func buildImageGenerationContent(req ImageGenerationRequest) ([]ContentPart, error) {
|
|
if req.MaskImagePath != "" && req.InputImagePath == "" {
|
|
return nil, stderrors.New("openrouter: mask image requires input image")
|
|
}
|
|
|
|
content := make([]ContentPart, 0, 6)
|
|
if req.InputImagePath != "" {
|
|
dataURI, err := imageFileDataURI(req.InputImagePath, req.InputMimeType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
content = append(content, ContentPart{
|
|
Type: "image_url",
|
|
ImageURL: &ImageURL{URL: dataURI},
|
|
})
|
|
}
|
|
if req.MaskImagePath != "" {
|
|
dataURI, err := imageFileDataURI(req.MaskImagePath, req.MaskMimeType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
content = append(content, ContentPart{Type: "text", Text: maskImageInstruction})
|
|
content = append(content, ContentPart{
|
|
Type: "image_url",
|
|
ImageURL: &ImageURL{URL: dataURI},
|
|
})
|
|
}
|
|
if req.ReferenceImagePath != "" {
|
|
dataURI, err := imageFileDataURI(req.ReferenceImagePath, req.ReferenceMimeType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
content = append(content, ContentPart{
|
|
Type: "image_url",
|
|
ImageURL: &ImageURL{URL: dataURI},
|
|
})
|
|
content = append(content, ContentPart{Type: "text", Text: referenceImageInstruction})
|
|
}
|
|
content = append(content, ContentPart{Type: "text", Text: req.Prompt})
|
|
return content, nil
|
|
}
|
|
|
|
func buildImageRecognitionContent(req ImageRecognitionRequest) ([]ContentPart, error) {
|
|
dataURI, err := imageFileDataURI(req.InputImagePath, req.InputMimeType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return []ContentPart{
|
|
{
|
|
Type: "image_url",
|
|
ImageURL: &ImageURL{URL: dataURI},
|
|
},
|
|
{Type: "text", Text: req.Prompt},
|
|
}, nil
|
|
}
|
|
|
|
func imageFileDataURI(path string, mimeType string) (string, error) {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return "", fmt.Errorf("openrouter: read input image: %w", err)
|
|
}
|
|
if mimeType == "" {
|
|
mimeType = http.DetectContentType(data)
|
|
}
|
|
return fmt.Sprintf("data:%s;base64,%s", mimeType, base64.StdEncoding.EncodeToString(data)), nil
|
|
}
|
|
|
|
func normalizePNGOutputPath(outputPath string) string {
|
|
ext := filepath.Ext(outputPath)
|
|
if strings.EqualFold(ext, ".png") {
|
|
return strings.TrimSuffix(outputPath, ext) + ".png"
|
|
}
|
|
if ext == "" {
|
|
return outputPath + ".png"
|
|
}
|
|
return strings.TrimSuffix(outputPath, ext) + ".png"
|
|
}
|
|
|
|
func normalizeImageSize(imageSize string) string {
|
|
imageSize = strings.TrimSpace(imageSize)
|
|
if strings.HasSuffix(strings.ToLower(imageSize), "k") {
|
|
return strings.TrimSuffix(imageSize, imageSize[len(imageSize)-1:]) + "K"
|
|
}
|
|
return imageSize
|
|
}
|
|
|
|
func imageOutputModalities(model string, override []string) []string {
|
|
if len(override) > 0 {
|
|
return normalizeModalities(override)
|
|
}
|
|
if isImageOnlyModel(model) {
|
|
return []string{"image"}
|
|
}
|
|
return []string{"image", "text"}
|
|
}
|
|
|
|
func normalizeModalities(modalities []string) []string {
|
|
normalized := make([]string, 0, len(modalities))
|
|
seen := make(map[string]struct{}, len(modalities))
|
|
for _, modality := range modalities {
|
|
modality = strings.ToLower(strings.TrimSpace(modality))
|
|
if modality == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[modality]; ok {
|
|
continue
|
|
}
|
|
seen[modality] = struct{}{}
|
|
normalized = append(normalized, modality)
|
|
}
|
|
return normalized
|
|
}
|
|
|
|
func isImageOnlyModel(model string) bool {
|
|
model = strings.ToLower(strings.TrimSpace(model))
|
|
imageOnlyPrefixes := []string{
|
|
"black-forest-labs/",
|
|
"sourceful/",
|
|
"bytedance-seed/",
|
|
}
|
|
for _, prefix := range imageOnlyPrefixes {
|
|
if strings.HasPrefix(model, prefix) {
|
|
return true
|
|
}
|
|
}
|
|
|
|
imageOnlyModels := []string{
|
|
"x-ai/grok-imagine-image-quality",
|
|
}
|
|
for _, candidate := range imageOnlyModels {
|
|
if model == candidate {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func encodePNGImageData(data []byte, targetWidth, targetHeight int) ([]byte, int, int, error) {
|
|
src, _, err := image.Decode(bytes.NewReader(data))
|
|
if err != nil {
|
|
return nil, 0, 0, fmt.Errorf("openrouter: decode generated image: %w", err)
|
|
}
|
|
|
|
dst := src
|
|
width := src.Bounds().Dx()
|
|
height := src.Bounds().Dy()
|
|
if targetWidth > 0 && targetHeight > 0 {
|
|
resized := image.NewRGBA(image.Rect(0, 0, targetWidth, targetHeight))
|
|
draw.ApproxBiLinear.Scale(resized, resized.Bounds(), src, src.Bounds(), draw.Over, nil)
|
|
dst = resized
|
|
width = targetWidth
|
|
height = targetHeight
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
if err := png.Encode(&buf, dst); err != nil {
|
|
return nil, 0, 0, fmt.Errorf("openrouter: encode png image: %w", err)
|
|
}
|
|
return buf.Bytes(), width, height, nil
|
|
}
|
|
|
|
func aspectRatioString(width, height int) string {
|
|
if width <= 0 || height <= 0 {
|
|
return ""
|
|
}
|
|
|
|
sourceRatio := math.Log(float64(width) / float64(height))
|
|
best := supportedAspectRatios[0]
|
|
bestDistance := math.Abs(sourceRatio - math.Log(float64(best.width)/float64(best.height)))
|
|
for _, candidate := range supportedAspectRatios[1:] {
|
|
distance := math.Abs(sourceRatio - math.Log(float64(candidate.width)/float64(candidate.height)))
|
|
if distance < bestDistance {
|
|
best = candidate
|
|
bestDistance = distance
|
|
}
|
|
}
|
|
return best.value
|
|
}
|
|
|
|
func decodeDataURL(dataURL string) (string, []byte, error) {
|
|
header, payload, ok := strings.Cut(dataURL, ",")
|
|
if !ok {
|
|
return "", nil, stderrors.New("openrouter: invalid image data url")
|
|
}
|
|
if !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") {
|
|
return "", nil, stderrors.New("openrouter: unsupported image data url")
|
|
}
|
|
|
|
mimeType := strings.TrimPrefix(strings.Split(header, ";")[0], "data:")
|
|
data, err := base64.StdEncoding.DecodeString(payload)
|
|
if err != nil {
|
|
return "", nil, fmt.Errorf("openrouter: decode image data: %w", err)
|
|
}
|
|
|
|
return mimeType, data, nil
|
|
}
|
|
|
|
// MaxImageDimension is the maximum width or height for images sent to the AI API.
|
|
const MaxImageDimension = 2048
|
|
|
|
// ResizeImageToFit reads an image, scales it so the longest side ≤ maxDim pixels,
|
|
// and writes the result as PNG.
|
|
func ResizeImageToFit(inputPath, outputPath string) error {
|
|
data, err := os.ReadFile(inputPath)
|
|
if err != nil {
|
|
return fmt.Errorf("resize: read input: %w", err)
|
|
}
|
|
src, _, err := image.Decode(bytes.NewReader(data))
|
|
if err != nil {
|
|
return fmt.Errorf("resize: decode: %w", err)
|
|
}
|
|
b := src.Bounds()
|
|
w, h := b.Dx(), b.Dy()
|
|
longest := w
|
|
if h > longest {
|
|
longest = h
|
|
}
|
|
dst := src
|
|
if longest > MaxImageDimension {
|
|
scale := float64(MaxImageDimension) / float64(longest)
|
|
newW := int(float64(w)*scale + 0.5)
|
|
newH := int(float64(h)*scale + 0.5)
|
|
resized := image.NewRGBA(image.Rect(0, 0, newW, newH))
|
|
draw.ApproxBiLinear.Scale(resized, resized.Bounds(), src, src.Bounds(), draw.Over, nil)
|
|
dst = resized
|
|
}
|
|
var buf bytes.Buffer
|
|
if err := png.Encode(&buf, dst); err != nil {
|
|
return fmt.Errorf("resize: encode png: %w", err)
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
|
|
return err
|
|
}
|
|
return os.WriteFile(outputPath, buf.Bytes(), 0644)
|
|
}
|
|
|
|
// RepairNearBlackBorder removes accidental model-added black canvas edges.
|
|
// It only acts when an entire edge row/column is almost uniformly black.
|
|
func RepairNearBlackBorder(path string) error {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return fmt.Errorf("repair border: read image: %w", err)
|
|
}
|
|
src, _, err := image.Decode(bytes.NewReader(data))
|
|
if err != nil {
|
|
return fmt.Errorf("repair border: decode image: %w", err)
|
|
}
|
|
b := src.Bounds()
|
|
w, h := b.Dx(), b.Dy()
|
|
if w < 8 || h < 8 {
|
|
return nil
|
|
}
|
|
|
|
img := image.NewRGBA(image.Rect(0, 0, w, h))
|
|
for y := 0; y < h; y++ {
|
|
for x := 0; x < w; x++ {
|
|
img.Set(x, y, src.At(b.Min.X+x, b.Min.Y+y))
|
|
}
|
|
}
|
|
|
|
maxBorder := minInt(32, minInt(w, h)/20)
|
|
if maxBorder < 2 {
|
|
maxBorder = 2
|
|
}
|
|
|
|
top := 0
|
|
for top < maxBorder && isNearBlackRow(img, top) {
|
|
top++
|
|
}
|
|
bottom := 0
|
|
for bottom < maxBorder && isNearBlackRow(img, h-1-bottom) {
|
|
bottom++
|
|
}
|
|
left := 0
|
|
for left < maxBorder && isNearBlackCol(img, left) {
|
|
left++
|
|
}
|
|
right := 0
|
|
for right < maxBorder && isNearBlackCol(img, w-1-right) {
|
|
right++
|
|
}
|
|
if top == 0 && bottom == 0 && left == 0 && right == 0 {
|
|
return nil
|
|
}
|
|
if top+bottom >= h || left+right >= w {
|
|
return nil
|
|
}
|
|
|
|
for y := 0; y < top; y++ {
|
|
copyRow(img, top, y)
|
|
}
|
|
for y := h - bottom; y < h; y++ {
|
|
copyRow(img, h-bottom-1, y)
|
|
}
|
|
for x := 0; x < left; x++ {
|
|
copyCol(img, left, x)
|
|
}
|
|
for x := w - right; x < w; x++ {
|
|
copyCol(img, w-right-1, x)
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
if err := png.Encode(&buf, img); err != nil {
|
|
return fmt.Errorf("repair border: encode png: %w", err)
|
|
}
|
|
if err := os.WriteFile(path, buf.Bytes(), 0644); err != nil {
|
|
return fmt.Errorf("repair border: write image: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func isNearBlackRow(img *image.RGBA, y int) bool {
|
|
w := img.Bounds().Dx()
|
|
black := 0
|
|
for x := 0; x < w; x++ {
|
|
if isNearBlack(img.RGBAAt(x, y)) {
|
|
black++
|
|
}
|
|
}
|
|
return black*100 >= w*96
|
|
}
|
|
|
|
func isNearBlackCol(img *image.RGBA, x int) bool {
|
|
h := img.Bounds().Dy()
|
|
black := 0
|
|
for y := 0; y < h; y++ {
|
|
if isNearBlack(img.RGBAAt(x, y)) {
|
|
black++
|
|
}
|
|
}
|
|
return black*100 >= h*96
|
|
}
|
|
|
|
func isNearBlack(c color.RGBA) bool {
|
|
return c.A > 0 && c.R <= 18 && c.G <= 18 && c.B <= 18
|
|
}
|
|
|
|
func copyRow(img *image.RGBA, srcY, dstY int) {
|
|
w := img.Bounds().Dx()
|
|
for x := 0; x < w; x++ {
|
|
img.SetRGBA(x, dstY, img.RGBAAt(x, srcY))
|
|
}
|
|
}
|
|
|
|
func copyCol(img *image.RGBA, srcX, dstX int) {
|
|
h := img.Bounds().Dy()
|
|
for y := 0; y < h; y++ {
|
|
img.SetRGBA(dstX, y, img.RGBAAt(srcX, y))
|
|
}
|
|
}
|
|
|
|
func minInt(a, b int) int {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|