main
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3import re
4
5from loguru import logger
6from pyrogram.types import Message
7
8from config import AI, PREFIX
9from database.kv import get_cf_kv
10
11
12async def get_video_model_configs(message: Message) -> list[dict]:
13 r"""Get model configs based on the message.
14
15 Model config is retrieved from CF-KV with key: {AI.IMG_MODEL_CONFIG_KEY}
16
17 A sample config:
18 {
19 "docs": "📽**AI视频**: `/gvid` + 提示词\n默认使用**Seedance-1.5-Pro**模型\n\n🔄使用以下命令强制切换模型:\n/sdt2v: Seedance 文生视频\n/sdf2v: Seedance 首帧生视频\n/sdfl2v: Seedance 首尾帧生视频",
20 "seedance": [
21 {
22 "model_name": "Seedance-1.5-Pro",
23 "api_type": "post",
24 "generation_type": "text_to_video",
25 "base_url": "https://ark.cn-beijing.volces.com",
26 "api_paths": {
27 "video_gen": "/api/v3/contents/generations/tasks",
28 "task_check": "/api/v3/contents/generations/tasks/%TASK_ID%"
29 },
30 "headers": {
31 "Authorization": "Bearer ARK_API_KEY",
32 "Content-Type": "application/json"
33 },
34 "body": {
35 "model": "doubao-seedance-1-5-pro-251215",
36 "content": [
37 {
38 "type": "text",
39 "text": "%PROMPT%"
40 }
41 ],
42 "generate_audio": true,
43 "resolution": "1080p",
44 "ratio": "adaptive",
45 "duration": -1,
46 "watermark": false
47 }
48 }
49 ],
50 "seedance-first-frame": [
51 {
52 "model_name": "Seedance-1.5-Pro",
53 "api_type": "post",
54 "generation_type": "first_frame",
55 "base_url": "https://ark.cn-beijing.volces.com",
56 "api_paths": {
57 "video_gen": "/api/v3/contents/generations/tasks",
58 "task_check": "/api/v3/contents/generations/tasks/%TASK_ID%"
59 },
60 "headers": {
61 "Authorization": "Bearer ARK_API_KEY",
62 "Content-Type": "application/json"
63 },
64 "body": {
65 "model": "doubao-seedance-1-5-pro-251215",
66 "content": [
67 {
68 "type": "text",
69 "text": "%PROMPT%"
70 },
71 {
72 "type": "image_url",
73 "role": "first_frame",
74 "image_url": {
75 "url": "%FIRST_FRAME%"
76 }
77 }
78 ],
79 "generate_audio": true,
80 "resolution": "720p",
81 "ratio": "adaptive",
82 "duration": -1,
83 "watermark": false
84 }
85 }
86 ],
87 "seedance-first-last-frame": [
88 {
89 "model_name": "Seedance-1.5-Pro",
90 "api_type": "post",
91 "generation_type": "first_last_frame",
92 "base_url": "https://ark.cn-beijing.volces.com",
93 "api_paths": {
94 "video_gen": "/api/v3/contents/generations/tasks",
95 "task_check": "/api/v3/contents/generations/tasks/%TASK_ID%"
96 },
97 "headers": {
98 "Authorization": "Bearer ARK_API_KEY",
99 "Content-Type": "application/json"
100 },
101 "body": {
102 "model": "doubao-seedance-1-5-pro-251215",
103 "content": [
104 {
105 "type": "text",
106 "text": "%PROMPT%"
107 },
108 {
109 "type": "image_url",
110 "role": "first_frame",
111 "image_url": {
112 "url": "%FIRST_FRAME%"
113 }
114 },
115 {
116 "type": "image_url",
117 "role": "last_frame",
118 "image_url": {
119 "url": "%LAST_FRAME%"
120 }
121 }
122 ],
123 "generate_audio": true,
124 "resolution": "720p",
125 "duration": -1,
126 "watermark": false
127 }
128 }
129 ]
130 }
131
132 Suppose this message is:
133 Message(text="/gvid prompt") -> use `AI.VIDEO_GENERATION_DEFAULT_MODEL` as model_alias
134 Message(text="/gvid @seedance-first-frame prompt") -> use `seedance-first-frame` as model_alias
135
136 Returns:
137 list of model config
138 """
139 texts = str(message.content).strip()
140 if matched := re.match(rf"^{PREFIX.AI_VIDEO_GENERATION}\s+@([a-zA-Z0-9_\-\.]+)(\s+)?", texts): # match /gvid @custom_model_id
141 model_alias = matched.group(1).strip()
142 return await get_config_by_model_alias(model_alias)
143 return await get_config_by_model_alias(AI.VIDEO_GENERATION_DEFAULT_MODEL)
144
145
146async def get_config_by_model_alias(model_alias: str) -> list[dict]:
147 """Get model config by model_alias.
148
149 Returns:
150 model_config
151 """
152 kv = await get_cf_kv(AI.VIDEO_MODEL_CONFIG_KEY, cache_ttl=600, silent=True)
153
154 custom_config = kv.get(model_alias, [])
155 if not custom_config:
156 logger.warning(f"Model `{model_alias}` is not configured in KV, using default config")
157 default_config = kv.get(AI.VIDEO_GENERATION_DEFAULT_MODEL, [])
158 if not default_config:
159 logger.warning(f"CF-KV key `{AI.VIDEO_MODEL_CONFIG_KEY}` does not has default `{AI.VIDEO_GENERATION_DEFAULT_MODEL}` field")
160 return default_config
161 return custom_config