118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
import json
|
||
import requests
|
||
import logging
|
||
|
||
from JiMeng.Kit.VolcanoConst import VOLCANO_API_KEY
|
||
|
||
# 配置日志
|
||
logger = logging.getLogger(__name__)
|
||
logger.level = logging.INFO
|
||
console_handler = logging.StreamHandler()
|
||
console_handler.setLevel(logging.INFO)
|
||
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||
console_handler.setFormatter(formatter)
|
||
logger.addHandler(console_handler)
|
||
|
||
class FenJingTouGenerator:
|
||
def __init__(self, api_key=VOLCANO_API_KEY):
|
||
"""
|
||
初始化分镜头提示词生成器
|
||
:param api_key: 火山引擎API密钥
|
||
"""
|
||
self.api_key = api_key
|
||
self.url = 'https://ark.cn-beijing.volces.com/api/v3/chat/completions'
|
||
self.headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {self.api_key}'
|
||
}
|
||
self.model = 'doubao-1-5-vision-pro-32k-250115'
|
||
|
||
def generate_fen_jing_tou(self, image_url, prompt_text, timeout=60):
|
||
"""
|
||
生成分镜头提示词
|
||
:param image_url: 首帧图片URL
|
||
:param prompt_text: 提示文本
|
||
:param timeout: 请求超时时间(秒)
|
||
:return: 生成的分镜头提示词内容，如果失败则返回None
|
||
"""
|
||
# 构建请求体
|
||
payload = {
|
||
'model': self.model,
|
||
'messages': [
|
||
{
|
||
'content': [
|
||
{
|
||
'image_url': {
|
||
'url': image_url
|
||
},
|
||
'type': 'image_url'
|
||
},
|
||
{
|
||
'text': prompt_text,
|
||
'type': 'text'
|
||
}
|
||
],
|
||
'role': 'user'
|
||
}
|
||
]
|
||
}
|
||
|
||
try:
|
||
# 发送POST请求
|
||
response = requests.post(
|
||
self.url,
|
||
headers=self.headers,
|
||
data=json.dumps(payload),
|
||
timeout=timeout
|
||
)
|
||
# 检查响应状态码
|
||
response.raise_for_status()
|
||
# 解析响应数据
|
||
result = response.json()
|
||
|
||
# 提取生成的内容
|
||
if 'choices' in result and len(result['choices']) > 0:
|
||
choice = result['choices'][0]
|
||
if 'message' in choice and 'content' in choice['message']:
|
||
content = choice['message']['content']
|
||
logger.info(f'分镜头提示词生成成功')
|
||
return content
|
||
else:
|
||
logger.error('响应中未找到message或content字段')
|
||
return None
|
||
else:
|
||
logger.error('响应中未找到choices字段或choices为空')
|
||
return None
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
logger.error(f'请求发生错误: {e}')
|
||
if hasattr(e, 'response') and e.response is not None:
|
||
logger.error(f'错误响应状态码: {e.response.status_code}')
|
||
logger.error(f'错误响应内容: {e.response.text}')
|
||
return None
|
||
|
||
except json.JSONDecodeError:
|
||
logger.error('响应解析失败: 返回的内容不是有效的JSON格式')
|
||
if 'response' in locals():
|
||
logger.error(f'响应内容: {response.text}')
|
||
return None
|
||
|
||
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 创建实例并运行
|
||
generator = FenJingTouGenerator()
|
||
|
||
"""主函数，用于测试分镜头提示词生成功能"""
|
||
image_url = 'https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/JiMeng/d4186502-5d6e-404b-890c-1f40fbbc1728.jpg'
|
||
prompt_text = '以这张图片为首帧，帮我生成一个5秒的视频提示词。提示词包括：镜号、运镜、画面内容'
|
||
|
||
# 生成分镜头提示词
|
||
content = generator.generate_fen_jing_tou(image_url, prompt_text)
|
||
|
||
# 打印结果
|
||
if content:
|
||
print(content)
|
||
else:
|
||
print('分镜头提示词生成失败') |