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_image_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生图: `/gen` 后接提示词即可生成\n\n⚙️模型配置:\n/gen: 默认模型 (Seedream)\n/z: Z-Image\n/qwi: Qwen-Image\n/sd: Seedream-4.5\n/flux: Flux",
20 "seedream":
21 [
22 {
23 "model_id": "doubao-seedream-4-5-251128",
24 "model_name": "Seedream-4.5",
25 "api_type": "openai",
26 "api_keys": "key1,key2,key3,...",
27 "max_reference_img": 14,
28 "client_config": { "base_url": "https://ark.cn-beijing.volces.com/api/v3" },
29 "generate_config": {
30 "size": "4K",
31 "stream": false,
32 "response_format": "url",
33 "extra_body": {
34 "watermark": false,
35 "sequential_image_generation": "auto",
36 "sequential_image_generation_options": { "max_images": 4 },
37 "optimize_prompt_options": { "mode": "standard" }
38 }
39 }
40 }
41 ],
42
43 "zimage":
44 [
45 {
46 "model_id": "z-image-turbo",
47 "model_name": "Z-Image",
48 "api_type": "openai",
49 "api_keys": "openai-key-1,openai-key-2,openai-key-3,...",
50 "client_config": { "base_url": "https://api.openai.com/v1" },
51 "generate_config": { "response_format": "b64_json" }
52 },
53 {
54 "model_name": "Z-Image",
55 "api_type": "post",
56 "base_url": "https://dashscope.aliyuncs.com",
57 "api_paths": {
58 "img_gen": "/api/v1/services/aigc/multimodal-generation/generation"
59 },
60 "headers": {
61 "Authorization": "Bearer DASHSCOPE_API_KEY",
62 "Content-Type": "application/json"
63 },
64 "body": {
65 "model": "z-image-turbo",
66 "input": {"messages":[{"role":"user","content":[{"text":"%PROMPT%"}]}]},
67 "parameters": {
68 "prompt_extend": true,
69 "size": "2048*1152"
70 }
71 }
72 },
73 {
74 "model_name": "Z-Image",
75 "api_type": "post",
76 "base_url": "https://api-inference.modelscope.cn",
77 "api_paths": {
78 "img_gen": "/v1/images/generations",
79 "task_check": "/v1/tasks/%TASK_ID%"
80 },
81 "headers": {
82 "Authorization": "Bearer MODELSCOPE_API_KEY",
83 "Content-Type": "application/json",
84 "X-ModelScope-Async-Mode": "true"
85 },
86 "body": {
87 "model": "Tongyi-MAI/Z-Image-Turbo",
88 "prompt": "%PROMPT%"
89 }
90 }
91 ],
92 "qwen-image":
93 [
94 {
95 "model_name": "Qwen-Image-2512",
96 "api_type": "post",
97 "base_url": "https://api-inference.modelscope.cn",
98 "api_paths": {
99 "img_gen": "/v1/images/generations",
100 "task_check": "/v1/tasks/%TASK_ID%"
101 },
102 "headers": {
103 "Authorization": "Bearer MODELSCOPE_API_KEY",
104 "Content-Type": "application/json",
105 "X-ModelScope-Async-Mode": "true"
106 },
107 "body": {
108 "model": "Qwen/Qwen-Image-2512",
109 "prompt": "%PROMPT%"
110 }
111 },
112 {
113 "model_name": "Qwen-Image",
114 "api_type": "post",
115 "base_url": "https://dashscope.aliyuncs.com",
116 "api_paths": {
117 "img_gen": "/api/v1/services/aigc/multimodal-generation/generation"
118 },
119 "headers": {
120 "Authorization": "Bearer DASHSCOPE_API_KEY",
121 "Content-Type": "application/json"
122 },
123 "body": {
124 "model": "qwen-image",
125 "input": { "messages": [{"role":"user","content": [{"text":"%PROMPT%"}]}]},
126 "parameters": { "watermark": false }
127 }
128 }
129 ]
130 }
131
132 Suppose this message is:
133 Message(text="/gen A cute cat") -> use `AI.IMG_GENERATION_DEFAULT_MODEL` as model_alias
134 Message(text="/gen @seedream hello") -> use `seedream` as model_alias
135
136 Returns: list of model config
137 [
138 {
139 "model_id": "doubao-seedream-4-5-251128",
140 "model_name": "Seedream-4.5",
141 "api_type": "openai",
142 "api_keys": "key1,key2,key3,...",
143 "client_config": { "base_url": "https://ark.cn-beijing.volces.com/api/v3" },
144 "generate_config": {
145 "size": "4K",
146 "stream": false,
147 "response_format": "url",
148 "extra_body": {
149 "watermark": false,
150 "sequential_image_generation": "auto",
151 "sequential_image_generation_options": { "max_images": 4 },
152 "optimize_prompt_options": { "mode": "standard" }
153 }
154 }
155 }
156 ]
157 """
158 texts = str(message.content).strip()
159 if matched := re.match(rf"^{PREFIX.AI_IMG_GENERATION}\s+@([a-zA-Z0-9_\-\.]+)(\s+)?", texts): # match /gen @custom_model_id
160 model_alias = matched.group(1).strip()
161 return await get_config_by_model_alias(model_alias)
162 return await get_config_by_model_alias(AI.IMG_GENERATION_DEFAULT_MODEL)
163
164
165async def get_config_by_model_alias(model_alias: str) -> list[dict]:
166 """Get model config by model_alias.
167
168 Returns:
169 model_config
170 """
171 kv = await get_cf_kv(AI.IMG_MODEL_CONFIG_KEY, cache_ttl=600, silent=True)
172
173 custom_config = kv.get(model_alias, [])
174 if not custom_config:
175 logger.warning(f"Model `{model_alias}` is not configured in KV, using default config")
176 default_config = kv.get(AI.IMG_GENERATION_DEFAULT_MODEL, [])
177 if not default_config:
178 logger.warning(f"CF-KV key `{AI.IMG_MODEL_CONFIG_KEY}` does not has default `{AI.IMG_GENERATION_DEFAULT_MODEL}` field")
179 return default_config
180 return custom_config