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=''))