Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: add prefixes to image models to solve the problem of duplicate models (#1469)

Browse Source 
* Add prefixes to image models to solve the problem of duplicate models

* Fix the issue that response_format is not set, causing the b64_json parameter to be ignored.
This commit is contained in:
Mo 2024-05-28 01:32:57 +08:00 committed by GitHub
parent 3be28da57b
commit 332c8db0b3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
relay/billing/ratio/image.go
View File

@ -49,3 +49,8 @@ var ImagePromptLengthLimitations = map[string]int{
"wanx-v1": 4000,
"cogview-3": 833,
}
var ImageOriginModelName = map[string]string{
"ali-stable-diffusion-xl": "stable-diffusion-xl",
"ali-stable-diffusion-v1.5": "stable-diffusion-v1.5",
}

7
relay/controller/image.go
View File

@ -55,6 +55,11 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
return openai.ErrorWrapper(err, "get_image_cost_ratio_failed", http.StatusInternalServerError)
}
imageModel := imageRequest.Model
// Convert the original image model
imageRequest.Model, _ = getMappedModelName(imageRequest.Model, billingratio.ImageOriginModelName)
c.Set("response_format", imageRequest.ResponseFormat)
var requestBody io.Reader
if isModelMapped || meta.ChannelType == channeltype.Azure { // make Azure channel request body
jsonStr, err := json.Marshal(imageRequest)
@ -89,7 +94,7 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
requestBody = bytes.NewBuffer(jsonStr)
}
modelRatio := billingratio.GetModelRatio(imageRequest.Model)
modelRatio := billingratio.GetModelRatio(imageModel)
groupRatio := billingratio.GetGroupRatio(meta.Group)
ratio := modelRatio * groupRatio
userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId)