30 lines
1.7 KiB
Go
30 lines
1.7 KiB
Go
package helper
|
|
|
|
import (
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/setting/model_setting"
|
|
"github.com/stretchr/testify/require"
|
|
"testing"
|
|
)
|
|
|
|
func TestImageGenerationResolutionMapping(t *testing.T) {
|
|
cfg := model_setting.GetGlobalSettings()
|
|
oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
|
|
t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
|
|
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2", "gpt-image-2-4k"}
|
|
for _, tc := range []struct{ size, tier string }{{"", "1k"}, {"auto", "1k"}, {"1024x1536", "1k"}, {"1536x1024", "1k"}, {"2048x1152", "2k"}, {"2048x2048", "2k"}, {"3840x2160", "4k"}, {"2160x3840", "4k"}, {"4k", "4k"}} {
|
|
info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "example-images", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
|
|
image := &dto.ImageRequest{Model: "gpt-image-2", Size: tc.size}
|
|
require.NoError(t, mapImageGenerationResolution(info, image))
|
|
require.Equal(t, "gpt-image-2-"+tc.tier, info.UpstreamModelName)
|
|
}
|
|
_, err := imageGenerationResolution("-1x1024")
|
|
require.Error(t, err)
|
|
_, err = imageGenerationResolution("garbage")
|
|
require.Error(t, err)
|
|
info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "default", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
|
|
require.NoError(t, mapImageGenerationResolution(info, &dto.ImageRequest{Model: "gpt-image-2", Size: "3840x2160"}))
|
|
require.Equal(t, "gpt-image-2", info.UpstreamModelName)
|
|
}
|