init
This commit is contained in:
219
services/service_qwen.py
Normal file
219
services/service_qwen.py
Normal file
@@ -0,0 +1,219 @@
|
||||
import time
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
api_key = '8e32d44e3007447cb4be6ee52c5d3110'
|
||||
|
||||
|
||||
class UploadInfo(BaseModel):
|
||||
fileName: str
|
||||
fileType: str
|
||||
|
||||
|
||||
class CreateInfo(BaseModel):
|
||||
taskId: str # 创建的任务 ID，可用于查询状态或获取结果
|
||||
taskStatus: str # 初始状态，可能为：QUEUED、RUNNING、FAILED
|
||||
clientId: str # 平台内部标识，用于排错，无需关注
|
||||
netWssUrl: str # WebSocket 地址（当前不稳定，不推荐使用）
|
||||
promptTips: str # ComfyUI 校验信息（字符串格式的 JSON），可用于识别配置异常节点
|
||||
|
||||
|
||||
class RunHubResponse(BaseModel):
|
||||
code: int # 状态码，0 表示成功
|
||||
msg: str # 提示信息
|
||||
data: UploadInfo | CreateInfo | str | None = None # 数据对象
|
||||
|
||||
class Config:
|
||||
extra = 'allow' # 允许添加额外字段
|
||||
|
||||
|
||||
async def upload(img_path: str) -> RunHubResponse:
|
||||
with open(img_path, 'rb') as f:
|
||||
img_data = f.read()
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field('apiKey', api_key)
|
||||
form.add_field('file', img_data, filename='image.jpg', content_type='image/jpeg')
|
||||
form.add_field('fileType', 'image')
|
||||
|
||||
url = 'https://www.runninghub.cn/task/openapi/upload'
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, data=form) as resp:
|
||||
response = await resp.json()
|
||||
|
||||
return RunHubResponse.model_validate(response)
|
||||
|
||||
|
||||
async def create(workflow_id: str, node_info_list: list[dict[str, str]]) -> RunHubResponse:
|
||||
url = 'https://www.runninghub.cn/task/openapi/create'
|
||||
json_data = {'apiKey': api_key, 'workflowId': workflow_id, 'nodeInfoList': node_info_list}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, json=json_data) as resp:
|
||||
response = await resp.json()
|
||||
|
||||
return RunHubResponse.model_validate(response)
|
||||
|
||||
|
||||
async def status(task_id: str) -> RunHubResponse:
|
||||
# 查询状态
|
||||
url = 'https://www.runninghub.cn/task/openapi/status'
|
||||
payload = {'apiKey': api_key, 'taskId': task_id}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, json=payload) as resp:
|
||||
response = await resp.json()
|
||||
|
||||
# ["QUEUED","RUNNING","FAILED","SUCCESS"]
|
||||
return RunHubResponse.model_validate(response)
|
||||
|
||||
|
||||
async def outputs(task_id: str) -> dict:
|
||||
# 获取结果
|
||||
url = 'https://www.runninghub.cn/task/openapi/outputs'
|
||||
payload = {'apiKey': api_key, 'taskId': task_id}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, json=payload) as resp:
|
||||
response = await resp.json()
|
||||
|
||||
return response
|
||||
|
||||
|
||||
async def 花纹提取_api(img_path: str, save_path: str, prompt: str = '') -> bool:
|
||||
"""
|
||||
异步花纹提取API
|
||||
|
||||
Args:
|
||||
img_path: 输入图片路径
|
||||
save_path: 输出图片路径
|
||||
prompt: 自定义提示词，为空则使用默认提示词
|
||||
|
||||
Returns:
|
||||
bool: 处理是否成功
|
||||
"""
|
||||
try:
|
||||
upload_res = await upload(img_path=img_path)
|
||||
if upload_res.code != 0 or not upload_res.data:
|
||||
logger.error(f"Qwen上传失败: code={upload_res.code}, msg={upload_res.msg}")
|
||||
return False
|
||||
|
||||
# 确保 data 是 UploadInfo 类型
|
||||
if not hasattr(upload_res.data, 'fileName'):
|
||||
logger.error(f"Qwen上传返回数据格式错误: {upload_res.data}")
|
||||
return False
|
||||
|
||||
logger.info(f"Qwen上传成功: {upload_res.data.fileName}")
|
||||
|
||||
workflow_id = '1980864078929379330'
|
||||
if len(prompt) == 0:
|
||||
prompt = '提取桌布上的花纹，自动补全空白，使得所有位置饱满并且完美衔接，去除所有的皱纹和扭曲和凸凹不平，图案自动摆正对齐并且铺平，使直线变得笔直，平行的花纹更有规律，没有残缺的花纹和折痕和断痕，铺满画布，完整的图案。简单的纯色背景'
|
||||
|
||||
node_info_list = [
|
||||
{
|
||||
'nodeId': '78',
|
||||
'fieldName': 'image',
|
||||
'fieldValue': upload_res.data.fileName,
|
||||
},
|
||||
{
|
||||
'nodeId': '103',
|
||||
'fieldName': 'text',
|
||||
'fieldValue': prompt,
|
||||
},
|
||||
]
|
||||
create_res = await create(workflow_id=workflow_id, node_info_list=node_info_list)
|
||||
|
||||
if create_res.code != 0 or not create_res.data:
|
||||
logger.error(f"Qwen任务创建失败: code={create_res.code}, msg={create_res.msg}")
|
||||
return False
|
||||
|
||||
# 确保 data 是 CreateInfo 类型
|
||||
if not hasattr(create_res.data, 'taskId'):
|
||||
logger.error(f"Qwen任务创建返回数据格式错误: {create_res.data}")
|
||||
return False
|
||||
|
||||
task_id = create_res.data.taskId
|
||||
logger.info(f"Qwen任务创建成功: {task_id}")
|
||||
|
||||
# 轮询检查状态
|
||||
max_retries = 120 # 最多等待10分钟（120次 * 5秒）
|
||||
retry_count = 0
|
||||
|
||||
while retry_count < max_retries:
|
||||
status_res = await status(task_id=task_id)
|
||||
if status_res.code == 0:
|
||||
if status_res.data == 'QUEUED':
|
||||
logger.info('Qwen队列排队中...')
|
||||
elif status_res.data == 'RUNNING':
|
||||
logger.info('Qwen正在处理中...')
|
||||
elif status_res.data == 'FAILED':
|
||||
logger.error(f'Qwen处理失败: {status_res}')
|
||||
return False
|
||||
elif status_res.data == 'SUCCESS':
|
||||
logger.info('Qwen处理完成，开始下载结果')
|
||||
outputs_res = await outputs(task_id=task_id)
|
||||
img_url = outputs_res['data'][0]['fileUrl']
|
||||
|
||||
# 下载结果图片
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(img_url) as resp:
|
||||
img_data = await resp.read()
|
||||
|
||||
with open(save_path, 'wb') as f:
|
||||
f.write(img_data)
|
||||
|
||||
logger.info(f"Qwen结果保存成功: {save_path}")
|
||||
try:
|
||||
from utils.api_cost_tracker import record
|
||||
record("qwen_enhance", count=1)
|
||||
except Exception:
|
||||
pass
|
||||
return True
|
||||
|
||||
await asyncio.sleep(5) # 每5秒检查一次
|
||||
retry_count += 1
|
||||
else:
|
||||
logger.error(f'Qwen处理失败: {status_res}')
|
||||
return False
|
||||
|
||||
logger.error(f"Qwen处理超时，超过{max_retries * 5}秒")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Qwen花纹提取异常: {e}")
|
||||
import traceback
|
||||
logger.error(f"异常堆栈: {traceback.format_exc()}")
|
||||
return False
|
||||
|
||||
|
||||
async def 清晰化_api(img_path: str, save_path: str) -> bool:
|
||||
"""
|
||||
高清增强：对透视矫正后的图案进行清晰化处理。
|
||||
使用与花纹提取相同的 ComfyUI 工作流，但提示词聚焦于清晰度增强。
|
||||
|
||||
Args:
|
||||
img_path: 输入图片路径（透视矫正后的结果）
|
||||
save_path: 输出图片路径
|
||||
|
||||
Returns:
|
||||
bool: 处理是否成功
|
||||
"""
|
||||
prompt = (
|
||||
"对这张已展平的图案进行高清增强处理："
|
||||
"提升整体清晰度和锐利度，修复模糊边缘，补全细节纹理，"
|
||||
"使图案线条清晰笔直，颜色鲜艳均匀，"
|
||||
"去除噪点和压缩痕迹，输出印刷级高质量平面图，"
|
||||
"背景保持纯白色，不要改变图案内容和构图。"
|
||||
)
|
||||
return await 花纹提取_api(img_path=img_path, save_path=save_path, prompt=prompt)
|
||||
|
||||
|
||||
# 测试代码（注释掉）
|
||||
# if __name__ == "__main__":
|
||||
# asyncio.run(花纹提取_api(img_path=r'1.jpg', save_path='save1.png', prompt=''))
|
||||
Reference in New Issue
Block a user