168 lines
5.7 KiB
Python
168 lines
5.7 KiB
Python
from http import HTTPStatus
|
||
from urllib.parse import urlparse, unquote
|
||
from pathlib import PurePosixPath
|
||
import requests
|
||
from dashscope import MultiModalConversation
|
||
from Config.Config import ALY_LLM_API_KEY
|
||
|
||
class QwenImageEditor:
|
||
"""Qwen Image Editing Toolkit"""
|
||
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-edit"
|
||
|
||
def edit_image(self, image_url, prompt, negative_prompt="", stream=False, watermark=True):
|
||
"""编辑图片
|
||
|
||
Args:
|
||
image_url: 原始图片URL
|
||
prompt: 编辑指令描述
|
||
negative_prompt: 负面提示词
|
||
stream: 是否流式返回
|
||
watermark: 是否添加水印
|
||
|
||
Returns:
|
||
dict: 包含编辑结果的字典
|
||
"""
|
||
try:
|
||
# 构建对话消息
|
||
messages = [{
|
||
"role": "user",
|
||
"content": [
|
||
{"image": image_url},
|
||
{"text": prompt}
|
||
]
|
||
}]
|
||
|
||
# 调用图片编辑API
|
||
response = MultiModalConversation.call(
|
||
api_key=self.api_key,
|
||
model=self.model,
|
||
messages=messages,
|
||
result_format='message',
|
||
stream=stream,
|
||
watermark=watermark,
|
||
negative_prompt=negative_prompt
|
||
)
|
||
|
||
# 处理响应结果
|
||
if response.status_code == HTTPStatus.OK:
|
||
# 提取编辑后的图片URL
|
||
edited_image_url = response.output.choices[0].message.content[0].get("image")
|
||
|
||
if not edited_image_url:
|
||
return {
|
||
'success': False,
|
||
'image_url': None,
|
||
'error_msg': 'API返回空结果，图片编辑失败'
|
||
}
|
||
|
||
return {
|
||
'success': True,
|
||
'image_url': edited_image_url,
|
||
'error_msg': None
|
||
}
|
||
else:
|
||
error_msg = f'API调用失败: status_code={response.status_code}, code={response.code}, message={response.message}'
|
||
return {
|
||
'success': False,
|
||
'image_url': None,
|
||
'error_msg': error_msg
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'image_url': None,
|
||
'error_msg': f'发生异常: {str(e)}'
|
||
}
|
||
|
||
def save_edited_image(self, image_url, save_dir='./'):
|
||
"""保存编辑后的图片到本地
|
||
|
||
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 edit_and_save_image(self, image_url, prompt, save_dir='./', negative_prompt=""):
|
||
"""编辑图片并保存到本地
|
||
|
||
Args:
|
||
image_url: 原始图片URL
|
||
prompt: 编辑指令描述
|
||
save_dir: 保存目录
|
||
negative_prompt: 负面提示词
|
||
|
||
Returns:
|
||
dict: 包含编辑和保存结果的字典
|
||
"""
|
||
# 编辑图片
|
||
edit_result = self.edit_image(image_url, prompt, negative_prompt)
|
||
|
||
if not edit_result['success']:
|
||
return edit_result
|
||
|
||
# 保存编辑后的图片
|
||
save_result = self.save_edited_image(edit_result['image_url'], save_dir)
|
||
|
||
if save_result['success']:
|
||
return {
|
||
'success': True,
|
||
'image_url': edit_result['image_url'],
|
||
'file_path': save_result['file_path'],
|
||
'error_msg': None
|
||
}
|
||
else:
|
||
return {
|
||
'success': False,
|
||
'image_url': edit_result['image_url'],
|
||
'file_path': None,
|
||
'error_msg': save_result['error_msg']
|
||
}
|
||
|
||
# 示例使用
|
||
if __name__ == "__main__":
|
||
# 创建编辑器实例
|
||
editor = QwenImageEditor()
|
||
|
||
# 示例参数
|
||
image_url = "https://dashscope.oss-cn-beijing.aliyuncs.com/images/dog_and_girl.jpeg"
|
||
edit_prompt = "将图中的人物改为站立姿势，弯腰握住狗的前爪"
|
||
|
||
print('----编辑图片，请等待任务执行----')
|
||
result = editor.edit_and_save_image(image_url, edit_prompt)
|
||
|
||
if result['success']:
|
||
print(f'编辑成功，图片URL: {result["image_url"]}')
|
||
print(f'保存成功，文件路径: {result["file_path"]}')
|
||
else:
|
||
print(f'编辑失败: {result["error_msg"]}')
|