feat(upstream): custom 模型按媒体分类配置参考文件上限

- ModelConfig 新增 max_reference_videos / max_reference_audios (配合既有 max_reference_images);admin 读写接口透传 - custom 模型表单去掉「参考模式」下拉,改为参考图/参考视频/ 参考音频三个数量输入;合计张数走 totalRefLimit 校验 - vivid 上游 /v1/uploads 按文件魔数自动选 kind(image/video/ audio)与 mime;openai 协议上游遇视频/音频参考自动跳过换号 - videoCaps 让 custom 模型的分类上限在画图台/文档徽标生效 - reference_mode 请求级 override 放行 "media" Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-14 14:58:34 +08:00
parent babfc71f13
commit 1be1b00511
9 changed files with 236 additions and 57 deletions
@@ -95,7 +95,7 @@ type EventLog struct {
 	Duration   string         `gorm:"size:32"`
 	Refs       int            `gorm:"not null;default:0"`
 	DeAI       bool           `gorm:"not null;default:false"` // 去AI特征 was applied (image only)
-	RefFiles   datatypes.JSON `gorm:"type:jsonb"` // relative paths of saved reference images, for回显 on reload
+	RefFiles   datatypes.JSON `gorm:"type:jsonb"`             // relative paths of saved reference images, for回显 on reload
 	Source     string         `gorm:"size:32;index"`
 	// AccountID is the provider token/account chosen to fulfil this generation,
 	// stamped when the upstream call begins. Drives the accounts view's live
@@ -139,7 +139,11 @@ type ModelConfig struct {
 	DurationPricesAgent datatypes.JSONMap `gorm:"type:jsonb;column:duration_prices_agent"`
 	Durations           datatypes.JSON    `gorm:"type:jsonb"`
 	MaxReferenceImages  int               `gorm:"not null;default:0"`
-	ReferenceMode       string            `gorm:"size:32;not null;default:'none'"`
+	// Per-kind reference caps for upstreams that accept video/audio refs
+	// (vivid-protocol custom upstreams). 0 = that kind is not accepted.
+	MaxReferenceVideos int    `gorm:"not null;default:0"`
+	MaxReferenceAudios int    `gorm:"not null;default:0"`
+	ReferenceMode      string `gorm:"size:32;not null;default:'none'"`
 	// Custom-upstream models (provider="custom"): UpstreamModel is the model name
 	// sent to the upstream OpenAI-compatible API; the base_url + key live on the
 	// matching custom account (pool="custom", meta.base_url). Empty for built-ins.
@@ -263,15 +267,15 @@ func AutoMigrateModels() []any {
 type Order struct {
 	ID          string    `gorm:"primaryKey;size:40"`
 	UserID      string    `gorm:"size:32;index;not null"`
-	Amount      float64   `gorm:"not null"`               // 充值金额(元)
-	Points      int       `gorm:"not null"`               // 到账积分
-	PayType     string    `gorm:"size:16"`                // wxpay | alipay | admin | cdk
-	Status      string    `gorm:"size:16;index;not null"` // pending | paid | cancelled
+	Amount      float64   `gorm:"not null"`                              // 充值金额(元)
+	Points      int       `gorm:"not null"`                              // 到账积分
+	PayType     string    `gorm:"size:16"`                               // wxpay | alipay | admin | cdk
+	Status      string    `gorm:"size:16;index;not null"`                // pending | paid | cancelled
 	Source      string    `gorm:"size:16;index;not null;default:'epay'"` // epay | admin | cdk
-	Remark      string    `gorm:"type:text"`              // e.g. 兑换码 code / 管理员操作说明
-	TradeNo     string    `gorm:"size:64;index"`          // 易支付平台订单号
-	PayInfo     string    `gorm:"type:text"`              // 二维码 url / 跳转 url
-	PayInfoType string    `gorm:"size:16"`                // qrcode | jump | html | ...
+	Remark      string    `gorm:"type:text"`                             // e.g. 兑换码 code / 管理员操作说明
+	TradeNo     string    `gorm:"size:64;index"`                         // 易支付平台订单号
+	PayInfo     string    `gorm:"type:text"`                             // 二维码 url / 跳转 url
+	PayInfoType string    `gorm:"size:16"`                               // qrcode | jump | html | ...
 	ExpiresAt   time.Time `gorm:"index"`
 	PaidAt      *time.Time
 	CreatedAt   time.Time
@@ -294,10 +298,10 @@ type ConcurrencyGroup struct {
 // node/edge graph as a versioned JSON blob; saves are whole-document
 // replacements (last write wins), which matches the single-editor frontend.
 type Canvas struct {
-	ID        string    `gorm:"primaryKey;size:40"`
-	UserID    string    `gorm:"size:32;index;not null"`
-	Title     string    `gorm:"size:120;not null;default:'未命名画布'"`
-	Data      string    `gorm:"type:text"`
+	ID        string `gorm:"primaryKey;size:40"`
+	UserID    string `gorm:"size:32;index;not null"`
+	Title     string `gorm:"size:120;not null;default:'未命名画布'"`
+	Data      string `gorm:"type:text"`
 	CreatedAt time.Time
 	UpdatedAt time.Time `gorm:"index"`
 }
@@ -178,14 +178,13 @@ func (c *Client) vividPayload(ctx context.Context, baseURL, apiKey, model, promp
 	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[].
+// vividUploadRef runs the two-step upload: POST /v1/uploads {kind,mime} → PUT
+// the bytes to the presigned url → returns the key for refs[]. kind is image /
+// video / audio, derived from the file magic (same classes the upstream
+// accepts: png/jpeg/webp · mp4/webm/mov · mp3/wav/m4a/ogg).
 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})
+	kind, mime := vividRefKindMime(data)
+	raw, _ := json.Marshal(map[string]any{"kind": kind, "mime": mime})
 	up, err := c.doJSON(ctx, http.MethodPost, baseURL+"/v1/uploads", apiKey, raw)
 	if err != nil {
 		return "", err
@@ -213,6 +212,58 @@ func (c *Client) vividUploadRef(ctx context.Context, baseURL, apiKey string, dat
 	return key, nil
 }

+// vividRefKindMime classifies a reference blob into the vivid "kind" (image /
+// video / audio) plus a concrete mime. Magic bytes first (DetectContentType
+// misses mov/m4a/ogg-audio), then sniffing, then a png default.
+func vividRefKindMime(data []byte) (kind, mime string) {
+	n := len(data)
+	// ISOBMFF box: ftyp — mp4/mov/m4a share the container; the brand decides.
+	if n >= 12 && string(data[4:8]) == "ftyp" {
+		brand := string(data[8:12])
+		switch {
+		case strings.Contains(brand, "M4A"), strings.Contains(brand, "M4B"):
+			return "audio", "audio/mp4"
+		case strings.Contains(brand, "qt"):
+			return "video", "video/quicktime"
+		default:
+			return "video", "video/mp4"
+		}
+	}
+	// WebM/MKV EBML
+	if n >= 4 && data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3 {
+		return "video", "video/webm"
+	}
+	// MP3: ID3 tag or MPEG sync
+	if n >= 3 && string(data[:3]) == "ID3" {
+		return "audio", "audio/mpeg"
+	}
+	if n >= 2 && data[0] == 0xFF && (data[1]&0xE0) == 0xE0 {
+		return "audio", "audio/mpeg"
+	}
+	// WAV / OGG / FLAC
+	if n >= 4 && string(data[:4]) == "RIFF" {
+		return "audio", "audio/wav"
+	}
+	if n >= 4 && string(data[:4]) == "OggS" {
+		return "audio", "audio/ogg"
+	}
+	if n >= 4 && string(data[:4]) == "fLaC" {
+		return "audio", "audio/x-flac"
+	}
+	// Sniffed type (images land here; also catches anything odd above missed).
+	if m := http.DetectContentType(data); m != "" && m != "application/octet-stream" {
+		switch {
+		case strings.HasPrefix(m, "video/"):
+			return "video", m
+		case strings.HasPrefix(m, "audio/"), m == "application/ogg":
+			return "audio", m
+		case strings.HasPrefix(m, "image/"):
+			return "image", m
+		}
+	}
+	return "image", "image/png"
+}
+
 // 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).
@@ -141,6 +141,56 @@ func TestVividErrorCodes(t *testing.T) {
 	}
 }

+func TestVividRefKindMime(t *testing.T) {
+	cases := []struct {
+		name string
+		data []byte
+		kind string
+	}{
+		{"png", []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 1, 2, 3, 4}, "image"},
+		{"jpeg", []byte{0xFF, 0xD8, 0xFF, 0xE0, 1, 2, 3, 4, 5, 6}, "image"},
+		{"mp4", []byte{0, 0, 0, 0, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'}, "video"},
+		{"mov", []byte{0, 0, 0, 0, 'f', 't', 'y', 'p', 'q', 't', ' ', ' '}, "video"},
+		{"webm", []byte{0x1A, 0x45, 0xDF, 0xA3, 1, 2, 3, 4}, "video"},
+		{"mp3-id3", []byte{'I', 'D', '3', 4, 0, 0, 1, 2}, "audio"},
+		{"mp3-sync", []byte{0xFF, 0xFB, 1, 2, 3, 4}, "audio"},
+		{"wav", []byte{'R', 'I', 'F', 'F', 1, 2, 3, 4, 'W', 'A', 'V', 'E'}, "audio"},
+		{"ogg", []byte{'O', 'g', 'g', 'S', 1, 2, 3, 4}, "audio"},
+		{"m4a", []byte{0, 0, 0, 0, 'f', 't', 'y', 'p', 'M', '4', 'A', ' '}, "audio"},
+		{"empty", []byte{}, "image"},
+	}
+	for _, tc := range cases {
+		kind, mime := vividRefKindMime(tc.data)
+		if kind != tc.kind {
+			t.Errorf("%s: kind=%s mime=%s", tc.name, kind, mime)
+		}
+	}
+}
+
+func TestVividUploadRefKindVideo(t *testing.T) {
+	var gotKind, gotMime string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/v1/uploads" {
+			var req map[string]any
+			_ = json.NewDecoder(r.Body).Decode(&req)
+			gotKind, _ = req["kind"].(string)
+			gotMime, _ = req["mime"].(string)
+			fmt.Fprintf(w, `{"key":"k1","url":"http://%s/up"}`, r.Host)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+	webm := []byte{0x1A, 0x45, 0xDF, 0xA3, 1, 2, 3, 4, 5, 6, 7, 8}
+	key, err := NewClient().vividUploadRef(context.Background(), srv.URL, "vk", webm)
+	if err != nil {
+		t.Fatalf("upload: %v", err)
+	}
+	if key != "k1" || gotKind != "video" || gotMime != "video/webm" {
+		t.Fatalf("key=%q kind=%q mime=%q", key, gotKind, gotMime)
+	}
+}
+
 func TestVividUploadRefFlow(t *testing.T) {
 	var putOK bool
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -102,6 +102,8 @@ func (s *AdminReadService) ModelsView(ctx context.Context) ([]map[string]any, er
 			"duration_prices_agent": map[string]any(item.DurationPricesAgent),
 			"durations":             repo.JSONStrings(item.Durations),
 			"max_reference_images":  item.MaxReferenceImages,
+			"max_reference_videos":  item.MaxReferenceVideos,
+			"max_reference_audios":  item.MaxReferenceAudios,
 			"reference_mode":        item.ReferenceMode,
 			"weight":                item.Weight,
 			"generation_count":      item.GenerationCount,
@@ -401,6 +401,8 @@ func (s *AdminWriteService) CreateModel(ctx context.Context, body map[string]any
 		DurationPricesAgent: jsonMap(body["duration_prices_agent"]),
 		Durations:           jsonArray(body["durations"]),
 		MaxReferenceImages:  intValue(body["max_reference_images"]),
+		MaxReferenceVideos:  intValue(body["max_reference_videos"]),
+		MaxReferenceAudios:  intValue(body["max_reference_audios"]),
 		ReferenceMode:       defaultString(strings.TrimSpace(stringValue(body["reference_mode"])), "none"),
 		Weight:              intValue(body["weight"]),
 		CreatedAt:           time.Now(),
@@ -477,6 +479,12 @@ func (s *AdminWriteService) UpdateModel(ctx context.Context, modelID string, bod
 	if _, ok := body["max_reference_images"]; ok {
 		patch["max_reference_images"] = intValue(body["max_reference_images"])
 	}
+	if _, ok := body["max_reference_videos"]; ok {
+		patch["max_reference_videos"] = intValue(body["max_reference_videos"])
+	}
+	if _, ok := body["max_reference_audios"]; ok {
+		patch["max_reference_audios"] = intValue(body["max_reference_audios"])
+	}
 	if _, ok := body["reference_mode"]; ok {
 		patch["reference_mode"] = defaultString(strings.TrimSpace(stringValue(body["reference_mode"])), "none")
 	}
@@ -95,9 +95,10 @@ func (s *UserGenerationService) Generate(ctx context.Context, user *model.User,
 	}
 }

-// validateReferenceMode mirrors the /v1 checks: the override must be "frame"
-// or "asset", the model must support references at all, and frame mode carries
-// at most 2 images (first+last frame).
+// validateReferenceMode mirrors the /v1 checks: the override must be "frame",
+// "asset" or "media", the model must support references at all, and frame mode
+// carries at most 2 images (first+last frame). "media" = 参考视频/音频 — refs
+// may be video or audio blobs (vivid-style upstreams that accept all kinds).
 func validateReferenceMode(rm string, modelItem *model.ModelConfig, refCount int) error {
 	rm = strings.TrimSpace(rm)
 	if rm == "" {
@@ -107,8 +108,8 @@ func validateReferenceMode(rm string, modelItem *model.ModelConfig, refCount int
 	if supported == "" || supported == "none" {
 		return fmt.Errorf("%w: reference_mode not supported for this model", ErrUnsupportedParams)
 	}
-	if rm != "frame" && rm != "asset" {
-		return fmt.Errorf("%w: reference_mode must be 'frame' or 'asset'", ErrUnsupportedParams)
+	if rm != "frame" && rm != "asset" && rm != "media" {
+		return fmt.Errorf("%w: reference_mode must be 'frame', 'asset' or 'media'", ErrUnsupportedParams)
 	}
 	if rm == "frame" && refCount > 2 {
 		return fmt.Errorf("%w: frame mode supports at most 2 reference images (first+last frame), got %d", ErrUnsupportedParams, refCount)
@@ -224,7 +224,7 @@ type V1VideoRequest struct {
 	AspectRatio     string
 	Resolution      string
 	ReferenceImages []string
-	ReferenceMode   string // "frame" or "asset", overrides model default
+	ReferenceMode   string // "frame", "asset" or "media", overrides model default
 	// BaseURL — see V1ImageRequest.BaseURL.
 	BaseURL string
 	// AccountID — see V1ImageRequest.AccountID.
@@ -1273,7 +1273,7 @@ func (s *V1Service) prepareImage(ctx context.Context, principal *APIPrincipal, i
 	}
 	refLimit := 0
 	if modelItem.ImageToImage {
-		refLimit = modelItem.MaxReferenceImages
+		refLimit = totalRefLimit(modelItem)
 		if refLimit <= 0 {
 			refLimit = 1
 		}
@@ -1362,7 +1362,7 @@ func (s *V1Service) prepareVideo(ctx context.Context, principal *APIPrincipal, i
 	} else if !ok {
 		return nil, "", "", "", 0, ErrNoProviderAccount
 	}
-	refLimit := modelItem.MaxReferenceImages
+	refLimit := totalRefLimit(modelItem)
 	if refLimit <= 0 {
 		refLimit = 10
 	}
@@ -2120,6 +2120,44 @@ func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, mod
 	return nil, "", lastErr
 }

+// totalRefLimit is the aggregate reference-file cap across all media kinds.
+// Models with per-kind caps (max_reference_videos/audios) total them; legacy
+// models only have max_reference_images, so the sum degenerates correctly.
+func totalRefLimit(m *model.ModelConfig) int {
+	return m.MaxReferenceImages + m.MaxReferenceVideos + m.MaxReferenceAudios
+}
+
+// refKinds counts decoded refs by media class (image / video / audio) using the
+// same magic-byte detection the leonardo upload path relies on.
+func refKinds(refs [][]byte) (images, videos, audios int) {
+	for _, r := range refs {
+		switch detectMediaType(r) {
+		case "video":
+			videos++
+		case "audio":
+			audios++
+		default:
+			images++
+		}
+	}
+	return
+}
+
+// checkCustomRefCaps validates decoded refs against a custom model's per-kind
+// caps (Chinese messages mirror the leonardo spec errors).
+func checkCustomRefCaps(modelItem *model.ModelConfig, refs [][]byte) (images, videos, audios int, err error) {
+	images, videos, audios = refKinds(refs)
+	switch {
+	case videos > modelItem.MaxReferenceVideos:
+		err = fmt.Errorf("最多 %d 段视频参考", modelItem.MaxReferenceVideos)
+	case audios > modelItem.MaxReferenceAudios:
+		err = fmt.Errorf("最多 %d 段音频参考", modelItem.MaxReferenceAudios)
+	case images > modelItem.MaxReferenceImages:
+		err = fmt.Errorf("最多 %d 张图片参考", modelItem.MaxReferenceImages)
+	}
+	return
+}
+
 // customAccountServes reports whether a custom (upstream) account is usable for a
 // given model id: active, not dead, has a base_url, and its meta.models list (csv
 // of model ids it serves) contains the id. An empty models list serves ALL ids.
@@ -2201,7 +2239,11 @@ func (s *V1Service) generateCustomImage(ctx context.Context, eventID string, mod
 	if s.custom == nil {
 		return nil, "", errors.New("custom client not configured")
 	}
-	refs, err := decodeReferenceImages(in.ReferenceImages, max(1, modelItem.MaxReferenceImages))
+	refs, err := decodeReferenceImages(in.ReferenceImages, max(1, totalRefLimit(modelItem)))
+	if err != nil {
+		return nil, "", err
+	}
+	_, vids, auds, err := checkCustomRefCaps(modelItem, refs)
 	if err != nil {
 		return nil, "", err
 	}
@@ -2229,10 +2271,17 @@ 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"])
+			vivid := strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid")
+			if !vivid && (vids > 0 || auds > 0) {
+				// OpenAI-protocol upstreams only accept image refs — skip this
+				// account; another vivid-protocol account may still serve.
+				lastErr = errors.New("upstream does not accept video/audio refs")
+				return false, true
+			}
 			var d []byte
 			var u string
 			var genErr error
-			if strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") {
+			if 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)
@@ -2291,9 +2340,13 @@ func (s *V1Service) generateCustomVideo(ctx context.Context, eventID string, mod
 		return nil, "", ErrNoProviderAccount
 	}
 	size := upstreamVideoSize(aspectRatio, resolution)
-	// Optional reference frames (image-to-video / first-last frames) — forwarded
-	// to the upstream as multipart input_reference[] files.
-	frames, err := decodeReferenceImages(in.ReferenceImages, max(1, modelItem.MaxReferenceImages))
+	// Optional reference files (images / videos / audios for vivid upstreams;
+	// images only for OpenAI-protocol ones).
+	frames, err := decodeReferenceImages(in.ReferenceImages, max(1, totalRefLimit(modelItem)))
+	if err != nil {
+		return nil, "", err
+	}
+	_, vids, auds, err := checkCustomRefCaps(modelItem, frames)
 	if err != nil {
 		return nil, "", err
 	}
@@ -2311,10 +2364,15 @@ 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"])
+			vivid := strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid")
+			if !vivid && (vids > 0 || auds > 0) {
+				lastErr = errors.New("upstream does not accept video/audio refs")
+				return false, true
+			}
 			var d []byte
 			var url string
 			var genErr error
-			if strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") {
+			if 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)