diff --git a/controller/relay-image.go b/controller/relay-image.go index 0ff18309..1f5138c0 100644 --- a/controller/relay-image.go +++ b/controller/relay-image.go @@ -10,6 +10,7 @@ import ( "net/http" "one-api/common" "one-api/model" + "strings" "github.com/gin-gonic/gin" ) @@ -32,6 +33,7 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode tokenId := c.GetInt("token_id") channelType := c.GetInt("channel") channelId := c.GetInt("channel_id") + apiVersion := c.GetString("api_version") userId := c.GetInt("id") group := c.GetString("group") @@ -101,8 +103,15 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode baseURL = c.GetString("base_url") } fullRequestURL := getFullRequestURL(baseURL, requestURL, channelType) + + // make Azure channel requestURL + if channelType == common.ChannelTypeAzure { + // url https://{resource_name}.openai.azure.com/openai/deployments/dall-e-3/images/generations?api-version=2023-06-01-preview + fullRequestURL = fmt.Sprintf("%s/%s%s?api-version=%s", fullRequestURL, imageModel, "/images/generations", apiVersion) + } + var requestBody io.Reader - if isModelMapped { + if isModelMapped || channelType == common.ChannelTypeAzure { // make Azure channel request body jsonStr, err := json.Marshal(imageRequest) if err != nil { return errorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) @@ -127,7 +136,13 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode if err != nil { return errorWrapper(err, "new_request_failed", http.StatusInternalServerError) } - req.Header.Set("Authorization", c.Request.Header.Get("Authorization")) + token := c.Request.Header.Get("Authorization") + if channelType == common.ChannelTypeAzure { // Azure authentication + token = strings.TrimPrefix(token, "Bearer ") + req.Header.Set("api-key", token) + } else { + req.Header.Set("Authorization", token) + } req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Header.Set("Accept", c.Request.Header.Get("Accept")) @@ -196,4 +211,4 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError) } return nil -} +} \ No newline at end of file diff --git a/controller/relay-utils.go b/controller/relay-utils.go index 391f28b4..84e1fc43 100644 --- a/controller/relay-utils.go +++ b/controller/relay-utils.go @@ -192,6 +192,11 @@ func getFullRequestURL(baseURL string, requestURL string, channelType int) strin } } + // Azure dall-e model compatibility + if channelType == common.ChannelTypeAzure && requestURL == "/v1/images/generations" { + fullRequestURL = fmt.Sprintf("%s%s", baseURL, "/openai/deployments") + } + return fullRequestURL } @@ -215,4 +220,4 @@ func postConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuo if totalQuota <= 0 { common.LogError(ctx, fmt.Sprintf("totalQuota consumed is %d, something is wrong", totalQuota)) } -} +} \ No newline at end of file diff --git a/controller/relay.go b/controller/relay.go index f91ba6da..58ee8381 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -133,12 +133,12 @@ type TextRequest struct { type ImageRequest struct { Model string `json:"model"` Prompt string `json:"prompt" binding:"required"` - N int `json:"n"` - Size string `json:"size"` - Quality string `json:"quality"` - ResponseFormat string `json:"response_format"` - Style string `json:"style"` - User string `json:"user"` + N int `json:"n,omitempty"` + Size string `json:"size,omitempty"` + Quality string `json:"quality,omitempty"` + ResponseFormat string `json:"response_format,omitempty"` + Style string `json:"style,omitempty"` + User string `json:"user,omitempty"` } type WhisperResponse struct {