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

Add prefixes to image models to solve the problem of duplicate models

Browse Source
This commit is contained in:
mo 2024-05-26 00:18:18 +08:00
parent 0acee9a065
commit c81d7204f4
2 changed files with 10 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",
}

6
relay/controller/image.go
View File

@ -55,6 +55,10 @@ 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)
var requestBody io.Reader
if isModelMapped || meta.ChannelType == channeltype.Azure { // make Azure channel request body
jsonStr, err := json.Marshal(imageRequest)
@ -89,7 +93,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)