262 lines
10 KiB
Python
262 lines
10 KiB
Python
"""
|
||
生成图片 API - 供 pywebview 前端调用
|
||
流程：上传图片 -> workflow_run -> 轮询 query_result -> 解析返回图片 URL
|
||
"""
|
||
|
||
import base64
|
||
import json
|
||
import os
|
||
import tempfile
|
||
import time
|
||
import re
|
||
import traceback
|
||
|
||
import requests
|
||
|
||
from coze import upload_file, workflow_run, query_result
|
||
from config import workflow_id
|
||
from ali_oss import upload_data_urls as oss_upload_data_urls
|
||
|
||
|
||
def _data_url_to_temp_file(data_url: str, allow_video: bool = False) -> str:
|
||
"""将 base64 data URL 转为临时文件路径"""
|
||
# data:image/png;base64,xxxx 或 data:video/mp4;base64,xxxx
|
||
if allow_video:
|
||
match = re.match(r'data:(?:image|video)/(\w+);base64,(.+)', data_url)
|
||
else:
|
||
match = re.match(r'data:image/(\w+);base64,(.+)', data_url)
|
||
if not match:
|
||
raise ValueError('无效的 data URL 格式')
|
||
mime = match.group(1).lower()
|
||
ext_map = {'png': 'png', 'webp': 'png', 'jpeg': 'jpg', 'jpg': 'jpg', 'mp4': 'mp4', 'webm': 'webm'}
|
||
ext = ext_map.get(mime, 'jpg')
|
||
data = base64.b64decode(match.group(2))
|
||
fd, path = tempfile.mkstemp(suffix=f'.{ext}')
|
||
try:
|
||
os.write(fd, data)
|
||
finally:
|
||
os.close(fd)
|
||
return path
|
||
|
||
|
||
def _upload_video(video_data_url: str) -> str:
|
||
"""上传视频，返回 file_id"""
|
||
if not video_data_url or not isinstance(video_data_url, str):
|
||
raise ValueError('无效的视频数据')
|
||
path = _data_url_to_temp_file(video_data_url, allow_video=True)
|
||
try:
|
||
resp = upload_file(path)
|
||
if resp.get('code') == 0 and resp.get('data', {}).get('id'):
|
||
return resp['data']['id']
|
||
raise RuntimeError(f'视频上传失败: {resp.get("msg", resp)}')
|
||
finally:
|
||
try:
|
||
os.unlink(path)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _upload_images(image_data_urls: list) -> list:
|
||
"""上传多张图片，返回 file_id 列表"""
|
||
file_ids = []
|
||
temp_paths = []
|
||
try:
|
||
for data_url in (image_data_urls or []):
|
||
if not data_url or not isinstance(data_url, str):
|
||
continue
|
||
if data_url.startswith("http"):
|
||
fd, path = tempfile.mkstemp(suffix=f'.png')
|
||
data_content = requests.get(data_url).content
|
||
os.write(fd, data_content)
|
||
else:
|
||
path = _data_url_to_temp_file(data_url)
|
||
temp_paths.append(path)
|
||
resp = upload_file(path)
|
||
if resp.get('code') == 0 and resp.get('data', {}).get('id'):
|
||
file_ids.append(resp['data']['id'])
|
||
else:
|
||
raise RuntimeError(f'上传失败: {resp.get("msg", resp)}')
|
||
return file_ids
|
||
finally:
|
||
for p in temp_paths:
|
||
try:
|
||
os.unlink(p)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _parse_output(output_str: str) -> list:
|
||
"""解析 query_result 中的 output，提取图片 URL 列表"""
|
||
try:
|
||
outer = json.loads(output_str)
|
||
inner_str = outer.get('Output', '{}')
|
||
inner = json.loads(inner_str)
|
||
data_str = inner.get('data', '[]')
|
||
data = json.loads(data_str)
|
||
urls = data.get("images")
|
||
long_image_url = data.get("merged_image_url")
|
||
return urls if isinstance(urls, list) else [],long_image_url
|
||
except Exception:
|
||
return [],None
|
||
|
||
|
||
def _parse_output_text(output_str: str) -> list:
|
||
"""解析 query_result 中的 output，提取图片 URL 列表"""
|
||
try:
|
||
outer = json.loads(output_str)
|
||
inner_str = outer.get('Output', '{}')
|
||
inner = json.loads(inner_str)
|
||
data_str = inner.get('data', '')
|
||
if isinstance(data_str,str):
|
||
data = json.loads(data_str).get("reverse_prompt")
|
||
return [data]
|
||
pattern = r'【?图像 \d+】.*?(?=【?图像 \d+】|$)'
|
||
segments = re.findall(pattern, data_str, re.DOTALL)
|
||
return segments
|
||
except Exception:
|
||
return ""
|
||
|
||
|
||
def _poll_until_done(wf_id: str, execute_id: str, interval: float =10, timeout: int = 600*2, is_text: bool = False) -> list:
|
||
"""轮询直到成功或超时，返回图片 URL 列表"""
|
||
start = time.time()
|
||
while time.time() - start < timeout:
|
||
resp = query_result(wf_id, execute_id)
|
||
if resp.get('code') != 0:
|
||
raise RuntimeError(f'查询失败: {resp.get("msg", resp)}')
|
||
items = resp.get('data') or []
|
||
if not items:
|
||
time.sleep(interval)
|
||
continue
|
||
item = items[0]
|
||
status = item.get('execute_status', '')
|
||
if status == 'Success':
|
||
output = item.get('output', '{}')
|
||
if is_text:
|
||
return _parse_output_text(output)
|
||
return _parse_output(output)
|
||
if status and status not in ('Running', 'Pending', ''):
|
||
raise RuntimeError(f'执行失败: {status}')
|
||
time.sleep(interval)
|
||
raise RuntimeError('生成超时')
|
||
|
||
|
||
def _res_to_2k(res: str) -> str:
|
||
"""统一分辨率为 2K/4K"""
|
||
r = (res or '2k').strip().upper()
|
||
return '4K' if r == '4K' else '2K'
|
||
|
||
|
||
def generate(params: dict) -> dict:
|
||
"""
|
||
生成图片
|
||
params: {
|
||
menu: int, # 1=图片反推 2=图片编辑 3=随机海报 4=克隆海报 5=服饰穿搭
|
||
prompt: str, # 用户自定义指令 (menu=1 必填)
|
||
ref_images: list, # base64 data URLs - 参考图
|
||
video: str, # base64 data URL - 视频 (menu=1 可选，仅支持1个)
|
||
model_images: list, # base64 data URLs - 多模特图 (menu!=1，最多5张)
|
||
name: str, desc: str, ratio: str, resolution: str, count: int,
|
||
language: str, style: str, batch_prompt: list, brand_name: str,
|
||
Ingredients: str, activity: str, mode: str, texts: list,
|
||
proc_images: list, layout_image: str,
|
||
}
|
||
返回: { success: bool, urls: list, prompts: list, error: str }
|
||
"""
|
||
original_urls = []
|
||
api_key = (params.get('api_key') or '').strip()
|
||
try:
|
||
menu = int(params.get('menu', 2))
|
||
ref_data = params.get('ref_images') or []
|
||
proc_data = params.get('proc_images') or []
|
||
layout_data = params.get('layout_image')
|
||
video_data = params.get('video')
|
||
model_data = params.get('model_images') or []
|
||
|
||
# menu 1: 反推词，简化参数
|
||
if menu == 1:
|
||
base_params = {'menu': 1, 'prompt': str(params.get('prompt', '')).strip() or ''}
|
||
|
||
ref_ids = _upload_images(ref_data) if ref_data else []
|
||
if ref_ids:
|
||
base_params['ref_images'] = [{'file_id': fid} for fid in ref_ids]
|
||
if video_data:
|
||
try:
|
||
vid = _upload_video(video_data)
|
||
base_params['video'] = {'file_id': vid}
|
||
except Exception as ve:
|
||
return {'success': False, 'urls': [], 'prompts': [], 'error': f'视频上传失败: {ve}'}
|
||
base_params = {k: v for k, v in base_params.items() if v is not None and v != ''}
|
||
base_params["api_key"] = api_key
|
||
resp = workflow_run(workflow_id, base_params)
|
||
if resp.get('code') != 0:
|
||
return {'success': False, 'urls': [], 'prompts': [], 'error': resp.get('msg', str(resp))}
|
||
execute_id = resp.get('execute_id')
|
||
if not execute_id:
|
||
return {'success': False, 'urls': [], 'prompts': [], 'error': '未返回 execute_id'}
|
||
prompts = _poll_until_done(workflow_id, str(execute_id), is_text=True)
|
||
# menu 1 返回提示词列表，统一转为字符串
|
||
# prompts = [str(x) for x in result_list] if result_list else []
|
||
# prompts = [prompts]
|
||
return {'success': True, 'urls': [], 'prompts': prompts, 'error': ''}
|
||
|
||
# menu != 1: 原有逻辑
|
||
all_ref = list(ref_data)
|
||
if layout_data:
|
||
all_ref = [layout_data] + list(ref_data)
|
||
all_originals = list(all_ref) + list(proc_data) + list(model_data)
|
||
|
||
if all_originals:
|
||
try:
|
||
original_urls = oss_upload_data_urls(all_originals, prefix="originals")
|
||
except Exception:
|
||
traceback.print_exc()
|
||
pass
|
||
|
||
ref_ids = _upload_images(all_ref) if all_ref else []
|
||
proc_ids = _upload_images(proc_data) if proc_data else []
|
||
model_ids = _upload_images(model_data) if model_data else []
|
||
|
||
res = _res_to_2k(params.get('resolution', '2K'))
|
||
|
||
base_params = {
|
||
'name': str(params.get('name', '')).strip(),
|
||
'ratio': str(params.get('ratio', '')).strip(),
|
||
'menu': menu,
|
||
'resolution': res,
|
||
'count': int(params.get('count', 1)),
|
||
'desc': str(params.get('desc', '')).strip(),
|
||
'language': str(params.get('language', '中文')).strip() or '中文',
|
||
'style': str(params.get('style', '')).strip(),
|
||
'prompt': str(params.get('prompt', '')).strip(),
|
||
'batch_prompt': params.get('batch_prompt') or [],
|
||
'brand_name': str(params.get('brand_name', '')).strip(),
|
||
'Ingredients': str(params.get('Ingredients', '')).strip(),
|
||
'activity': str(params.get('activity', '')).strip(),
|
||
'mode': str(params.get('mode', '1')).strip() or '1',
|
||
'texts': params.get('text') or [],
|
||
"api_key" : api_key
|
||
}
|
||
|
||
if ref_ids:
|
||
base_params['ref_images'] = [{'file_id': fid} for fid in ref_ids]
|
||
if proc_ids:
|
||
base_params['proc_images'] = [{'file_id': fid} for fid in proc_ids]
|
||
if model_ids:
|
||
base_params['model_images'] = [{'file_id': fid} for fid in model_ids]
|
||
base_params = {k: v for k, v in base_params.items() if v is not None and v != ''}
|
||
|
||
resp = workflow_run(workflow_id, base_params)
|
||
if resp.get('code') != 0:
|
||
return {'success': False, 'urls': [], 'error': resp.get('msg', str(resp))}
|
||
execute_id = resp.get('execute_id')
|
||
if not execute_id:
|
||
return {'success': False, 'urls': [], 'error': '未返回 execute_id'}
|
||
|
||
urls,long_image_url = _poll_until_done(workflow_id, str(execute_id))
|
||
return {'success': True, 'urls': urls, 'original_urls': original_urls, 'error': '',"long_image_url":long_image_url}
|
||
except Exception as e:
|
||
return {'success': False, 'urls': [], 'prompts': [], 'original_urls': [], 'error': str(e),"long_image_url":""}
|
||
|
||
|