diff --git a/backend/internal/provider/openai_image.go b/backend/internal/provider/openai_image.go index 4b4eaec..1ce6895 100644 --- a/backend/internal/provider/openai_image.go +++ b/backend/internal/provider/openai_image.go @@ -3,12 +3,18 @@ package provider import ( "bytes" "context" + "encoding/base64" "encoding/json" "fmt" "image-gen-service/internal/diagnostic" "image-gen-service/internal/model" "io" + "math" + "mime/multipart" "net/http" + "net/textproto" + "regexp" + "strconv" "strings" "time" ) @@ -25,6 +31,12 @@ type openAIImagesGenerationRequest struct { N int `json:"n,omitempty"` } +type openAIImageReference struct { + Name string + Content []byte + MIME string +} + func NewOpenAIImageProvider(config *model.ProviderConfig) (*OpenAIImageProvider, error) { base, err := NewOpenAIProvider(config) if err != nil { @@ -42,9 +54,6 @@ func (p *OpenAIImageProvider) ValidateParams(params map[string]interface{}) erro if strings.TrimSpace(prompt) == "" { return fmt.Errorf("prompt 不能为空") } - if raw, ok := params["reference_images"].([]interface{}); ok && len(raw) > 0 { - return fmt.Errorf("OpenAI Images 当前仅支持文本生图") - } count, ok := toInt(params["count"]) if !ok { @@ -55,10 +64,8 @@ func (p *OpenAIImageProvider) ValidateParams(params map[string]interface{}) erro } size, _ := params["size"].(string) - switch strings.TrimSpace(strings.ToLower(size)) { - case "", "auto", "1024x1024", "1024x1536", "1536x1024": - default: - return fmt.Errorf("size 仅支持 auto、1024x1024、1024x1536、1536x1024") + if !isValidOpenAIImageSize(size) { + return fmt.Errorf("size 仅支持 auto 或 宽x高 格式") } quality, _ := params["quality"].(string) @@ -86,21 +93,35 @@ func (p *OpenAIImageProvider) Generate(ctx context.Context, params map[string]in if err != nil { return nil, err } + refImages, err := collectOpenAIImageReferences(params["reference_images"]) + if err != nil { + return nil, err + } diagnostic.Logf(params, "request_prepare", - "provider=%s model=%s size=%q quality=%q count=%d prompt_hash=%s prompt_preview=%q", + "provider=%s model=%s size=%q quality=%q count=%d ref_image_count=%d prompt_hash=%s prompt_preview=%q", p.Name(), modelID, reqBody.Size, reqBody.Quality, reqBody.N, + len(refImages), diagnostic.PromptHash(promptPreview), diagnostic.Preview(promptPreview, 160), ) - respBytes, headers, err := p.doImagesGenerationRequest(ctx, reqBody, params) - if err != nil { - return nil, err + var respBytes []byte + var headers http.Header + if len(refImages) > 0 { + respBytes, headers, err = p.doImagesEditRequest(ctx, reqBody, refImages, params) + if err != nil { + return nil, err + } + } else { + respBytes, headers, err = p.doImagesGenerationRequest(ctx, reqBody, params) + if err != nil { + return nil, err + } } images, summary, err := p.extractImages(ctx, respBytes) @@ -141,12 +162,9 @@ func (p *OpenAIImageProvider) buildImagesGenerationRequestBody(modelID string, p body := &openAIImagesGenerationRequest{ Model: modelID, Prompt: prompt, - Size: "auto", + Size: resolveOpenAIImageSize(modelID, params), N: 1, } - if size, _ := params["size"].(string); strings.TrimSpace(size) != "" { - body.Size = strings.TrimSpace(strings.ToLower(size)) - } if quality, _ := params["quality"].(string); strings.TrimSpace(quality) != "" { body.Quality = strings.TrimSpace(strings.ToLower(quality)) } @@ -181,7 +199,6 @@ func (p *OpenAIImageProvider) doImagesGenerationRequest(ctx context.Context, bod req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json") req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(p.config.APIKey)) - req.Header.Set("Connection", "close") if strings.TrimSpace(p.userAgent) != "" { req.Header.Set("User-Agent", p.userAgent) } @@ -231,3 +248,354 @@ func (p *OpenAIImageProvider) doImagesGenerationRequest(ctx context.Context, bod return respBody, resp.Header.Clone(), nil } + +func (p *OpenAIImageProvider) doImagesEditRequest(ctx context.Context, body *openAIImagesGenerationRequest, refs []openAIImageReference, params map[string]interface{}) ([]byte, http.Header, error) { + requestURL := strings.TrimRight(strings.TrimSpace(p.apiBase), "/") + "/images/edits" + fields := openAIImageEditFields(body) + imageFieldCount := len(refs) + if imageFieldCount > 2 { + imageFieldCount = 2 + } + diagnostic.Logf(params, "request_payload", + "url=%s fields=%d image_count=%d", + diagnostic.RedactSensitive(requestURL), + len(fields), + imageFieldCount, + ) + + maxRetries := providerMaxRetries(p.config) + var elapsed time.Duration + resp, _, err := doRequestWithRetry(ctx, params, p.Name(), maxRetries, func(attempt int) (*http.Response, error) { + reader, contentType := openAIImageEditBody(fields, refs) + req, buildErr := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, reader) + if buildErr != nil { + return nil, fmt.Errorf("构建 OpenAI Images Edit 请求失败: %w", buildErr) + } + + req.Header.Set("Content-Type", contentType) + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(p.config.APIKey)) + if strings.TrimSpace(p.userAgent) != "" { + req.Header.Set("User-Agent", p.userAgent) + } + + startedAt := time.Now() + resp, doErr := p.httpClient.Do(req) + elapsed = time.Since(startedAt) + return resp, doErr + }) + if err != nil { + return nil, nil, fmt.Errorf("doRequest: error sending request: %w", err) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, resp.Header.Clone(), fmt.Errorf("读取 OpenAI Images Edit 响应失败: %w", err) + } + + requestID := extractRequestIDFromHeaders(resp.Header) + diagnostic.Logf(params, "response_headers", + "status=%s elapsed=%s request_id=%s headers=%q", + resp.Status, + elapsed, + requestID, + diagnostic.Preview(strings.Join(headerLines(resp.Header), " | "), 1000), + ) + diagnostic.Logf(params, "response_body", + "status=%s elapsed=%s request_id=%s body=%q", + resp.Status, + elapsed, + requestID, + diagnostic.RedactSensitive(string(respBody)), + ) + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + bodyPreview := diagnostic.Preview(parseOpenAIError(respBody), 1200) + if requestID == "" { + requestID = diagnostic.ExtractRequestID(string(respBody)) + } + return nil, resp.Header.Clone(), fmt.Errorf("OpenAI HTTP %d request_id=%s body=%s", resp.StatusCode, requestID, bodyPreview) + } + + if len(respBody) == 0 { + return nil, resp.Header.Clone(), fmt.Errorf("接口未返回内容") + } + + return respBody, resp.Header.Clone(), nil +} + +func openAIImageEditFields(body *openAIImagesGenerationRequest) map[string]string { + fields := map[string]string{ + "model": body.Model, + "prompt": body.Prompt, + "size": body.Size, + "n": strconv.Itoa(body.N), + } + // Edits 只发送官方通用字段,避免代理或官方接口因未知字段返回 400。 + if body.Quality != "" { + fields["quality"] = body.Quality + } + return fields +} + +func openAIImageEditBody(fields map[string]string, refs []openAIImageReference) (io.Reader, string) { + reader, writer := io.Pipe() + multipartWriter := multipart.NewWriter(writer) + contentType := multipartWriter.FormDataContentType() + + go func() { + err := writeOpenAIImageEditMultipart(multipartWriter, fields, refs) + if closeErr := multipartWriter.Close(); err == nil { + err = closeErr + } + _ = writer.CloseWithError(err) + }() + + return reader, contentType +} + +func writeOpenAIImageEditMultipart(writer *multipart.Writer, fields map[string]string, refs []openAIImageReference) error { + for key, value := range fields { + if err := writer.WriteField(key, value); err != nil { + return fmt.Errorf("构建 OpenAI Images Edit 字段失败: %w", err) + } + } + + for idx, ref := range refs { + fieldName := "image" + if idx == 1 { + fieldName = "mask" + } else if idx > 1 { + break + } + + header := make(textproto.MIMEHeader) + header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, escapeMultipartFilename(ref.Name))) + header.Set("Content-Type", ref.MIME) + part, err := writer.CreatePart(header) + if err != nil { + return fmt.Errorf("构建 OpenAI Images Edit 图片字段失败: %w", err) + } + if _, err := part.Write(ref.Content); err != nil { + return fmt.Errorf("写入 OpenAI Images Edit 图片字段失败: %w", err) + } + } + + return nil +} + +func collectOpenAIImageReferences(raw interface{}) ([]openAIImageReference, error) { + refImgs, ok := raw.([]interface{}) + if !ok || len(refImgs) == 0 { + return nil, nil + } + + refs := make([]openAIImageReference, 0, len(refImgs)) + for idx, ref := range refImgs { + var imgBytes []byte + switch v := ref.(type) { + case string: + base64Data := v + if strings.Contains(base64Data, ",") { + base64Data = strings.SplitN(base64Data, ",", 2)[1] + } + decoded, err := base64.StdEncoding.DecodeString(base64Data) + if err != nil { + return nil, fmt.Errorf("解码第 %d 张参考图失败: %w", idx+1, err) + } + imgBytes = decoded + case []byte: + imgBytes = v + default: + continue + } + + if len(imgBytes) == 0 { + continue + } + mimeType := http.DetectContentType(imgBytes) + if mimeType != "image/png" { + return nil, fmt.Errorf("第 %d 张参考图不是有效的 PNG 图片,OpenAI Edit 仅支持 PNG 格式", idx+1) + } + refs = append(refs, openAIImageReference{ + Name: fmt.Sprintf("reference-%d.png", idx+1), + Content: imgBytes, + MIME: mimeType, + }) + } + + return refs, nil +} + +func escapeMultipartFilename(name string) string { + return strings.NewReplacer("\\", "\\\\", `"`, "\\\"").Replace(name) +} + +func isValidOpenAIImageSize(size string) bool { + size = strings.TrimSpace(strings.ToLower(size)) + if size == "" || size == "auto" { + return true + } + return regexp.MustCompile(`^[1-9][0-9]{1,4}x[1-9][0-9]{1,4}$`).MatchString(size) +} + +func resolveOpenAIImageSize(modelID string, params map[string]interface{}) string { + aspectRatio := firstStringParam(params, "aspect_ratio", "aspectRatio", "aspect") + resolution := firstStringParam(params, "resolution_level", "imageSize", "image_size", "resolution") + model := strings.ToLower(strings.TrimSpace(modelID)) + if size, _ := params["size"].(string); strings.TrimSpace(size) != "" { + return normalizeExplicitOpenAIImageSize(model, strings.TrimSpace(strings.ToLower(size)), aspectRatio, resolution) + } + + if strings.Contains(model, "dall-e-3") { + return resolveDalle3Size(aspectRatio) + } + if strings.Contains(model, "dall-e-2") { + return "1024x1024" + } + if strings.Contains(model, "gpt-image-2") { + return computeDynamicOpenAIImageSize(aspectRatio, resolution) + } + return resolveStandardGPTImageSize(aspectRatio) +} + +func normalizeExplicitOpenAIImageSize(model, size, aspectRatio, resolution string) string { + if strings.Contains(model, "dall-e-3") { + switch size { + case "1024x1024", "1792x1024", "1024x1792": + return size + default: + return resolveDalle3Size(aspectRatio) + } + } + + if strings.Contains(model, "dall-e-2") { + switch size { + case "256x256", "512x512", "1024x1024": + return size + default: + return "1024x1024" + } + } + + if strings.Contains(model, "gpt-image-2") { + if size == "auto" { + return computeDynamicOpenAIImageSize(aspectRatio, resolution) + } + return size + } + + switch size { + case "1024x1024", "1536x1024", "1024x1536": + return size + default: + return resolveStandardGPTImageSize(aspectRatio) + } +} + +func firstStringParam(params map[string]interface{}, keys ...string) string { + for _, key := range keys { + if val, ok := params[key].(string); ok && strings.TrimSpace(val) != "" { + return strings.TrimSpace(val) + } + } + return "" +} + +func resolveDalle3Size(aspectRatio string) string { + w, h, ok := parseAspectRatio(aspectRatio) + if !ok { + return "1024x1024" + } + if w == h { + return "1024x1024" + } + if w > h { + return "1792x1024" + } + return "1024x1792" +} + +func resolveStandardGPTImageSize(aspectRatio string) string { + w, h, ok := parseAspectRatio(aspectRatio) + if !ok { + return "auto" + } + if w == h { + return "1024x1024" + } + if w > h { + return "1536x1024" + } + return "1024x1536" +} + +// 根据用户选择的宽高比和分辨率档位,为 gpt-image-2 代理计算实际 WxH。 +// 计算流程:先确定长边,再按宽高比推导短边,最后做 16 像素对齐和总像素上限保护。 +func computeDynamicOpenAIImageSize(aspectRatio, resolution string) string { + wRatio, hRatio, ok := parseAspectRatio(aspectRatio) + if !ok { + return "auto" + } + + // gpt-image-2 代理支持更灵活的 WxH。这里用用户选择的分辨率作为长边, + // 再按比例计算短边,并统一向下取 16 的倍数,避免上游拒绝非对齐尺寸。 + longEdge := 2048 + switch strings.ToUpper(strings.TrimSpace(resolution)) { + case "1K": + longEdge = 1280 + case "2K", "": + longEdge = 2048 + case "4K": + longEdge = 3840 + } + + width := longEdge + height := int(float64(width) * float64(hRatio) / float64(wRatio)) + if hRatio > wRatio { + height = longEdge + width = int(float64(height) * float64(wRatio) / float64(hRatio)) + } + + width = roundDownToMultiple(width, 16) + height = roundDownToMultiple(height, 16) + + // 按 banana-slides 的经验保留官方/代理常见的像素上限保护: + // 超过约 8.29MP 时等比缩小,再次对齐到 16 的倍数。 + const maxPixels = 8294400 + if width*height > maxPixels { + scale := math.Sqrt(float64(maxPixels) / float64(width*height)) + width = roundDownToMultiple(int(float64(width)*scale), 16) + height = roundDownToMultiple(int(float64(height)*scale), 16) + } + + if width < 16 || height < 16 { + return "auto" + } + return fmt.Sprintf("%dx%d", width, height) +} + +func parseAspectRatio(value string) (int, int, bool) { + parts := strings.Split(strings.TrimSpace(value), ":") + if len(parts) != 2 { + return 0, 0, false + } + w, errW := strconv.Atoi(strings.TrimSpace(parts[0])) + h, errH := strconv.Atoi(strings.TrimSpace(parts[1])) + if errW != nil || errH != nil || w <= 0 || h <= 0 { + return 0, 0, false + } + return w, h, true +} + +func roundDownToMultiple(value, multiple int) int { + if multiple <= 0 { + return value + } + rounded := (value / multiple) * multiple + if rounded < multiple { + return multiple + } + return rounded +} diff --git a/backend/internal/provider/openai_image_test.go b/backend/internal/provider/openai_image_test.go new file mode 100644 index 0000000..df308eb --- /dev/null +++ b/backend/internal/provider/openai_image_test.go @@ -0,0 +1,239 @@ +package provider + +import ( + "encoding/base64" + "encoding/json" + "image-gen-service/internal/model" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +const tinyPNGBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=" + +func TestResolveOpenAIImageSize(t *testing.T) { + tests := []struct { + name string + model string + params map[string]interface{} + want string + }{ + { + name: "explicit size wins", + model: "gpt-image-2-all", + params: map[string]interface{}{ + "size": "1024x1536", + "aspectRatio": "16:9", + "imageSize": "2K", + }, + want: "1024x1536", + }, + { + name: "gpt image 2 uses dynamic 16:9 2K", + model: "gpt-image-2-all", + params: map[string]interface{}{ + "aspectRatio": "16:9", + "imageSize": "2K", + }, + want: "2048x1152", + }, + { + name: "standard gpt image uses supported landscape size", + model: "gpt-image-1", + params: map[string]interface{}{ + "aspectRatio": "16:9", + "imageSize": "2K", + }, + want: "1536x1024", + }, + { + name: "dalle 3 maps portrait", + model: "dall-e-3", + params: map[string]interface{}{ + "aspectRatio": "9:16", + }, + want: "1024x1792", + }, + { + name: "dalle 3 normalizes explicit auto", + model: "dall-e-3", + params: map[string]interface{}{ + "size": "auto", + "aspectRatio": "16:9", + }, + want: "1792x1024", + }, + { + name: "dalle 2 normalizes unsupported explicit size", + model: "dall-e-2", + params: map[string]interface{}{ + "size": "1536x1024", + }, + want: "1024x1024", + }, + { + name: "standard gpt image normalizes unsupported explicit size", + model: "gpt-image-1", + params: map[string]interface{}{ + "size": "2048x1152", + "aspectRatio": "16:9", + }, + want: "1536x1024", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := resolveOpenAIImageSize(tt.model, tt.params); got != tt.want { + t.Fatalf("resolveOpenAIImageSize() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestOpenAIImageProviderGenerateUsesGenerations(t *testing.T) { + var seenPath string + var seenAuth string + var seenBody openAIImagesGenerationRequest + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seenPath = r.URL.Path + seenAuth = r.Header.Get("Authorization") + if err := json.NewDecoder(r.Body).Decode(&seenBody); err != nil { + t.Fatalf("decode request body: %v", err) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]string{{"b64_json": tinyPNGBase64}}, + }) + })) + defer server.Close() + + p, err := NewOpenAIImageProvider(&model.ProviderConfig{ + ProviderName: "openai-image", + APIBase: server.URL, + APIKey: "test-key", + TimeoutSeconds: 5, + }) + if err != nil { + t.Fatalf("NewOpenAIImageProvider: %v", err) + } + + result, err := p.Generate(t.Context(), map[string]interface{}{ + "prompt": "test prompt", + "model_id": "gpt-image-2-all", + "aspectRatio": "16:9", + "imageSize": "2K", + "count": 1, + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if seenPath != "/v1/images/generations" { + t.Fatalf("path = %q, want /v1/images/generations", seenPath) + } + if seenAuth != "Bearer test-key" { + t.Fatalf("Authorization = %q", seenAuth) + } + if seenBody.Model != "gpt-image-2-all" || seenBody.Size != "2048x1152" || seenBody.Prompt != "test prompt" { + t.Fatalf("unexpected body: %+v", seenBody) + } + if len(result.Images) != 1 { + t.Fatalf("image count = %d, want 1", len(result.Images)) + } +} + +func TestOpenAIImageProviderGenerateWithReferenceUsesEdits(t *testing.T) { + refBytes, err := base64.StdEncoding.DecodeString(tinyPNGBase64) + if err != nil { + t.Fatalf("decode ref: %v", err) + } + + var seenPath string + var seenFields = map[string]string{} + var seenImageFileCount int + var seenMaskFileCount int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seenPath = r.URL.Path + if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") { + t.Fatalf("Content-Type = %q", r.Header.Get("Content-Type")) + } + if err := r.ParseMultipartForm(2 << 20); err != nil { + t.Fatalf("ParseMultipartForm: %v", err) + } + for key, values := range r.MultipartForm.Value { + if len(values) > 0 { + seenFields[key] = values[0] + } + } + seenImageFileCount = len(r.MultipartForm.File["image"]) + seenMaskFileCount = len(r.MultipartForm.File["mask"]) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]string{{"b64_json": tinyPNGBase64}}, + }) + })) + defer server.Close() + + p, err := NewOpenAIImageProvider(&model.ProviderConfig{ + ProviderName: "openai-image", + APIBase: server.URL, + APIKey: "test-key", + TimeoutSeconds: 5, + }) + if err != nil { + t.Fatalf("NewOpenAIImageProvider: %v", err) + } + + result, err := p.Generate(t.Context(), map[string]interface{}{ + "prompt": "edit prompt", + "model_id": "gpt-image-2-all", + "aspect_ratio": "1:1", + "resolution_level": "1K", + "count": 1, + "reference_images": []interface{}{refBytes, refBytes, refBytes}, + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if seenPath != "/v1/images/edits" { + t.Fatalf("path = %q, want /v1/images/edits", seenPath) + } + if seenFields["model"] != "gpt-image-2-all" || seenFields["prompt"] != "edit prompt" || seenFields["size"] != "1280x1280" { + t.Fatalf("unexpected fields: %+v", seenFields) + } + if _, ok := seenFields["input_fidelity"]; ok { + t.Fatalf("input_fidelity should not be sent to edits: %+v", seenFields) + } + if seenImageFileCount != 1 { + t.Fatalf("image file count = %d, want 1", seenImageFileCount) + } + if seenMaskFileCount != 1 { + t.Fatalf("mask file count = %d, want 1", seenMaskFileCount) + } + if len(result.Images) != 1 { + t.Fatalf("image count = %d, want 1", len(result.Images)) + } +} + +func TestOpenAIImageProviderRejectsNonPNGReference(t *testing.T) { + p, err := NewOpenAIImageProvider(&model.ProviderConfig{ + ProviderName: "openai-image", + APIBase: "http://example.test/v1", + APIKey: "test-key", + TimeoutSeconds: 5, + }) + if err != nil { + t.Fatalf("NewOpenAIImageProvider: %v", err) + } + + _, err = p.Generate(t.Context(), map[string]interface{}{ + "prompt": "edit prompt", + "model_id": "gpt-image-2-all", + "aspect_ratio": "1:1", + "resolution_level": "1K", + "count": 1, + "reference_images": []interface{}{[]byte("not an image")}, + }) + if err == nil || !strings.Contains(err.Error(), "OpenAI Edit 仅支持 PNG 格式") { + t.Fatalf("Generate error = %v, want PNG validation error", err) + } +}