61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
from http import HTTPStatus
|
||
from urllib.parse import urlparse, unquote
|
||
from pathlib import PurePosixPath
|
||
|
||
import requests
|
||
from dashscope import ImageSynthesis
|
||
import oss2
|
||
from WxMini.Milvus.Config.MulvusConfig import *
|
||
|
||
|
||
def generate_image(prompt, model="wanx2.1-t2i-turbo", n=1, size='1024*1024'):
|
||
# 初始化 OSS 客户端
|
||
auth = oss2.Auth(ACCESS_KEY_ID, ACCESS_KEY_SECRET)
|
||
bucket = oss2.Bucket(auth, ENDPOINT, BUCKET_NAME)
|
||
"""
|
||
根据提示词生成图片并直接上传到 OSS
|
||
:param prompt: 提示词
|
||
:param model: 使用的模型，默认为 "wanx2.1-t2i-turbo"
|
||
:param n: 生成图片的数量，默认为 1
|
||
:param size: 图片尺寸，默认为 '1024*1024'
|
||
:return: (success, message) success 为 True 表示成功，False 表示失败；message 为提示信息
|
||
"""
|
||
try:
|
||
rsp = ImageSynthesis.call(api_key=MODEL_API_KEY, model=model, prompt=prompt, n=n, size=size)
|
||
if rsp.status_code == HTTPStatus.OK:
|
||
# 遍历生成的图片
|
||
for result in rsp.output.results:
|
||
# 获取图片 URL
|
||
image_url = result.url
|
||
# 下载图片数据
|
||
image_data = requests.get(image_url).content
|
||
# 生成 OSS 文件名
|
||
file_name = PurePosixPath(unquote(urlparse(image_url).path)).parts[-1]
|
||
# 上传到 OSS
|
||
key = 'Upload/' + file_name
|
||
bucket.put_object(key, image_data)
|
||
return True, f'{OSS_PREFIX}{key}'
|
||
else:
|
||
# 生成失败，返回错误信息
|
||
error_message = f"同步调用失败，状态码: {rsp.status_code}, 错误码: {rsp.code}, 错误信息: {rsp.message}"
|
||
print(error_message)
|
||
return False, error_message
|
||
except Exception as e:
|
||
# 捕获异常，返回错误信息
|
||
error_message = f"生成图片时发生异常: {str(e)}"
|
||
print(error_message)
|
||
return False, error_message
|
||
|
||
|
||
# 获取图片的宽高
|
||
def getImgWidthHeight(img_url):
|
||
try:
|
||
url = f'{img_url}?x-oss-process=image/info'
|
||
response = requests.get(url)
|
||
jo = response.json()
|
||
width = int(jo['ImageWidth']['value'])
|
||
height = int(jo['ImageHeight']['value'])
|
||
return width, height
|
||
except:
|
||
return 182, 182
|