feat: sync deployed v1.0.0-rc.38-custom.67

relay/helper/image_generation_resolution.go

@@ -0,0 +1,59 @@
+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
+}