wonder
50f95adb98
- ImageGenConfig 新增 timeout/max_retries/retry_delay 配置项（默认 120s/2次/5s） - 替换 http.DefaultClient 为带超时的 imageHTTPClient - callImageAPI 对 5xx 错误自动重试，避免 504 等瞬时故障直接失败
319 lines
9.2 KiB
Go
Executable File
319 lines
9.2 KiB
Go
Executable File
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"image/color"
|
||
"image/png"
|
||
"io"
|
||
"log"
|
||
"mime/multipart"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"gen2d/internal/config"
|
||
)
|
||
|
||
// imgCfg 保存文生图配置，由 main 通过 InitImageGenConfig 注入。
|
||
var imgCfg config.ImageGenConfig
|
||
|
||
// imageHTTPClient 带超时的 HTTP 客户端，由 InitImageGenConfig 初始化。
|
||
var imageHTTPClient *http.Client
|
||
|
||
// InitImageGenConfig 注入文生图配置。
|
||
func InitImageGenConfig(cfg config.ImageGenConfig) {
|
||
imgCfg = cfg
|
||
timeout := time.Duration(cfg.Timeout) * time.Second
|
||
if timeout <= 0 {
|
||
timeout = 120 * time.Second
|
||
}
|
||
imageHTTPClient = &http.Client{Timeout: timeout}
|
||
}
|
||
|
||
// ======================== OpenAI 兼容 Images API 类型 ========================
|
||
|
||
// imageGenRequest OpenAI 兼容文生图请求体。
|
||
type imageGenRequest struct {
|
||
Model string `json:"model"`
|
||
Prompt string `json:"prompt"`
|
||
N int `json:"n,omitempty"`
|
||
Size string `json:"size,omitempty"`
|
||
Quality string `json:"quality,omitempty"`
|
||
}
|
||
|
||
// imageGenResponse OpenAI 兼容文生图响应体。
|
||
type imageGenResponse struct {
|
||
Data []struct {
|
||
URL string `json:"url"`
|
||
B64JSON string `json:"b64_json"`
|
||
} `json:"data"`
|
||
}
|
||
|
||
// ======================== 文生图 ========================
|
||
|
||
// GenerateImages 调用 OpenAI 兼容 Images API 生成图片，未配置 key 时回退 mock。
|
||
func GenerateImages(ctx context.Context, prompt string, params AssetParams) ([]GeneratedImage, error) {
|
||
count := 1
|
||
if params.Frames.Directions > 0 && params.Frames.FramesPerDirection > 0 {
|
||
count = params.Frames.Directions * params.Frames.FramesPerDirection
|
||
}
|
||
|
||
if imgCfg.APIKey != "" {
|
||
width, height := imgCfg.Width, imgCfg.Height
|
||
if params.Resolution > 0 {
|
||
width = params.Resolution
|
||
height = params.Resolution
|
||
}
|
||
log.Println("[inference] calling image gen API")
|
||
return callImageAPI(ctx, prompt, count, width, height)
|
||
}
|
||
|
||
size := params.Resolution
|
||
if size <= 0 {
|
||
size = 64
|
||
}
|
||
log.Println("[inference] image API key not configured, using mock")
|
||
return generateMockImages(size, count)
|
||
}
|
||
|
||
// callImageAPI 调用 OpenAI 兼容 Images API，5xx 错误自动重试。
|
||
func callImageAPI(ctx context.Context, prompt string, count, width, height int) ([]GeneratedImage, error) {
|
||
reqBody := imageGenRequest{
|
||
Model: imgCfg.Model,
|
||
Prompt: prompt,
|
||
N: count,
|
||
Size: fmt.Sprintf("%dx%d", width, height),
|
||
Quality: imgCfg.Quality,
|
||
}
|
||
|
||
body, err := json.Marshal(reqBody)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("marshal request: %w", err)
|
||
}
|
||
|
||
url := strings.TrimRight(imgCfg.BaseURL, "/") + "/images/generations"
|
||
maxRetries := imgCfg.MaxRetries
|
||
if maxRetries <= 0 {
|
||
maxRetries = 2
|
||
}
|
||
retryDelay := time.Duration(imgCfg.RetryDelay) * time.Second
|
||
if retryDelay <= 0 {
|
||
retryDelay = 5 * time.Second
|
||
}
|
||
|
||
var lastErr error
|
||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||
if attempt > 0 {
|
||
log.Printf("[inference] retrying image API (attempt %d/%d)", attempt, maxRetries)
|
||
select {
|
||
case <-ctx.Done():
|
||
return nil, fmt.Errorf("context cancelled during retry: %w", ctx.Err())
|
||
case <-time.After(retryDelay):
|
||
}
|
||
}
|
||
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create request: %w", err)
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("Authorization", "Bearer "+imgCfg.APIKey)
|
||
|
||
resp, err := imageHTTPClient.Do(req)
|
||
if err != nil {
|
||
lastErr = fmt.Errorf("send request: %w", err)
|
||
continue
|
||
}
|
||
|
||
if resp.StatusCode == http.StatusOK {
|
||
return parseImageResponse(ctx, resp.Body, width, height)
|
||
}
|
||
|
||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||
resp.Body.Close()
|
||
|
||
if resp.StatusCode >= 500 {
|
||
lastErr = fmt.Errorf("image api error %d: %s", resp.StatusCode, string(b))
|
||
continue
|
||
}
|
||
return nil, fmt.Errorf("image api error %d: %s", resp.StatusCode, string(b))
|
||
}
|
||
|
||
return nil, lastErr
|
||
}
|
||
|
||
// parseImageResponse 解析 OpenAI 兼容图片响应（b64_json 或 url）。
|
||
func parseImageResponse(ctx context.Context, r io.Reader, width, height int) ([]GeneratedImage, error) {
|
||
var genResp imageGenResponse
|
||
if err := json.NewDecoder(r).Decode(&genResp); err != nil {
|
||
return nil, fmt.Errorf("decode response: %w", err)
|
||
}
|
||
|
||
images := make([]GeneratedImage, 0, len(genResp.Data))
|
||
for i, d := range genResp.Data {
|
||
var data []byte
|
||
switch {
|
||
case d.B64JSON != "":
|
||
var err error
|
||
data, err = base64.StdEncoding.DecodeString(d.B64JSON)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("decode base64 image %d: %w", i, err)
|
||
}
|
||
case d.URL != "":
|
||
var err error
|
||
data, err = downloadImage(ctx, d.URL)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("download image %d: %w", i, err)
|
||
}
|
||
default:
|
||
return nil, fmt.Errorf("image %d: no data or url in response", i)
|
||
}
|
||
images = append(images, GeneratedImage{
|
||
Data: data,
|
||
Width: width,
|
||
Height: height,
|
||
Format: "png",
|
||
})
|
||
}
|
||
return images, nil
|
||
}
|
||
|
||
// downloadImage 从 URL 下载图片数据。
|
||
func downloadImage(ctx context.Context, url string) ([]byte, error) {
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create download request: %w", err)
|
||
}
|
||
resp, err := imageHTTPClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("download: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
return nil, fmt.Errorf("download status %d", resp.StatusCode)
|
||
}
|
||
return io.ReadAll(resp.Body)
|
||
}
|
||
|
||
// ======================== 图片编辑 ========================
|
||
|
||
// EditImages 图片编辑接口，调用 OpenAI 兼容 Images Edits API（multipart/form-data）。
|
||
func EditImages(ctx context.Context, imageData []byte, prompt string, count int) ([]GeneratedImage, error) {
|
||
if imgCfg.APIKey == "" {
|
||
return nil, fmt.Errorf("image API key not configured")
|
||
}
|
||
|
||
var buf bytes.Buffer
|
||
writer := multipart.NewWriter(&buf)
|
||
|
||
part, err := writer.CreateFormFile("image", "image.png")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create form file: %w", err)
|
||
}
|
||
if _, err := part.Write(imageData); err != nil {
|
||
return nil, fmt.Errorf("write image data: %w", err)
|
||
}
|
||
|
||
writer.WriteField("prompt", prompt)
|
||
writer.WriteField("model", imgCfg.Model)
|
||
writer.WriteField("n", fmt.Sprintf("%d", count))
|
||
writer.WriteField("size", fmt.Sprintf("%dx%d", imgCfg.Width, imgCfg.Height))
|
||
|
||
if err := writer.Close(); err != nil {
|
||
return nil, fmt.Errorf("close multipart writer: %w", err)
|
||
}
|
||
|
||
url := strings.TrimRight(imgCfg.BaseURL, "/") + "/images/edits"
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, &buf)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create request: %w", err)
|
||
}
|
||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||
req.Header.Set("Authorization", "Bearer "+imgCfg.APIKey)
|
||
|
||
resp, err := imageHTTPClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("send request: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||
return nil, fmt.Errorf("image edit api error %d: %s", resp.StatusCode, string(b))
|
||
}
|
||
|
||
return parseImageResponse(ctx, resp.Body, imgCfg.Width, imgCfg.Height)
|
||
}
|
||
|
||
// ======================== 质检 ========================
|
||
|
||
var QualityChecker = defaultCheckQuality
|
||
|
||
func CheckQuality(ctx context.Context, images []GeneratedImage, style map[string]string) (bool, string, error) {
|
||
return QualityChecker(ctx, images, style)
|
||
}
|
||
|
||
func defaultCheckQuality(ctx context.Context, images []GeneratedImage, style map[string]string) (bool, string, error) {
|
||
return true, "", nil
|
||
}
|
||
|
||
func NewCountedQualityChecker(passOnRetry int) func(context.Context, []GeneratedImage, map[string]string) (bool, string, error) {
|
||
var callCount int
|
||
return func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
|
||
callCount++
|
||
if callCount >= passOnRetry {
|
||
return true, "", nil
|
||
}
|
||
return false, fmt.Sprintf("style inconsistent (attempt %d)", callCount), nil
|
||
}
|
||
}
|
||
|
||
func AlwaysFailQualityChecker() func(context.Context, []GeneratedImage, map[string]string) (bool, string, error) {
|
||
return func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
|
||
return false, "style inconsistent", nil
|
||
}
|
||
}
|
||
|
||
// ======================== Mock 回退 ========================
|
||
|
||
func generateMockImages(size, count int) ([]GeneratedImage, error) {
|
||
images := make([]GeneratedImage, count)
|
||
for i := 0; i < count; i++ {
|
||
data, err := generateMockImage(size, i)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("generate mock image %d: %w", i, err)
|
||
}
|
||
images[i] = GeneratedImage{Data: data, Width: size, Height: size, Format: "png"}
|
||
}
|
||
return images, nil
|
||
}
|
||
|
||
func generateMockImage(size int, seed int) ([]byte, error) {
|
||
img := image.NewRGBA(image.Rect(0, 0, size, size))
|
||
r := uint8((seed*47 + 13) % 256)
|
||
g := uint8((seed*83 + 37) % 256)
|
||
b := uint8((seed*61 + 71) % 256)
|
||
for y := 0; y < size; y++ {
|
||
for x := 0; x < size; x++ {
|
||
img.Set(x, y, color.RGBA{R: r, G: g, B: b, A: 255})
|
||
}
|
||
}
|
||
var buf bytes.Buffer
|
||
if err := png.Encode(&buf, img); err != nil {
|
||
return nil, err
|
||
}
|
||
return buf.Bytes(), nil
|
||
}
|
||
|
||
func generateRandomBytes(n int) ([]byte, error) {
|
||
b := make([]byte, n)
|
||
_, err := rand.Read(b)
|
||
return b, err
|
||
}
|