chiyi/image2api
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(upstream): 自定义上游支持 Vivid 协议(/v1/generate NDJSON)

Browse Source 
- custom 上游新增 protocol 选项(openai|vivid);vivid 走
  POST /v1/generate NDJSON 流:跳过心跳空行,取末行终态,
  断流携 jobId 重连不重复扣费;参考图先 POST /v1/uploads
  拿预签名再 PUT 直传;错误码 4011/4002/4293/5001/4001 分类
- 添加上游弹窗加协议二选一;账号 meta 存 protocol 并可回显
- 附 vivid 链路单测(成功流/URL 透传/失败帧/错误码/两步上传)

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
chiyi
2026-10-14 12:18:26 +08:00
co-authored by Devin
parent ae49b95e0f
commit babfc71f13
6 changed files with 513 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/tokens.go
+26 -15
View File
@@ -2,24 +2,24 @@ package service
import (
"context"
"sync"
"sync/atomic"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"log"
"strings"
"sync"
"sync/atomic"
"time"
"backend/internal/model"
"backend/internal/provider/adobe"
"backend/internal/provider/chatgpt"
"backend/internal/provider/creativefabrica"
"backend/internal/provider/grok"
"backend/internal/provider/imagine"
"backend/internal/provider/krea"
"backend/internal/provider/leonardo"
"backend/internal/provider/grok"
"backend/internal/provider/runway"
"backend/internal/repo"
@@ -28,14 +28,14 @@ import (
)
var validTokenPools = map[string]string{
"chatgpt": "openai",
"adobe": "adobe",
"runway": "runway",
"leonardo": "leonardo",
"krea": "krea",
"imagine": "imagine",
"grok": "grok",
"custom": "custom",
"chatgpt": "openai",
"adobe": "adobe",
"runway": "runway",
"leonardo": "leonardo",
"krea": "krea",
"imagine": "imagine",
"grok": "grok",
"custom": "custom",
"creativefabrica": "creativefabrica",
}
@@ -1036,7 +1036,6 @@ func (s *TokenService) ImportGrokToken(ctx context.Context, ssoToken, tokenID st
return item, nil
}
func (s *TokenService) checkPendingGrok(tokenID, ssoToken string) {
defer func() {
if r := recover(); r != nil {
@@ -1217,9 +1216,10 @@ func (s *TokenService) checkPendingCreativeFabrica(tokenID, cookie string) {
// csv list of model ids it serves (empty = all), plus optional weight and
// per-account concurrency. No probe — the account goes active immediately and is
// matched to custom models by id at generation time. Calls go direct (no proxy).
func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey, models, name string, weight, concurrency int, tokenID string) (*model.TokenAccount, error) {
func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey, models, name, protocol string, weight, concurrency int, tokenID string) (*model.TokenAccount, error) {
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
apiKey = strings.TrimSpace(apiKey)
protocol = strings.ToLower(strings.TrimSpace(protocol))
// Edit mode: tokenID points at an existing custom account. base_url required;
// a blank key keeps the stored one.
if strings.TrimSpace(tokenID) != "" {
@@ -1230,7 +1230,14 @@ func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey,
if baseURL == "" {
return nil, errors.New("base_url required")
}
meta := datatypes.JSONMap{"base_url": baseURL}
proto := protocol
if proto == "" {
proto = strings.ToLower(strings.TrimSpace(stringValue(existing.Meta["protocol"])))
}
if proto != "vivid" {
proto = "openai"
}
meta := datatypes.JSONMap{"base_url": baseURL, "protocol": proto}
if m := strings.TrimSpace(models); m != "" {
meta["models"] = m
}
@@ -1248,7 +1255,10 @@ func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey,
if baseURL == "" || apiKey == "" {
return nil, errors.New("base_url and key required")
}
meta := datatypes.JSONMap{"base_url": baseURL}
if protocol != "vivid" {
protocol = "openai"
}
meta := datatypes.JSONMap{"base_url": baseURL, "protocol": protocol}
if m := strings.TrimSpace(models); m != "" {
meta["models"] = m
}
@@ -1952,6 +1962,7 @@ func accountRow(item model.TokenAccount, inFlight int64) map[string]any {
"weight": item.Weight,
"concurrency": item.Concurrency,
"base_url": emptyToNil(strings.TrimSpace(stringValue(item.Meta["base_url"]))),
"protocol": emptyToNil(strings.TrimSpace(stringValue(item.Meta["protocol"]))),
"models": strings.TrimSpace(stringValue(item.Meta["models"])),
}
}
backend/internal/service/v1.go
+30 -5
View File
@@ -23,8 +23,8 @@ import (
"backend/internal/model"
"backend/internal/provider/adobe"
"backend/internal/provider/chatgpt"
"backend/internal/provider/custom"
"backend/internal/provider/creativefabrica"
"backend/internal/provider/custom"
"backend/internal/provider/grok"
"backend/internal/provider/imagine"
"backend/internal/provider/krea"
@@ -32,6 +32,7 @@ import (
"backend/internal/provider/runway"
"backend/internal/repo"
"backend/internal/storage"
"gorm.io/gorm"
)
@@ -2228,7 +2229,14 @@ func (s *V1Service) generateCustomImage(ctx context.Context, eventID string, mod
_ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail)
_ = s.tokens.TouchLastUsed(ctx, token.ID)
baseURL := stringValue(token.Meta["base_url"])
d, u, genErr := s.custom.GenerateImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, quality, refs, !urlOnly)
var d []byte
var u string
var genErr error
if strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") {
d, u, genErr = s.custom.GenerateVividImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, aspectRatio, vividQuality(resolution), refs, !urlOnly)
} else {
d, u, genErr = s.custom.GenerateImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, quality, refs, !urlOnly)
}
if genErr == nil {
_, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{
"last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0,
@@ -2303,7 +2311,14 @@ func (s *V1Service) generateCustomVideo(ctx context.Context, eventID string, mod
_ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail)
_ = s.tokens.TouchLastUsed(ctx, token.ID)
baseURL := stringValue(token.Meta["base_url"])
d, url, genErr := s.custom.GenerateVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, durationSeconds, frames, downloadResult)
var d []byte
var url string
var genErr error
if strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") {
d, url, genErr = s.custom.GenerateVividVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, aspectRatio, vividQuality(resolution), durationSeconds, frames, downloadResult)
} else {
d, url, genErr = s.custom.GenerateVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, durationSeconds, frames, downloadResult)
}
if genErr == nil {
_, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{
"last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0,
@@ -2401,6 +2416,16 @@ func upstreamVideoSize(aspectRatio, resolution string) string {
return fmt.Sprintf("%dx%d", short, short*h/w)
}
// vividQuality maps a resolution tier to the quality string a Vivid-protocol
// upstream expects ("1K"/"2K"/"4K"); empty → "1K" (required field).
func vividQuality(resolution string) string {
switch strings.ToUpper(strings.TrimSpace(resolution)) {
case "2K", "4K":
return strings.ToUpper(strings.TrimSpace(resolution))
}
return "1K"
}
// upstreamQuality maps a resolution tier to the OpenAI quality enum.
func upstreamQuality(resolution string) string {
switch strings.ToUpper(strings.TrimSpace(resolution)) {
@@ -2975,8 +3000,8 @@ const leonardoPrivateSuffix = "-不卡人脸"
type leonardoVideoSpec struct {
upstream string
// long/short 是长边/短边像素,按比例组合成 16:9 或 9:16。
long int
short int
long int
short int
maxImages int
maxAudios int
maxVideos int