156 lines
5.3 KiB
Python
156 lines
5.3 KiB
Python
from http import HTTPStatus
|
||
from urllib.parse import urlparse, unquote
|
||
from pathlib import PurePosixPath
|
||
import requests
|
||
# pip install -U dashscope
|
||
from dashscope import ImageSynthesis
|
||
from Config.Config import ALY_LLM_API_KEY
|
||
|
||
class QwenImageGenerator:
|
||
|
||
def __init__(self, api_key=None):
|
||
"""初始化生成器
|
||
|
||
Args:
|
||
api_key: 阿里云API密钥，如果为None则使用配置文件中的密钥
|
||
"""
|
||
self.api_key = api_key or ALY_LLM_API_KEY
|
||
self.model = "qwen-image"
|
||
|
||
def generate_image(self, prompt, n=1, size='1328*1328'):
|
||
"""生成图片
|
||
|
||
Args:
|
||
prompt: 图片描述提示词
|
||
n: 生成图片数量
|
||
size: 图片尺寸
|
||
|
||
Returns:
|
||
dict: 包含生成结果的字典，格式如下：
|
||
{
|
||
'success': bool, # 操作是否成功
|
||
'images': list, # 图片URL列表
|
||
'error_msg': str # 错误信息，如果成功则为None
|
||
}
|
||
"""
|
||
try:
|
||
# 调用文心一言图片生成API
|
||
rsp = ImageSynthesis.call(
|
||
api_key=self.api_key,
|
||
model=self.model,
|
||
prompt=prompt,
|
||
n=n,
|
||
size=size
|
||
)
|
||
|
||
# 处理响应结果
|
||
if rsp.status_code == HTTPStatus.OK:
|
||
# 提取图片URL
|
||
image_urls = [result.url for result in rsp.output.results]
|
||
return {
|
||
'success': True,
|
||
'images': image_urls,
|
||
'error_msg': None
|
||
}
|
||
else:
|
||
error_msg = f'API调用失败: status_code={rsp.status_code}, code={rsp.code}, message={rsp.message}'
|
||
return {
|
||
'success': False,
|
||
'images': [],
|
||
'error_msg': error_msg
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'images': [],
|
||
'error_msg': f'发生异常: {str(e)}'
|
||
}
|
||
|
||
def save_image_from_url(self, image_url, save_dir='./'):
|
||
"""从URL保存图片到本地
|
||
|
||
Args:
|
||
image_url: 图片URL
|
||
save_dir: 保存目录
|
||
|
||
Returns:
|
||
dict: 包含保存结果的字典
|
||
"""
|
||
try:
|
||
# 从URL解析文件名
|
||
file_name = PurePosixPath(unquote(urlparse(image_url).path)).parts[-1]
|
||
save_path = f'{save_dir}{file_name}'
|
||
|
||
# 下载并保存图片
|
||
with open(save_path, 'wb+') as f:
|
||
f.write(requests.get(image_url).content)
|
||
|
||
return {
|
||
'success': True,
|
||
'file_path': save_path,
|
||
'error_msg': None
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'file_path': None,
|
||
'error_msg': f'保存图片失败: {str(e)}'
|
||
}
|
||
|
||
def generate_and_save_images(self, prompt, n=1, size='1328*1328', save_dir='./'):
|
||
"""生成图片并保存到本地
|
||
|
||
Args:
|
||
prompt: 图片描述提示词
|
||
n: 生成图片数量
|
||
size: 图片尺寸
|
||
save_dir: 保存目录
|
||
|
||
Returns:
|
||
dict: 包含生成和保存结果的字典
|
||
"""
|
||
# 生成图片
|
||
generate_result = self.generate_image(prompt, n, size)
|
||
|
||
if not generate_result['success']:
|
||
return generate_result
|
||
|
||
# 保存所有生成的图片
|
||
saved_files = []
|
||
for image_url in generate_result['images']:
|
||
save_result = self.save_image_from_url(image_url, save_dir)
|
||
if save_result['success']:
|
||
saved_files.append(save_result['file_path'])
|
||
else:
|
||
return {
|
||
'success': False,
|
||
'images': generate_result['images'],
|
||
'saved_files': saved_files,
|
||
'error_msg': save_result['error_msg']
|
||
}
|
||
|
||
return {
|
||
'success': True,
|
||
'images': generate_result['images'],
|
||
'saved_files': saved_files,
|
||
'error_msg': None
|
||
}
|
||
|
||
# 示例使用
|
||
if __name__ == "__main__":
|
||
# 创建生成器实例
|
||
generator = QwenImageGenerator()
|
||
|
||
# 示例提示词
|
||
prompt = """一副典雅庄重的对联悬挂于厅堂之中，房间是个安静古典的中式布置，桌子上放着一些青花瓷，对联上左书"云端书声琅琅跨山海"，右书"智脉笔影翩翩启星辰"， 横批"东师理想"，字体飘逸，中间挂在一着一副中国风的画作，内容是岳阳楼。"""
|
||
|
||
print('----生成图片，请等待任务执行----')
|
||
result = generator.generate_and_save_images(prompt)
|
||
|
||
if result['success']:
|
||
print(f'生成成功，图片URL: {result["images"]}')
|
||
print(f'保存成功，文件路径: {result["saved_files"]}')
|
||
else:
|
||
print(f'生成失败: {result["error_msg"]}')
|
||
|