feature: 支持 replicate 的绘图

.github/workflows/ci.yml

@@ -1,13 +1,13 @@
 name: CI
 
 # This setup assumes that you run the unit tests with code coverage in the same
-# workflow that will also print the coverage report as comment to the pull request. 
+# workflow that will also print the coverage report as comment to the pull request.
 # Therefore, you need to trigger this workflow when a pull request is (re)opened or
 # when new code is pushed to the branch of the pull request. In addition, you also
-# need to trigger this workflow when new code is pushed to the main branch because 
+# need to trigger this workflow when new code is pushed to the main branch because
 # we need to upload the code coverage results as artifact for the main branch as
 # well since it will be the baseline code coverage.
-# 
+#
 # We do not want to trigger the workflow for pushes to *any* branch because this
 # would trigger our jobs twice on pull requests (once from "push" event and once
 # from "pull_request->synchronize")
@@ -31,7 +31,7 @@ jobs:
         with:
           go-version: ^1.22
 
-      # When you execute your unit tests, make sure to use the "-coverprofile" flag to write a 
+      # When you execute your unit tests, make sure to use the "-coverprofile" flag to write a
       # coverage profile to a file. You will need the name of the file (e.g. "coverage.txt")
       # in the next step as well as the next job.
       - name: Test

.gitignore

@@ -9,4 +9,5 @@ logs
 data
 /web/node_modules
 cmd.md
-.env
+.env
+/one-api

relay/adaptor.go

@@ -16,6 +16,7 @@ import (
 	"github.com/songquanpeng/one-api/relay/adaptor/openai"
 	"github.com/songquanpeng/one-api/relay/adaptor/palm"
 	"github.com/songquanpeng/one-api/relay/adaptor/proxy"
+	"github.com/songquanpeng/one-api/relay/adaptor/replicate"
 	"github.com/songquanpeng/one-api/relay/adaptor/tencent"
 	"github.com/songquanpeng/one-api/relay/adaptor/vertexai"
 	"github.com/songquanpeng/one-api/relay/adaptor/xunfei"
@@ -61,6 +62,8 @@ func GetAdaptor(apiType int) adaptor.Adaptor {
 		return &vertexai.Adaptor{}
 	case apitype.Proxy:
 		return &proxy.Adaptor{}
+	case apitype.Replicate:
+		return &replicate.Adaptor{}
 	}
 	return nil
 }

relay/adaptor/openai/util.go

@@ -1,8 +1,16 @@
 package openai
 
-import "github.com/songquanpeng/one-api/relay/model"
+import (
+	"context"
+	"fmt"
+
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay/model"
+)
 
 func ErrorWrapper(err error, code string, statusCode int) *model.ErrorWithStatusCode {
+	logger.Error(context.TODO(), fmt.Sprintf("[%s]%+v", code, err))
+
 	Error := model.Error{
 		Message: err.Error(),
 		Type:    "one_api_error",

relay/adaptor/replicate/adaptor.go

@@ -0,0 +1,85 @@
+package replicate
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"slices"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay/adaptor"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/meta"
+	"github.com/songquanpeng/one-api/relay/model"
+	"github.com/songquanpeng/one-api/relay/relaymode"
+)
+
+type Adaptor struct {
+	meta *meta.Meta
+}
+
+// ConvertImageRequest implements adaptor.Adaptor.
+func (*Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) {
+	return DrawImageRequest{
+		Input: ImageInput{
+			Steps:           25,
+			Prompt:          request.Prompt,
+			Guidance:        3,
+			Seed:            int(time.Now().UnixNano()),
+			SafetyTolerance: 5,
+			NImages:         1, // replicate will always return 1 image
+			Width:           1440,
+			Height:          1440,
+			AspectRatio:     "1:1",
+		},
+	}, nil
+}
+
+func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) {
+	return nil, errors.New("not implemented")
+}
+
+func (a *Adaptor) Init(meta *meta.Meta) {
+	a.meta = meta
+}
+
+func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
+	if !slices.Contains(ModelList, meta.OriginModelName) {
+		return "", errors.Errorf("model %s not supported", meta.OriginModelName)
+	}
+
+	return fmt.Sprintf("https://api.replicate.com/v1/models/%s/predictions", meta.OriginModelName), nil
+}
+
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error {
+	adaptor.SetupCommonRequestHeader(c, req, meta)
+	req.Header.Set("Authorization", "Bearer "+meta.APIKey)
+	return nil
+}
+
+func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) {
+	logger.Info(c, "send image request to replicate")
+	return adaptor.DoRequestHelper(a, c, meta, requestBody)
+}
+
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) {
+	switch meta.Mode {
+	case relaymode.ImagesGenerations:
+		err, usage = ImageHandler(c, resp)
+	default:
+		err = openai.ErrorWrapper(errors.New("not implemented"), "not_implemented", http.StatusInternalServerError)
+	}
+
+	return
+}
+
+func (a *Adaptor) GetModelList() []string {
+	return ModelList
+}
+
+func (a *Adaptor) GetChannelName() string {
+	return "replicate"
+}

