Merge pull request #2190 from Sh1n3zZ/support-replicate-channel

feat: replicate channel flux model

common/api_type.go

@@ -71,6 +71,8 @@ func ChannelType2APIType(channelType int) (int, bool) {
 		apiType = constant.APITypeSubmodel
 	case constant.ChannelTypeMiniMax:
 		apiType = constant.APITypeMiniMax
+	case constant.ChannelTypeReplicate:
+		apiType = constant.APITypeReplicate
 	}
 	if apiType == -1 {
 		return constant.APITypeOpenAI, false

constant/api_type.go

@@ -34,5 +34,6 @@ const (
 	APITypeMoonshot
 	APITypeSubmodel
 	APITypeMiniMax
+	APITypeReplicate
 	APITypeDummy // this one is only for count, do not add any channel after this
 )

constant/channel.go

@@ -53,6 +53,7 @@ const (
 	ChannelTypeSubmodel       = 53
 	ChannelTypeDoubaoVideo    = 54
 	ChannelTypeSora           = 55
+	ChannelTypeReplicate      = 56
 	ChannelTypeDummy          // this one is only for count, do not add any channel after this
 
 )
@@ -114,6 +115,7 @@ var ChannelBaseURLs = []string{
 	"https://llm.submodel.ai",                   //53
 	"https://ark.cn-beijing.volces.com",         //54
 	"https://api.openai.com",                    //55
+	"https://api.replicate.com",                 //56
 }
 
 var ChannelTypeNames = map[int]string{
@@ -169,6 +171,7 @@ var ChannelTypeNames = map[int]string{
 	ChannelTypeSubmodel:       "Submodel",
 	ChannelTypeDoubaoVideo:    "DoubaoVideo",
 	ChannelTypeSora:           "Sora",
+	ChannelTypeReplicate:      "Replicate",
 }
 
 func GetChannelTypeName(channelType int) string {

relay/channel/replicate/constants.go

+package replicate
+
+const (
+	// ChannelName identifies the replicate channel.
+	ChannelName = "replicate"
+	// ModelFlux11Pro is the default image generation model supported by this channel.
+	ModelFlux11Pro = "black-forest-labs/flux-1.1-pro"
+)
+
+var ModelList = []string{
+	ModelFlux11Pro,
+}

relay/channel/replicate/dto.go

+package replicate
+
+type PredictionResponse struct {
+	Status string           `json:"status"`
+	Output any              `json:"output"`
+	Error  *PredictionError `json:"error"`
+}
+
+type PredictionError struct {
+	Code    string `json:"code"`
+	Message string `json:"message"`
+	Detail  string `json:"detail"`
+}
+
+type FileUploadResponse struct {
+	Urls struct {
+		Get string `json:"get"`
+	} `json:"urls"`
+}

relay/image_handler.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
 	"github.com/QuantumNous/new-api/dto"
 	"github.com/QuantumNous/new-api/logger"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
@@ -92,12 +93,17 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
 		httpResp = resp.(*http.Response)
 		info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
 		if httpResp.StatusCode != http.StatusOK {
+			if httpResp.StatusCode == http.StatusCreated && info.ApiType == constant.APITypeReplicate {
+				// replicate channel returns 201 Created when using Prefer: wait, treat it as success.
+				httpResp.StatusCode = http.StatusOK
+			} else {
 				newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
 				// reset status code 重置状态码
 				service.ResetStatusCode(newAPIError, statusCodeMappingStr)
 				return newAPIError
 			}
 		}
+	}
 
 	usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
 	if newAPIError != nil {

relay/relay_adaptor.go

@@ -26,6 +26,7 @@ import (
 	"github.com/QuantumNous/new-api/relay/channel/openai"
 	"github.com/QuantumNous/new-api/relay/channel/palm"
 	"github.com/QuantumNous/new-api/relay/channel/perplexity"
+	"github.com/QuantumNous/new-api/relay/channel/replicate"
 	"github.com/QuantumNous/new-api/relay/channel/siliconflow"
 	"github.com/QuantumNous/new-api/relay/channel/submodel"
 	taskali "github.com/QuantumNous/new-api/relay/channel/task/ali"
@@ -113,6 +114,8 @@ func GetAdaptor(apiType int) channel.Adaptor {
 		return &submodel.Adaptor{}
 	case constant.APITypeMiniMax:
 		return &minimax.Adaptor{}
+	case constant.APITypeReplicate:
+		return &replicate.Adaptor{}
 	}
 	return nil
 }

setting/ratio_setting/model_ratio.go

@@ -272,6 +272,7 @@ var defaultModelPrice = map[string]float64{
 	"suno_lyrics":                    0.01,
 	"dall-e-3":                       0.04,
 	"imagen-3.0-generate-002":        0.03,
+	"black-forest-labs/flux-1.1-pro": 0.04,
 	"gpt-4-gizmo-*":                  0.1,
 	"mj_video":                       0.8,
 	"mj_imagine":                     0.1,

web/src/constants/channel.constants.js

@@ -179,6 +179,11 @@ export const CHANNEL_OPTIONS = [
     color: 'green',
     label: 'Sora',
   },
+  {
+    value: 56,
+    color: 'blue',
+    label: 'Replicate',
+  },
 ];
 
 export const MODEL_TABLE_PAGE_SIZE = 10;

web/src/helpers/render.jsx

@@ -55,6 +55,7 @@ import {
   Kling,
   Jimeng,
   Perplexity,
+  Replicate,
 } from '@lobehub/icons';
 
 import {
@@ -342,6 +343,8 @@ export function getChannelIcon(channelType) {
       return <Jimeng.Color size={iconSize} />;
     case 54: // 豆包视频 Doubao Video
       return <Doubao.Color size={iconSize} />;
+    case 56: // Replicate
+      return <Replicate size={iconSize} />;
     case 8: // 自定义渠道
     case 22: // 知识库：FastGPT
       return <FastGPT.Color size={iconSize} />;