Skip to content

Commit c3777fc

Browse files
timerringtimerring
authored
feat: baidu ernie (#264)
* feat: baidu ernie * docs: update docs
1 parent fd2426d commit c3777fc

6 files changed

Lines changed: 109 additions & 3 deletions

File tree

‎README.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
<img src="assets/minimax-text.svg" alt="Minimax" width="60" height="60" />
2222
<img src="assets/siliconcloud-color.svg" alt="SiliconFlow" width="15" height="60" />
2323
<img src="assets/siliconcloud-text.svg" alt="SiliconFlow" width="100" height="60" />
24+
<img src="assets/wenxin-color.svg" alt="Baidu ERNIE" width="60" height="60" />
2425
</div>
2526

2627
## 1. Introduction
@@ -49,6 +50,7 @@
4950
- `Minimax image-01`
5051
- `Kwai Kolors`
5152
- `Tencent Hunyuan`
53+
- `Baidu ERNIE`
5254

5355
项目架构流程如下：
5456

@@ -215,6 +217,12 @@ MLLM 模型主要用于自动切片后的切片标题生成，此功能默认关
215217
216218
请自行[注册账号](https://console.cloud.tencent.com/cam/capi)并申请 API Key，填写到 `bilive.toml` 文件中对应的 `TENCENT_SECRET_ID` 和 `TENCENT_SECRET_KEY` 中。
217219

220+
##### 3.3.4 Baidu ERNIE 模型
221+
222+
> 如需使用 Baidu ERNIE 模型，请将 `IMAGE_GEN_MODEL` 参数设置为 `baidu`。
223+
224+
请自行[注册账号](https://console.bce.baidu.com/iam/key/list)并申请 API Key，填写到 `bilive.toml` 文件中对应的 `BAIDU_API_KEY` 中。
225+
218226
#### 4. bilitool 登录
219227

220228
> 由于一般日志打印不出二维码效果（docker 的日志不确定是否能打印，等发布新image时再修改，docker 版本请先参考文档 [bilive](https://bilive.timerring.com)，本 README 只针对源码部署），所以这步需要提前在机器上安装 [bilitool](https://github.com/timerring/bilitool):

‎assets/wenxin-color.svg‎

Lines changed: 1 addition & 0 deletions
Loading

‎bilive.toml‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,9 @@ qwen_api_key = "" # Apply for your own Qwen API key at https://bailian.console.a
3737

3838
[cover]
3939
generate_cover = false # whether to generate cover
40-
image_gen_model = "minimax" # the image generation model, can be "minimax" or "siliconflow" or "tencent"
40+
image_gen_model = "minimax" # the image generation model, can be "minimax" or "siliconflow" or "tencent" or "baidu"
4141
minimax_api_key = "" # Apply for your own Minimax API key at https://platform.minimaxi.com/user-center/basic-information/interface-key
4242
siliconflow_api_key = "" # Apply for your own SiliconFlow API key at https://cloud.siliconflow.cn/i/3Szr5BVg
4343
tencent_secret_id = "" # Apply for your own Tencent Cloud API key at https://console.cloud.tencent.com/cam/capi
44-
tencent_secret_key = "" # Apply for your own Tencent Cloud secret key as above
44+
tencent_secret_key = "" # Apply for your own Tencent Cloud secret key as above
45+
baidu_api_key = "" # Apply for your own Baidu API key at https://console.bce.baidu.com/iam/key/list

‎src/config.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,4 +77,5 @@ def get_interface_config():
7777
MINIMAX_API_KEY = config.get('cover', {}).get('minimax_api_key')
7878
SILICONFLOW_API_KEY = config.get('cover', {}).get('siliconflow_api_key')
7979
TENCENT_SECRET_ID = config.get('cover', {}).get('tencent_secret_id')
80-
TENCENT_SECRET_KEY = config.get('cover', {}).get('tencent_secret_key')
80+
TENCENT_SECRET_KEY = config.get('cover', {}).get('tencent_secret_key')
81+
BAIDU_API_KEY = config.get('cover', {}).get('baidu_api_key')

‎src/cover/cover_generator.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,10 @@ def wrapper(video_path):
6363
from .image_model_sdk.tencent_sdk import hunyuan_generate_cover
6464

6565
return hunyuan_generate_cover(cover_path)
66+
elif model_type == "baidu":
67+
from .image_model_sdk.baidu_sdk import baidu_generate_cover
68+
69+
return baidu_generate_cover(cover_path)
6670
else:
6771
upload_log.error(f"Unsupported model type: {model_type}")
6872
return None

‎src/cover/image_model_sdk/baidu_sdk.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import requests
2+
import json
3+
import os
4+
import base64
5+
from PIL import Image
6+
from io import BytesIO
7+
import time
8+
from src.upload.bilitool.bilitool.model.model import Model
9+
from src.config import BAIDU_API_KEY
10+
11+
def cover_up(img: str):
12+
"""Upload the cover image
13+
Parameters
14+
----------
15+
- img: img path or stream
16+
Returns
17+
-------
18+
- url: str
19+
the url of the cover image in bili server
20+
"""
21+
from PIL import Image
22+
from io import BytesIO
23+
request = requests.Session()
24+
request.headers = {
25+
'user-agent': "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/63.0.3239.108",
26+
'referer': "https://www.bilibili.com/",
27+
'connection': 'keep-alive'
28+
}
29+
model = Model().get_config()
30+
request.cookies.set('SESSDATA', model['cookies']['SESSDATA'])
31+
with Image.open(img) as im:
32+
# you should keep the image ratio 16:10
33+
xsize, ysize = im.size
34+
if xsize / ysize > 1.6:
35+
delta = xsize - ysize * 1.6
36+
region = im.crop((delta / 2, 0, xsize - delta / 2, ysize))
37+
else:
38+
delta = ysize - xsize * 10 / 16
39+
region = im.crop((0, delta / 2, xsize, ysize - delta / 2))
40+
buffered = BytesIO()
41+
region.save(buffered, format=im.format)
42+
r = request.post(
43+
url="https://member.bilibili.com/x/vu/web/cover/up",
44+
data={
45+
"cover": b"data:image/jpeg;base64,"
46+
+ (base64.b64encode(buffered.getvalue())),
47+
"csrf": model['cookies']['bili_jct']
48+
},
49+
timeout=30,
50+
)
51+
buffered.close()
52+
res = r.json()
53+
if res.get("data") is None:
54+
raise Exception(res)
55+
print(res["data"]["url"], flush=True)
56+
return res["data"]["url"]
57+
58+
def baidu_generate_cover(your_file_path):
59+
try:
60+
cover_url = cover_up(your_file_path)
61+
62+
url = "https://qianfan.baidubce.com/v2/images/generations"
63+
payload = json.dumps({
64+
"model": "irag-1.0",
65+
"prompt": "这是视频截图，请根据该图生成对应的动漫类型的封面",
66+
"refer_image": cover_url
67+
})
68+
headers = {
69+
'Content-Type': 'application/json',
70+
'Authorization': f'Bearer {BAIDU_API_KEY}'
71+
}
72+
73+
response = requests.request("POST", url, headers=headers, data=payload)
74+
if response.status_code == 200:
75+
image_url=response.json()['data'][0]['url']
76+
img_data = requests.get(image_url).content
77+
cover_name = time.strftime("%Y%m%d%H%M%S") + ".png"
78+
temp_cover_path = os.path.join(os.path.dirname(your_file_path), cover_name)
79+
with open(temp_cover_path, "wb") as handler:
80+
handler.write(img_data)
81+
os.remove(your_file_path)
82+
return temp_cover_path
83+
else:
84+
print(response.text, flush=True)
85+
return None
86+
except Exception as e:
87+
print(e, flush=True)
88+
return None
89+
90+
if __name__ == '__main__':
91+
print(baidu_generate_cover(""))

0 commit comments

Comments
 (0)