diff --git a/dto/openai_image.go b/dto/openai_image.go index 547b0d18420f..cdde7b37e612 100644 --- a/dto/openai_image.go +++ b/dto/openai_image.go @@ -2,6 +2,7 @@ package dto import ( "encoding/json" + "fmt" "reflect" "strings" @@ -182,3 +183,99 @@ type ImageData struct { B64Json string `json:"b64_json"` RevisedPrompt string `json:"revised_prompt"` } + +func (i *ImageRequest) InputImageSources() ([]types.FileSource, error) { + values := make([]string, 0) + for _, raw := range []json.RawMessage{i.Image, i.Images} { + parsed, err := parseImageSourceValues(raw) + if err != nil { + return nil, err + } + values = append(values, parsed...) + } + return fileSourcesFromImageValues(values), nil +} + +func parseImageSourceValues(raw json.RawMessage) ([]string, error) { + if len(raw) == 0 || common.GetJsonType(raw) == "null" { + return nil, nil + } + + var values []string + switch common.GetJsonType(raw) { + case "string": + var value string + if err := common.Unmarshal(raw, &value); err != nil { + return nil, err + } + values = append(values, value) + case "array": + var items []json.RawMessage + if err := common.Unmarshal(raw, &items); err != nil { + return nil, err + } + for _, item := range items { + itemValues, err := parseImageSourceValues(item) + if err != nil { + return nil, err + } + values = append(values, itemValues...) + } + case "object": + value, err := parseImageSourceObject(raw) + if err != nil { + return nil, err + } + if value != "" { + values = append(values, value) + } + default: + return nil, fmt.Errorf("image input must be a string, object, or array") + } + + return values, nil +} + +func fileSourcesFromImageValues(values []string) []types.FileSource { + sources := make([]types.FileSource, 0, len(values)) + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" { + continue + } + sources = append(sources, types.NewFileSourceFromData(value, "")) + } + return sources +} + +func parseImageSourceObject(raw json.RawMessage) (string, error) { + var item map[string]json.RawMessage + if err := common.Unmarshal(raw, &item); err != nil { + return "", err + } + + for _, key := range []string{"url", "image_url", "b64_json", "base64", "data"} { + rawValue, ok := item[key] + if !ok || common.GetJsonType(rawValue) == "null" { + continue + } + if key == "image_url" && common.GetJsonType(rawValue) == "object" { + if value, err := parseImageSourceObject(rawValue); err != nil || value != "" { + return value, err + } + continue + } + if common.GetJsonType(rawValue) != "string" { + continue + } + var value string + if err := common.Unmarshal(rawValue, &value); err != nil { + return "", err + } + if strings.TrimSpace(value) != "" { + return value, nil + } + } + + return "", nil +} diff --git a/dto/openai_image_test.go b/dto/openai_image_test.go new file mode 100644 index 000000000000..c5c493628975 --- /dev/null +++ b/dto/openai_image_test.go @@ -0,0 +1,51 @@ +package dto + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/types" + "github.com/stretchr/testify/require" +) + +func TestImageRequestInputImageSources(t *testing.T) { + raw := []byte(`{ + "model":"gemini-2.5-flash-image", + "prompt":"edit", + "image":"https://example.com/input.png", + "images":[ + "data:image/png;base64,aW1hZ2U=", + {"url":"https://example.com/second.webp"}, + {"image_url":{"url":"https://example.com/openai-style.jpg"}}, + {"b64_json":"aW1hZ2Uy"} + ] + }`) + + var req ImageRequest + require.NoError(t, common.Unmarshal(raw, &req)) + + sources, err := req.InputImageSources() + require.NoError(t, err) + require.Len(t, sources, 5) + + _, ok := sources[0].(*types.URLSource) + require.True(t, ok) + require.Equal(t, "https://example.com/input.png", sources[0].GetRawData()) + + _, ok = sources[1].(*types.Base64Source) + require.True(t, ok) + require.Equal(t, "data:image/png;base64,aW1hZ2U=", sources[1].GetRawData()) + + require.Equal(t, "https://example.com/second.webp", sources[2].GetRawData()) + require.Equal(t, "https://example.com/openai-style.jpg", sources[3].GetRawData()) + require.Equal(t, "aW1hZ2Uy", sources[4].GetRawData()) +} + +func TestImageRequestInputImageSourcesRejectsScalarJSON(t *testing.T) { + var req ImageRequest + require.NoError(t, common.Unmarshal([]byte(`{"model":"m","prompt":"p","image":123}`), &req)) + + _, err := req.InputImageSources() + require.Error(t, err) + require.Contains(t, err.Error(), "image input must be") +} diff --git a/relay/channel/gemini/adaptor.go b/relay/channel/gemini/adaptor.go index 680c4ee484ec..8fb21727363d 100644 --- a/relay/channel/gemini/adaptor.go +++ b/relay/channel/gemini/adaptor.go @@ -7,11 +7,13 @@ import ( "net/http" "strings" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel/openai" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/model_setting" "github.com/QuantumNous/new-api/setting/reasoning" "github.com/QuantumNous/new-api/types" @@ -58,10 +60,26 @@ func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInf } func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) { + sources, err := request.InputImageSources() + if err != nil { + return nil, fmt.Errorf("invalid image input: %w", err) + } + + if model_setting.IsGeminiModelSupportImagine(info.UpstreamModelName) { + return convertOpenAIImageRequestToGeminiGenerateContent(c, info, request, sources) + } + + if len(sources) > 0 { + return nil, errors.New("input images are supported only by Gemini image generation models") + } + if !strings.HasPrefix(info.UpstreamModelName, "imagen") { - return nil, errors.New("not supported model for image generation, only imagen models are supported") + return nil, errors.New("not supported model for image generation, only imagen and Gemini image models are supported") } + return convertOpenAIImageRequestToImagenPredict(request), nil +} +func convertOpenAIImageRequestToImagenPredict(request dto.ImageRequest) dto.GeminiImageRequest { // convert size to aspect ratio but allow user to specify aspect ratio aspectRatio := "1:1" // default aspect ratio size := strings.TrimSpace(request.Size) @@ -120,7 +138,119 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf geminiRequest.Parameters.ImageSize = imageSize } - return geminiRequest, nil + return geminiRequest +} + +func convertOpenAIImageRequestToGeminiGenerateContent(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest, sources []types.FileSource) (*dto.GeminiChatRequest, error) { + if request.Stream != nil && *request.Stream { + return nil, errors.New("streaming is not supported for Gemini image generation through images/generations") + } + + parts := make([]dto.GeminiPart, 0, len(sources)+1) + if strings.TrimSpace(request.Prompt) != "" { + parts = append(parts, dto.GeminiPart{Text: request.Prompt}) + } + for _, source := range sources { + base64Data, mimeType, err := service.GetBase64Data(c, source, "formatting image generation input for Gemini") + if err != nil { + return nil, fmt.Errorf("get image input from '%s' failed: %w", source.GetIdentifier(), err) + } + normalizedMimeType := strings.ToLower(mimeType) + if !strings.HasPrefix(normalizedMimeType, "image/") { + return nil, fmt.Errorf("mime type is not supported for Gemini image generation: '%s', url: '%s'", mimeType, source.GetIdentifier()) + } + if _, ok := geminiSupportedMimeTypes[normalizedMimeType]; !ok { + return nil, fmt.Errorf("mime type is not supported by Gemini: '%s', url: '%s', supported types are: %v", mimeType, source.GetIdentifier(), getSupportedMimeTypesList()) + } + parts = append(parts, dto.GeminiPart{ + InlineData: &dto.GeminiInlineData{ + MimeType: mimeType, + Data: base64Data, + }, + }) + } + if len(parts) == 0 { + return nil, errors.New("prompt or image is required") + } + + if request.N != nil && *request.N > 1 { + return nil, errors.New("Gemini image generation supports only n=1") + } + + info.RelayMode = constant.RelayModeImagesGenerations + geminiRequest := dto.GeminiChatRequest{ + Contents: []dto.GeminiChatContent{ + { + Role: "user", + Parts: parts, + }, + }, + GenerationConfig: dto.GeminiChatGenerationConfig{ + ResponseModalities: []string{"TEXT", "IMAGE"}, + }, + SafetySettings: make([]dto.GeminiChatSafetySettings, 0, len(SafetySettingList)), + } + + for _, category := range SafetySettingList { + geminiRequest.SafetySettings = append(geminiRequest.SafetySettings, dto.GeminiChatSafetySettings{ + Category: category, + Threshold: model_setting.GetGeminiSafetySetting(category), + }) + } + + imageConfig := map[string]interface{}{} + if aspectRatio := geminiAspectRatioFromSize(request.Size); aspectRatio != "" { + imageConfig["aspectRatio"] = aspectRatio + } + if imageSize := geminiImageSizeFromQuality(request.Quality); imageSize != "" { + imageConfig["imageSize"] = imageSize + } + if len(imageConfig) > 0 { + imageConfigBytes, err := common.Marshal(imageConfig) + if err != nil { + return nil, fmt.Errorf("failed to marshal image config: %w", err) + } + geminiRequest.GenerationConfig.ImageConfig = imageConfigBytes + } + + return &geminiRequest, nil +} + +func geminiAspectRatioFromSize(size string) string { + size = strings.TrimSpace(size) + if size == "" { + return "" + } + if strings.Contains(size, ":") { + return size + } + switch size { + case "256x256", "512x512", "1024x1024": + return "1:1" + case "1536x1024": + return "3:2" + case "1024x1536": + return "2:3" + case "1024x1792": + return "9:16" + case "1792x1024": + return "16:9" + default: + return "" + } +} + +func geminiImageSizeFromQuality(quality string) string { + switch strings.TrimSpace(quality) { + case "": + return "" + case "hd", "high", "2K": + return "2K" + case "standard", "medium", "low", "auto", "1K": + return "1K" + default: + return "1K" + } } func (a *Adaptor) Init(info *relaycommon.RelayInfo) { @@ -263,6 +393,10 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom return GeminiImageHandler(c, info, resp) } + if info.RelayMode == constant.RelayModeImagesGenerations { + return GeminiGenerateContentImageHandler(c, info, resp) + } + // check if the model is an embedding model if strings.HasPrefix(info.UpstreamModelName, "text-embedding") || strings.HasPrefix(info.UpstreamModelName, "embedding") || diff --git a/relay/channel/gemini/image_generation_test.go b/relay/channel/gemini/image_generation_test.go new file mode 100644 index 000000000000..3bfb8c156385 --- /dev/null +++ b/relay/channel/gemini/image_generation_test.go @@ -0,0 +1,82 @@ +package gemini + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestConvertOpenAIImageRequestToGeminiGenerateContent(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + RelayMode: relayconstant.RelayModeImagesGenerations, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gemini-2.5-flash-image", + }, + } + + var req dto.ImageRequest + require.NoError(t, common.Unmarshal([]byte(`{ + "model":"gemini-2.5-flash-image", + "prompt":"make it cinematic", + "image":"data:image/png;base64,aW1hZ2U=", + "size":"1792x1024", + "quality":"high" + }`), &req)) + + got, err := (&Adaptor{}).ConvertImageRequest(c, info, req) + require.NoError(t, err) + + geminiReq, ok := got.(*dto.GeminiChatRequest) + require.True(t, ok) + require.Len(t, geminiReq.Contents, 1) + require.Len(t, geminiReq.Contents[0].Parts, 2) + require.Equal(t, "make it cinematic", geminiReq.Contents[0].Parts[0].Text) + require.NotNil(t, geminiReq.Contents[0].Parts[1].InlineData) + require.Equal(t, "image/png", geminiReq.Contents[0].Parts[1].InlineData.MimeType) + require.Equal(t, "aW1hZ2U=", geminiReq.Contents[0].Parts[1].InlineData.Data) + require.Equal(t, []string{"TEXT", "IMAGE"}, geminiReq.GenerationConfig.ResponseModalities) + + var imageConfig map[string]string + require.NoError(t, common.Unmarshal(geminiReq.GenerationConfig.ImageConfig, &imageConfig)) + require.Equal(t, "16:9", imageConfig["aspectRatio"]) + require.Equal(t, "2K", imageConfig["imageSize"]) +} + +func TestGeminiGenerateContentImageHandler(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{ + "candidates":[{ + "content":{ + "parts":[ + {"text":"done"}, + {"inlineData":{"mimeType":"image/png","data":"Z2VuZXJhdGVk"}} + ] + } + }], + "usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3,"totalTokenCount":5} + }`)), + } + info := &relaycommon.RelayInfo{RelayMode: relayconstant.RelayModeImagesGenerations} + + usage, err := GeminiGenerateContentImageHandler(c, info, resp) + require.Nil(t, err) + require.Equal(t, 5, usage.TotalTokens) + require.Contains(t, recorder.Body.String(), `"b64_json":"Z2VuZXJhdGVk"`) + require.NotContains(t, recorder.Body.String(), `"choices"`) +} diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index e39826dd64e7..333ca39b1ee8 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1652,7 +1652,7 @@ func GeminiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. }) } - jsonResponse, jsonErr := json.Marshal(openAIResponse) + jsonResponse, jsonErr := common.Marshal(openAIResponse) if jsonErr != nil { return nil, types.NewError(jsonErr, types.ErrorCodeBadResponseBody) } @@ -1675,6 +1675,69 @@ func GeminiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. return usage, nil } +func GeminiGenerateContentImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { + responseBody, readErr := io.ReadAll(resp.Body) + if readErr != nil { + return nil, types.NewOpenAIError(readErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } + _ = resp.Body.Close() + + var geminiResponse dto.GeminiChatResponse + if jsonErr := common.Unmarshal(responseBody, &geminiResponse); jsonErr != nil { + return nil, types.NewOpenAIError(jsonErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } + if len(geminiResponse.Candidates) == 0 { + usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { + return &usage, types.NewOpenAIError( + errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason), + types.ErrorCodePromptBlocked, + http.StatusBadRequest, + ) + } + return &usage, types.NewOpenAIError(errors.New("empty response from Gemini API"), types.ErrorCodeEmptyResponse, http.StatusInternalServerError) + } + + openAIResponse := dto.ImageResponse{ + Created: common.GetTimestamp(), + Data: make([]dto.ImageData, 0), + } + for _, candidate := range geminiResponse.Candidates { + for _, part := range candidate.Content.Parts { + if part.InlineData == nil || !strings.HasPrefix(strings.ToLower(part.InlineData.MimeType), "image/") { + continue + } + openAIResponse.Data = append(openAIResponse.Data, dto.ImageData{ + B64Json: part.InlineData.Data, + }) + } + } + if len(openAIResponse.Data) == 0 { + return nil, types.NewOpenAIError(errors.New("no images generated"), types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } + + jsonResponse, jsonErr := common.Marshal(openAIResponse) + if jsonErr != nil { + return nil, types.NewError(jsonErr, types.ErrorCodeBadResponseBody) + } + + c.Writer.Header().Set("Content-Type", "application/json") + c.Writer.WriteHeader(resp.StatusCode) + _, _ = c.Writer.Write(jsonResponse) + + usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + if usage.TotalTokens == 0 { + const imageTokens = 258 + generatedImages := len(openAIResponse.Data) + usage = dto.Usage{ + PromptTokens: imageTokens * generatedImages, + CompletionTokens: 0, + TotalTokens: imageTokens * generatedImages, + } + } + return &usage, nil +} + type GeminiModelsResponse struct { Models []dto.GeminiModel `json:"models"` NextPageToken string `json:"nextPageToken"` diff --git a/relay/channel/vertex/adaptor.go b/relay/channel/vertex/adaptor.go index 7f087c21b90c..6f32da720a3b 100644 --- a/relay/channel/vertex/adaptor.go +++ b/relay/channel/vertex/adaptor.go @@ -354,6 +354,9 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom if strings.HasPrefix(info.UpstreamModelName, "imagen") { return gemini.GeminiImageHandler(c, info, resp) } + if info.RelayMode == constant.RelayModeImagesGenerations { + return gemini.GeminiGenerateContentImageHandler(c, info, resp) + } return gemini.GeminiChatHandler(c, info, resp) } case RequestModeOpenSource: diff --git a/setting/model_setting/gemini.go b/setting/model_setting/gemini.go index dea7131b9bbf..b56d46056e3b 100644 --- a/setting/model_setting/gemini.go +++ b/setting/model_setting/gemini.go @@ -30,6 +30,7 @@ var defaultGeminiSettings = GeminiSettings{ "gemini-3-pro-image-preview", "gemini-2.5-flash-image", "gemini-3.1-flash-image-preview", + "nano-banana-pro-preview", }, ThinkingAdapterEnabled: false, ThinkingAdapterBudgetTokensPercentage: 0.6,