63 lines
2.0 KiB
Go
63 lines
2.0 KiB
Go
|
package memoir
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"data-terminal-cloud/internal/extraction"
|
||
|
|
"encoding/base64"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type ArkIllustrator struct {
|
||
|
|
BaseURL, APIKey, Model string
|
||
|
|
Client *http.Client
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewIllustrator(base, key, model string) *ArkIllustrator {
|
||
|
|
return &ArkIllustrator{strings.TrimRight(base, "/"), key, model, &http.Client{Timeout: 120 * time.Second}}
|
||
|
|
}
|
||
|
|
func (a *ArkIllustrator) Generate(ctx context.Context, prompt string) ([]byte, error) {
|
||
|
|
body, _ := json.Marshal(map[string]any{"model": a.Model, "prompt": prompt + "。温暖的水彩情景插画，米色纸张肌理，禁止文字和水印以外的标记。这是基于叙述的想象性插画，不是真实历史照片。", "size": "2K", "output_format": "jpeg", "response_format": "b64_json", "watermark": true})
|
||
|
|
req, e := http.NewRequestWithContext(ctx, "POST", a.BaseURL+"/images/generations", bytes.NewReader(body))
|
||
|
|
if e != nil {
|
||
|
|
return nil, e
|
||
|
|
}
|
||
|
|
req.Header.Set("Authorization", "Bearer "+a.APIKey)
|
||
|
|
req.Header.Set("Content-Type", "application/json")
|
||
|
|
resp, e := a.Client.Do(req)
|
||
|
|
if e != nil {
|
||
|
|
return nil, errors.New("MEMOIR_IMAGE_UNAVAILABLE")
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
if resp.StatusCode != 200 {
|
||
|
|
var failure struct {
|
||
|
|
Error struct {
|
||
|
|
Code string `json:"code"`
|
||
|
|
} `json:"error"`
|
||
|
|
}
|
||
|
|
_ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&failure)
|
||
|
|
if failure.Error.Code == "AccountOverdueError" {
|
||
|
|
return nil, extraction.ErrModelAccountOverdue
|
||
|
|
}
|
||
|
|
return nil, errors.New("MEMOIR_IMAGE_REJECTED")
|
||
|
|
}
|
||
|
|
var out struct {
|
||
|
|
Data []struct {
|
||
|
|
Base64 string `json:"b64_json"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
if json.NewDecoder(io.LimitReader(resp.Body, 16<<20)).Decode(&out) != nil || len(out.Data) != 1 {
|
||
|
|
return nil, errors.New("MEMOIR_IMAGE_INVALID")
|
||
|
|
}
|
||
|
|
b, e := base64.StdEncoding.DecodeString(out.Data[0].Base64)
|
||
|
|
if e != nil || len(b) > 10<<20 || http.DetectContentType(b) != "image/jpeg" {
|
||
|
|
return nil, errors.New("MEMOIR_IMAGE_INVALID")
|
||
|
|
}
|
||
|
|
return b, nil
|
||
|
|
}
|