60 lines
1.6 KiB
Go
60 lines
1.6 KiB
Go
package helper
|
|
|
|
import (
|
|
"fmt"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/setting/model_setting"
|
|
"slices"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
func imageGenerationResolution(size string) (string, error) {
|
|
size = strings.ToLower(strings.TrimSpace(size))
|
|
if size == "" || size == "auto" || size == "1k" {
|
|
return "1k", nil
|
|
}
|
|
if size == "2k" || size == "4k" {
|
|
return size, nil
|
|
}
|
|
dimensions := strings.Split(size, "x")
|
|
if len(dimensions) != 2 {
|
|
return "", fmt.Errorf("size must be auto, 1k, 2k, 4k or WIDTHxHEIGHT")
|
|
}
|
|
width, errWidth := strconv.Atoi(dimensions[0])
|
|
height, errHeight := strconv.Atoi(dimensions[1])
|
|
if errWidth != nil || errHeight != nil || width <= 0 || height <= 0 {
|
|
return "", fmt.Errorf("invalid image size")
|
|
}
|
|
short := min(width, height)
|
|
if short <= 1024 {
|
|
return "1k", nil
|
|
}
|
|
if short <= 2048 {
|
|
return "2k", nil
|
|
}
|
|
return "4k", nil
|
|
}
|
|
|
|
func mapImageGenerationResolution(info *relaycommon.RelayInfo, request dto.Request) error {
|
|
cfg := model_setting.GetGlobalSettings()
|
|
image, ok := request.(*dto.ImageRequest)
|
|
if !ok || cfg.ImageGenerationGroup == "" || info.UsingGroup != cfg.ImageGenerationGroup || !slices.Contains(cfg.ImageGenerationModels, info.OriginModelName) {
|
|
return nil
|
|
}
|
|
name := info.OriginModelName
|
|
for _, suffix := range []string{"-1k", "-2k", "-4k"} {
|
|
if strings.HasSuffix(name, suffix) {
|
|
return nil
|
|
}
|
|
}
|
|
tier, err := imageGenerationResolution(image.Size)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
info.UpstreamModelName = name + "-" + tier
|
|
info.IsModelMapped = true
|
|
return nil
|
|
}
|