打包项目

2026-03-20 18:13:47 +08:00
parent 7c6d9dd989
commit 63d3e4047c
59 changed files with 1872 additions and 4187 deletions
--- a/source_code/generate_api.py
+++ b/source_code/generate_api.py
@@ -0,0 +1,261 @@
+"""
+生成图片 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":""}
+
+