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

- 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:
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
@@ -6,6 +6,7 @@ import (
"strings" "strings"
"backend/internal/service" "backend/internal/service"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"gorm.io/gorm" "gorm.io/gorm"
) )
@@ -141,6 +142,7 @@ func (h *ProviderAdminHandler) ImportCustomAccount(c *gin.Context) {
APIKey string `json:"api_key"` APIKey string `json:"api_key"`
Models string `json:"models"` Models string `json:"models"`
Name string `json:"name"` Name string `json:"name"`
Protocol string `json:"protocol"` // openai (default) | vivid
Weight int `json:"weight"` Weight int `json:"weight"`
Concurrency int `json:"concurrency"` Concurrency int `json:"concurrency"`
ID string `json:"id"` ID string `json:"id"`
@@ -157,7 +159,7 @@ func (h *ProviderAdminHandler) ImportCustomAccount(c *gin.Context) {
if key == "" { if key == "" {
key = body.APIKey key = body.APIKey
} }
item, err := h.tokens.ImportCustomAccount(c.Request.Context(), baseURL, key, body.Models, body.Name, body.Weight, body.Concurrency, body.ID) item, err := h.tokens.ImportCustomAccount(c.Request.Context(), baseURL, key, body.Models, body.Name, body.Protocol, body.Weight, body.Concurrency, body.ID)
if err != nil { if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
return return
@@ -395,13 +397,21 @@ func (h *ProviderAdminHandler) AccountsList(c *gin.Context) {
isFree := plan == "" || plan == "free" isFree := plan == "" || plan == "free"
switch planFilter { switch planFilter {
case "vip": case "vip":
if isFree { continue } if isFree {
continue
}
case "free": case "free":
if !isFree { continue } if !isFree {
continue
}
case "sub": case "sub":
if !isSub { continue } if !isSub {
continue
}
case "master": case "master":
if isFree || isSub { continue } if isFree || isSub {
continue
}
} }
} }
if q != "" { if q != "" {
+248
View File
@@ -6,6 +6,7 @@
package custom package custom
import ( import (
"bufio"
"bytes" "bytes"
"context" "context"
"encoding/json" "encoding/json"
@@ -120,6 +121,253 @@ func (c *Client) GenerateImage(ctx context.Context, baseURL, apiKey, model, prom
return imageFromResponse(ctx, body, downloadResult) return imageFromResponse(ctx, body, downloadResult)
} }
// ── Vivid 协议 ──────────────────────────────────────────────────────────────
// Vivid 上游用自家 API:POST /v1/generate 返回 application/x-ndjson 流 —— 首行
// {"jobId":...,"status":"running"},等待期间每 25s 一个空行心跳,末行为终态
// (succeeded 带 result 预签名 URL / failed 带 error)。参考素材走两步:
// POST /v1/uploads {kind,mime} → PUT 直传预签名 URL → 拿 key 放进 refs。
// 断流可携 jobId 重发重连。错误体为 {"error":{"code":"4011|4001|4002|4293|5001"}}。
// GenerateVividImage submits an image job to a Vivid-protocol upstream.
// quality is the tier string the upstream expects ("1K"/"2K"/"4K"), ratio the
// "W:H" aspect. refs are pre-uploaded via /v1/uploads into key strings.
func (c *Client) GenerateVividImage(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, refs [][]byte, downloadResult bool) ([]byte, string, error) {
payload, err := c.vividPayload(ctx, baseURL, apiKey, model, prompt, ratio, quality, 0, refs)
if err != nil {
return nil, "", err
}
return c.vividGenerate(ctx, baseURL, apiKey, payload, downloadResult)
}
// GenerateVividVideo submits a video job (duration required by ranged models)
// to a Vivid-protocol upstream.
func (c *Client) GenerateVividVideo(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, durationSeconds int, refs [][]byte, downloadResult bool) ([]byte, string, error) {
payload, err := c.vividPayload(ctx, baseURL, apiKey, model, prompt, ratio, quality, durationSeconds, refs)
if err != nil {
return nil, "", err
}
return c.vividGenerate(ctx, baseURL, apiKey, payload, downloadResult)
}
func (c *Client) vividPayload(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, durationSeconds int, refs [][]byte) (map[string]any, error) {
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
if baseURL == "" || apiKey == "" {
return nil, ErrAuth
}
payload := map[string]any{"model": model, "prompt": prompt}
if quality != "" {
payload["quality"] = quality
}
if ratio != "" {
payload["ratio"] = ratio
}
if durationSeconds > 0 {
payload["duration"] = durationSeconds
}
if len(refs) > 0 {
keys := make([]string, 0, len(refs))
for _, ref := range refs {
key, err := c.vividUploadRef(ctx, baseURL, apiKey, ref)
if err != nil {
return nil, err
}
keys = append(keys, key)
}
payload["refs"] = keys
}
return payload, nil
}
// vividUploadRef runs the two-step upload: POST /v1/uploads {kind:"image",mime}
// → PUT the bytes to the presigned url → returns the key for refs[].
func (c *Client) vividUploadRef(ctx context.Context, baseURL, apiKey string, data []byte) (string, error) {
mime := http.DetectContentType(data)
if mime == "" || mime == "application/octet-stream" {
mime = "image/png"
}
raw, _ := json.Marshal(map[string]any{"kind": "image", "mime": mime})
up, err := c.doJSON(ctx, http.MethodPost, baseURL+"/v1/uploads", apiKey, raw)
if err != nil {
return "", err
}
putURL := strings.TrimSpace(stringValue(up["url"]))
key := strings.TrimSpace(stringValue(up["key"]))
if putURL == "" || key == "" {
return "", fmt.Errorf("%w: upload session missing url/key", ErrTemporaryUpstream)
}
req, err := http.NewRequest(http.MethodPut, putURL, bytes.NewReader(data))
if err != nil {
return "", err
}
req = req.WithContext(ctx)
req.Header.Set("Content-Type", mime) // must match the signed mime exactly
resp, err := httpClient().Do(req)
if err != nil {
return "", fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return "", fmt.Errorf("%w: ref upload %d %s", ErrTemporaryUpstream, resp.StatusCode, clip(body, 120))
}
return key, nil
}
// vividGenerate POSTs /v1/generate and reads the NDJSON stream to the terminal
// frame. If the stream drops before terminal, it reconnects once with the
// jobId (same-request resume — no new task, no double charge).
func (c *Client) vividGenerate(ctx context.Context, baseURL, apiKey string, payload map[string]any, downloadResult bool) ([]byte, string, error) {
var last map[string]any
for attempt := 0; attempt < 3; attempt++ {
if attempt > 0 {
if sleepCtx(ctx, 2*time.Second) != nil {
return nil, "", ctx.Err()
}
}
term, err := c.vividStreamOnce(ctx, baseURL, apiKey, payload)
if err == nil && term != nil {
last = term
break
}
if err != nil && !errors.Is(err, ErrTemporaryUpstream) {
return nil, "", err
}
// Temporary error mid-stream: if we hold a jobId, reconnect mode retries
// without creating a new task.
if job := strings.TrimSpace(stringValue(payload["jobId"])); job != "" {
payload = map[string]any{"jobId": job}
continue
}
if attempt == 2 {
return nil, "", err
}
}
if last == nil {
return nil, "", fmt.Errorf("%w: no terminal frame", ErrTemporaryUpstream)
}
status := strings.ToLower(strings.TrimSpace(stringValue(last["status"])))
switch status {
case "succeeded", "completed", "success":
resultURL := strings.TrimSpace(stringValue(last["result"]))
if resultURL == "" {
return nil, "", fmt.Errorf("%w: succeeded frame missing result", ErrTemporaryUpstream)
}
if !downloadResult {
return nil, resultURL, nil
}
data, err := c.vividDownload(ctx, resultURL)
if err != nil {
return nil, "", err
}
return data, resultURL, nil
default:
reason := strings.TrimSpace(stringValue(last["error"]))
if reason == "" {
reason = "upstream job " + status
}
if isCreditError(reason) || strings.Contains(reason, "积分不足") || strings.Contains(reason, "余额") {
return nil, "", fmt.Errorf("%w: %s", ErrQuotaExhausted, clip([]byte(reason), 160))
}
return nil, "", errors.New("custom: " + clip([]byte(reason), 160))
}
}
// vividStreamOnce performs one /v1/generate request and scans the NDJSON body
// to the last frame. Returns the terminal frame, or a temporary error on
// transport failure (caller's reconnect logic decides what to do with jobId).
func (c *Client) vividStreamOnce(ctx context.Context, baseURL, apiKey string, payload map[string]any) (map[string]any, error) {
raw, _ := json.Marshal(payload)
req, err := http.NewRequest(http.MethodPost, baseURL+"/v1/generate", bytes.NewReader(raw))
if err != nil {
return nil, err
}
req = req.WithContext(ctx)
req.Header.Set("Authorization", "Bearer "+apiKey)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Wait-Seconds", "7200")
resp, err := httpClient().Do(req)
if err != nil {
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
body, _ := io.ReadAll(resp.Body)
return nil, vividError(resp.StatusCode, body)
}
var last map[string]any
sc := bufio.NewScanner(resp.Body)
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
for sc.Scan() {
line := bytes.TrimSpace(sc.Bytes())
if len(line) == 0 {
continue // 25s heartbeat empty line
}
var frame map[string]any
if json.Unmarshal(line, &frame) != nil {
continue
}
last = frame
// Remember the jobId for potential reconnect.
if job := strings.TrimSpace(stringValue(frame["jobId"])); job != "" {
payload["jobId"] = job
}
}
if err := sc.Err(); err != nil {
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
}
return last, nil
}
// vividError maps the {"error":{"code":"NNNN"}} envelope to our error classes.
func vividError(status int, body []byte) error {
var e struct {
Error struct {
Code string `json:"code"`
Message string `json:"message"`
} `json:"error"`
}
code := ""
msg := clip(body, 160)
if json.Unmarshal(body, &e) == nil {
code = strings.TrimSpace(e.Error.Code)
if m := strings.TrimSpace(e.Error.Message); m != "" {
msg = clip([]byte(m), 160)
}
}
switch code {
case "4011":
return fmt.Errorf("%w: %s", ErrAuth, msg)
case "4002":
return fmt.Errorf("%w: %s", ErrQuotaExhausted, msg)
case "4293", "5001":
return fmt.Errorf("%w: %s", ErrTemporaryUpstream, msg)
case "4001":
return errors.New("custom: " + msg) // 参数/业务拒绝 → 永久错,不换号
}
return mapStatus(status, body)
}
// vividDownload fetches the presigned result URL — no auth header (it's signed).
func (c *Client) vividDownload(ctx context.Context, url string) ([]byte, error) {
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
resp, err := httpClient().Do(req)
if err != nil {
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("%w: result download %d", ErrTemporaryUpstream, resp.StatusCode)
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if len(data) == 0 {
return nil, fmt.Errorf("%w: empty result", ErrTemporaryUpstream)
}
return data, nil
}
// GenerateVideo drives the upstream Sora-style async video API: // GenerateVideo drives the upstream Sora-style async video API:
// POST /v1/videos → poll GET /v1/videos/{id} → GET /v1/videos/{id}/content. // POST /v1/videos → poll GET /v1/videos/{id} → GET /v1/videos/{id}/content.
// Reference frames (image-to-video / first-last frames) are sent as multipart // Reference frames (image-to-video / first-last frames) are sent as multipart
@@ -0,0 +1,179 @@
package custom
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
// vividServer serves /v1/uploads, /v1/generate (ndjson) and a fake presigned
// result URL, all under one httptest server.
func vividServer(t *testing.T, frames []string, httpStatus int, errCode string) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/v1/uploads":
var req struct {
Kind string `json:"kind"`
Mime string `json:"mime"`
}
_ = json.NewDecoder(r.Body).Decode(&req)
if req.Kind != "image" || req.Mime == "" {
t.Errorf("bad upload request: %+v", req)
}
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"uuid":"u1","key":"u/refs/image/u1.png","url":"%s","expiresAt":9999999999}`,
// presigned URL points back at this server's /put endpoint
"http://"+r.Host+"/put")
case "/put":
if r.Method != http.MethodPut {
t.Errorf("expected PUT, got %s", r.Method)
}
if ct := r.Header.Get("Content-Type"); !strings.HasPrefix(ct, "image/") {
t.Errorf("bad content-type %q", ct)
}
io.Copy(io.Discard, r.Body)
w.WriteHeader(http.StatusOK)
case "/v1/generate":
if httpStatus != 0 {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(httpStatus)
fmt.Fprintf(w, `{"error":{"code":%q,"message":"boom"}}`, errCode)
return
}
w.Header().Set("Content-Type", "application/x-ndjson")
for _, f := range frames {
fmt.Fprintln(w, f)
fmt.Fprintln(w) // heartbeat line
}
case "/result.png":
w.Write([]byte("PNGDATA"))
default:
http.NotFound(w, r)
}
}))
}
func TestVividGenerateImageSuccess(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/v1/generate":
w.Header().Set("Content-Type", "application/x-ndjson")
fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`)
fmt.Fprintln(w) // heartbeat
fmt.Fprintf(w, `{"jobId":"j1","status":"succeeded","result":"http://%s/result.png","creditsCharged":10}`+"\n", r.Host)
case "/result.png":
w.Write([]byte("PNGDATA"))
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
data, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, true)
if err != nil {
t.Fatalf("generate: %v", err)
}
if string(data) != "PNGDATA" {
t.Fatalf("data=%q", data)
}
if url == "" {
t.Fatal("empty url")
}
}
func TestVividGenerateImageURLPassthrough(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/x-ndjson")
fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`)
fmt.Fprintln(w, `{"jobId":"j1","status":"succeeded","result":"https://fs.example.com/x.png"}`)
}))
defer srv.Close()
data, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false)
if err != nil {
t.Fatalf("generate: %v", err)
}
if data != nil {
t.Fatalf("expected no bytes, got %d", len(data))
}
if url != "https://fs.example.com/x.png" {
t.Fatalf("url=%q", url)
}
}
func TestVividGenerateFailedFrame(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/x-ndjson")
fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`)
fmt.Fprintln(w, `{"jobId":"j1","status":"failed","error":"积分不足","refunded":true}`)
}))
defer srv.Close()
_, _, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false)
if !errors.Is(err, ErrQuotaExhausted) {
t.Fatalf("want quota, got %v", err)
}
}
func TestVividErrorCodes(t *testing.T) {
cases := []struct {
status int
code string
want error
}{
{401, "4011", ErrAuth},
{402, "4002", ErrQuotaExhausted},
{429, "4293", ErrTemporaryUpstream},
{500, "5001", ErrTemporaryUpstream},
}
for _, tc := range cases {
srv := vividServer(t, nil, tc.status, tc.code)
_, _, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false)
srv.Close()
if !errors.Is(err, tc.want) {
t.Fatalf("code %s: want %v, got %v", tc.code, tc.want, err)
}
}
}
func TestVividUploadRefFlow(t *testing.T) {
var putOK bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/v1/uploads":
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"key":"k123","url":"http://%s/up"}`, r.Host)
case "/up":
if r.Method == http.MethodPut {
putOK = true
}
w.WriteHeader(http.StatusOK)
case "/v1/generate":
var req map[string]any
_ = json.NewDecoder(r.Body).Decode(&req)
refs, _ := req["refs"].([]any)
if len(refs) != 1 || refs[0] != "k123" {
t.Errorf("refs=%v", req["refs"])
}
w.Header().Set("Content-Type", "application/x-ndjson")
fmt.Fprintln(w, `{"jobId":"j1","status":"succeeded","result":"https://x/1.png"}`)
}
}))
defer srv.Close()
png := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A} // fake png magic
_, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", [][]byte{png}, false)
if err != nil {
t.Fatalf("generate: %v", err)
}
if !putOK {
t.Fatal("PUT never happened")
}
if url == "" {
t.Fatal("no url")
}
}
+26 -15
View File
@@ -2,24 +2,24 @@ package service
import ( import (
"context" "context"
"sync"
"sync/atomic"
"encoding/base64" "encoding/base64"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"log" "log"
"strings" "strings"
"sync"
"sync/atomic"
"time" "time"
"backend/internal/model" "backend/internal/model"
"backend/internal/provider/adobe" "backend/internal/provider/adobe"
"backend/internal/provider/chatgpt" "backend/internal/provider/chatgpt"
"backend/internal/provider/creativefabrica" "backend/internal/provider/creativefabrica"
"backend/internal/provider/grok"
"backend/internal/provider/imagine" "backend/internal/provider/imagine"
"backend/internal/provider/krea" "backend/internal/provider/krea"
"backend/internal/provider/leonardo" "backend/internal/provider/leonardo"
"backend/internal/provider/grok"
"backend/internal/provider/runway" "backend/internal/provider/runway"
"backend/internal/repo" "backend/internal/repo"
@@ -28,14 +28,14 @@ import (
) )
var validTokenPools = map[string]string{ var validTokenPools = map[string]string{
"chatgpt": "openai", "chatgpt": "openai",
"adobe": "adobe", "adobe": "adobe",
"runway": "runway", "runway": "runway",
"leonardo": "leonardo", "leonardo": "leonardo",
"krea": "krea", "krea": "krea",
"imagine": "imagine", "imagine": "imagine",
"grok": "grok", "grok": "grok",
"custom": "custom", "custom": "custom",
"creativefabrica": "creativefabrica", "creativefabrica": "creativefabrica",
} }
@@ -1036,7 +1036,6 @@ func (s *TokenService) ImportGrokToken(ctx context.Context, ssoToken, tokenID st
return item, nil return item, nil
} }
func (s *TokenService) checkPendingGrok(tokenID, ssoToken string) { func (s *TokenService) checkPendingGrok(tokenID, ssoToken string) {
defer func() { defer func() {
if r := recover(); r != nil { 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 // 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 // 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). // 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), "/") baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
apiKey = strings.TrimSpace(apiKey) apiKey = strings.TrimSpace(apiKey)
protocol = strings.ToLower(strings.TrimSpace(protocol))
// Edit mode: tokenID points at an existing custom account. base_url required; // Edit mode: tokenID points at an existing custom account. base_url required;
// a blank key keeps the stored one. // a blank key keeps the stored one.
if strings.TrimSpace(tokenID) != "" { if strings.TrimSpace(tokenID) != "" {
@@ -1230,7 +1230,14 @@ func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey,
if baseURL == "" { if baseURL == "" {
return nil, errors.New("base_url required") 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 != "" { if m := strings.TrimSpace(models); m != "" {
meta["models"] = m meta["models"] = m
} }
@@ -1248,7 +1255,10 @@ func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey,
if baseURL == "" || apiKey == "" { if baseURL == "" || apiKey == "" {
return nil, errors.New("base_url and key required") 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 != "" { if m := strings.TrimSpace(models); m != "" {
meta["models"] = m meta["models"] = m
} }
@@ -1952,6 +1962,7 @@ func accountRow(item model.TokenAccount, inFlight int64) map[string]any {
"weight": item.Weight, "weight": item.Weight,
"concurrency": item.Concurrency, "concurrency": item.Concurrency,
"base_url": emptyToNil(strings.TrimSpace(stringValue(item.Meta["base_url"]))), "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"])), "models": strings.TrimSpace(stringValue(item.Meta["models"])),
} }
} }
+30 -5
View File
@@ -23,8 +23,8 @@ import (
"backend/internal/model" "backend/internal/model"
"backend/internal/provider/adobe" "backend/internal/provider/adobe"
"backend/internal/provider/chatgpt" "backend/internal/provider/chatgpt"
"backend/internal/provider/custom"
"backend/internal/provider/creativefabrica" "backend/internal/provider/creativefabrica"
"backend/internal/provider/custom"
"backend/internal/provider/grok" "backend/internal/provider/grok"
"backend/internal/provider/imagine" "backend/internal/provider/imagine"
"backend/internal/provider/krea" "backend/internal/provider/krea"
@@ -32,6 +32,7 @@ import (
"backend/internal/provider/runway" "backend/internal/provider/runway"
"backend/internal/repo" "backend/internal/repo"
"backend/internal/storage" "backend/internal/storage"
"gorm.io/gorm" "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.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail)
_ = s.tokens.TouchLastUsed(ctx, token.ID) _ = s.tokens.TouchLastUsed(ctx, token.ID)
baseURL := stringValue(token.Meta["base_url"]) 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 { if genErr == nil {
_, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{ _, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{
"last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0, "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.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail)
_ = s.tokens.TouchLastUsed(ctx, token.ID) _ = s.tokens.TouchLastUsed(ctx, token.ID)
baseURL := stringValue(token.Meta["base_url"]) 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 { if genErr == nil {
_, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{ _, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{
"last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0, "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) 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. // upstreamQuality maps a resolution tier to the OpenAI quality enum.
func upstreamQuality(resolution string) string { func upstreamQuality(resolution string) string {
switch strings.ToUpper(strings.TrimSpace(resolution)) { switch strings.ToUpper(strings.TrimSpace(resolution)) {
@@ -2975,8 +3000,8 @@ const leonardoPrivateSuffix = "-不卡人脸"
type leonardoVideoSpec struct { type leonardoVideoSpec struct {
upstream string upstream string
// long/short 是长边/短边像素,按比例组合成 16:9 或 9:16。 // long/short 是长边/短边像素,按比例组合成 16:9 或 9:16。
long int long int
short int short int
maxImages int maxImages int
maxAudios int maxAudios int
maxVideos int maxVideos int
+15 -1
View File
@@ -9,6 +9,7 @@ const emit = defineEmits(['close', 'imported'])
const isEdit = !!props.account const isEdit = !!props.account
const name = ref(props.account?.email || '') const name = ref(props.account?.email || '')
const baseUrl = ref(props.account?.base_url || '') const baseUrl = ref(props.account?.base_url || '')
const protocol = ref(props.account?.protocol === 'vivid' ? 'vivid' : 'openai')
const key = ref('') // edit: blank = keep existing key const key = ref('') // edit: blank = keep existing key
const allModels = ref([]) // existing models to pick from const allModels = ref([]) // existing models to pick from
const selected = ref(props.account?.models ? String(props.account.models).split(',').map((x) => x.trim()).filter(Boolean) : []) const selected = ref(props.account?.models ? String(props.account.models).split(',').map((x) => x.trim()).filter(Boolean) : [])
@@ -43,6 +44,7 @@ async function submit() {
base_url: baseUrl.value.trim(), base_url: baseUrl.value.trim(),
key: key.value.trim(), // blank in edit = keep existing key: key.value.trim(), // blank in edit = keep existing
models: selected.value.join(','), models: selected.value.join(','),
protocol: protocol.value,
weight: Number(weight.value) || 0, weight: Number(weight.value) || 0,
concurrency: Number(concurrency.value) || 1, concurrency: Number(concurrency.value) || 1,
})) }))
@@ -65,7 +67,7 @@ async function submit() {
@click.self="emit('close')"> @click.self="emit('close')">
<div class="card !shadow-xl mt-14 mb-14 w-full max-w-lg"> <div class="card !shadow-xl mt-14 mb-14 w-full max-w-lg">
<div class="px-5 py-4 border-b border-slate-100 flex items-center justify-between"> <div class="px-5 py-4 border-b border-slate-100 flex items-center justify-between">
<h2 class="text-sm font-semibold">{{ isEdit ? '编辑上游' : '添加上游(自定义 OpenAI 兼容)' }}</h2> <h2 class="text-sm font-semibold">{{ isEdit ? '编辑上游' : '添加上游' }}</h2>
<button @click="emit('close')" class="text-slate-400 hover:text-slate-700 transition-colors"> <button @click="emit('close')" class="text-slate-400 hover:text-slate-700 transition-colors">
<Icon name="close" class="w-5 h-5" /> <Icon name="close" class="w-5 h-5" />
</button> </button>
@@ -79,6 +81,18 @@ async function submit() {
<label class="text-xs text-slate-500">备注名</label> <label class="text-xs text-slate-500">备注名</label>
<input v-model="name" class="field" placeholder="例如:我的中转 / xx-api" /> <input v-model="name" class="field" placeholder="例如:我的中转 / xx-api" />
</div> </div>
<div>
<label class="text-xs text-slate-500">协议</label>
<div class="flex gap-2">
<button v-for="p in [{id:'openai',label:'OpenAI',desc:'/v1/images/* 标准格式'},{id:'vivid',label:'Vivid',desc:'/v1/generate NDJSON 流'}]" :key="p.id" type="button"
@click="protocol = p.id"
class="flex-1 rounded-lg ring-1 px-3 py-2 text-left transition-colors"
:class="protocol === p.id ? 'ring-sky-400 bg-sky-500/10' : 'ring-slate-200 bg-white hover:ring-slate-300'">
<div class="text-xs font-medium" :class="protocol === p.id ? 'text-sky-700' : 'text-slate-700'">{{ p.label }}</div>
<div class="text-[11px] text-slate-400 font-mono mt-0.5">{{ p.desc }}</div>
</button>
</div>
</div>
<div> <div>
<label class="text-xs text-slate-500">v1 URL <span class="text-rose-500">*</span></label> <label class="text-xs text-slate-500">v1 URL <span class="text-rose-500">*</span></label>
<input v-model="baseUrl" class="field font-mono text-xs" placeholder="https://api.example.com(无需 /v1 结尾)" /> <input v-model="baseUrl" class="field font-mono text-xs" placeholder="https://api.example.com(无需 /v1 结尾)" />