relay/adaptor/replicate/constant.go

@@ -0,0 +1,58 @@
+package replicate
+
+// ModelList is a list of models that can be used with Replicate.
+//
+// https://replicate.com/pricing
+var ModelList = []string{
+	// -------------------------------------
+	// image model
+	// -------------------------------------
+	"black-forest-labs/flux-1.1-pro",
+	"black-forest-labs/flux-1.1-pro-ultra",
+	"black-forest-labs/flux-canny-dev",
+	"black-forest-labs/flux-canny-pro",
+	"black-forest-labs/flux-depth-dev",
+	"black-forest-labs/flux-depth-pro",
+	"black-forest-labs/flux-dev",
+	"black-forest-labs/flux-dev-lora",
+	"black-forest-labs/flux-fill-dev",
+	"black-forest-labs/flux-fill-pro",
+	"black-forest-labs/flux-pro",
+	"black-forest-labs/flux-redux-dev",
+	"black-forest-labs/flux-redux-schnell",
+	"black-forest-labs/flux-schnell",
+	"black-forest-labs/flux-schnell-lora",
+	"ideogram-ai/ideogram-v2",
+	"ideogram-ai/ideogram-v2-turbo",
+	"recraft-ai/recraft-v3",
+	"recraft-ai/recraft-v3-svg",
+	"stability-ai/stable-diffusion-3",
+	"stability-ai/stable-diffusion-3.5-large",
+	"stability-ai/stable-diffusion-3.5-large-turbo",
+	"stability-ai/stable-diffusion-3.5-medium",
+	// -------------------------------------
+	// language model
+	// -------------------------------------
+	// "ibm-granite/granite-20b-code-instruct-8k",  // TODO: implement the adaptor
+	// "ibm-granite/granite-3.0-2b-instruct",  // TODO: implement the adaptor
+	// "ibm-granite/granite-3.0-8b-instruct",  // TODO: implement the adaptor
+	// "ibm-granite/granite-8b-code-instruct-128k",  // TODO: implement the adaptor
+	// "meta/llama-2-13b",  // TODO: implement the adaptor
+	// "meta/llama-2-13b-chat",  // TODO: implement the adaptor
+	// "meta/llama-2-70b",  // TODO: implement the adaptor
+	// "meta/llama-2-70b-chat",  // TODO: implement the adaptor
+	// "meta/llama-2-7b",  // TODO: implement the adaptor
+	// "meta/llama-2-7b-chat",  // TODO: implement the adaptor
+	// "meta/meta-llama-3.1-405b-instruct",  // TODO: implement the adaptor
+	// "meta/meta-llama-3-70b",  // TODO: implement the adaptor
+	// "meta/meta-llama-3-70b-instruct",  // TODO: implement the adaptor
+	// "meta/meta-llama-3-8b",  // TODO: implement the adaptor
+	// "meta/meta-llama-3-8b-instruct",  // TODO: implement the adaptor
+	// "mistralai/mistral-7b-instruct-v0.2",  // TODO: implement the adaptor
+	// "mistralai/mistral-7b-v0.1",  // TODO: implement the adaptor
+	// "mistralai/mixtral-8x7b-instruct-v0.1",  // TODO: implement the adaptor
+	// -------------------------------------
+	// video model
+	// -------------------------------------
+	// "minimax/video-01",  // TODO: implement the adaptor
+}