413 lines
16 KiB
Python
413 lines
16 KiB
Python
"""
|
||
图片生成蓝图：生成、历史、下载、拼接、版本
|
||
"""
|
||
import os
|
||
import sys
|
||
import json
|
||
import re
|
||
import io
|
||
import base64
|
||
import tempfile
|
||
import threading
|
||
import subprocess
|
||
import requests
|
||
from urllib.parse import urlparse, quote
|
||
|
||
from flask import Blueprint, request, jsonify, Response, send_file
|
||
from PIL import Image
|
||
|
||
from app_common import get_db, login_required, current_user_id, BASE_DIR
|
||
from config import STITCH_WORKFLOW_ID,client_name
|
||
|
||
image_bp = Blueprint('image', __name__)
|
||
|
||
|
||
def _load_image_from_url_or_data(url_or_data):
|
||
"""从 http(s) URL 或 data URL 加载为 PIL Image，失败返回 None"""
|
||
if not url_or_data or not isinstance(url_or_data, str):
|
||
return None
|
||
try:
|
||
if url_or_data.startswith('data:'):
|
||
m = re.match(r'data:image/[^;]+;base64,(.+)', url_or_data, re.DOTALL)
|
||
if not m:
|
||
return None
|
||
raw = base64.b64decode(m.group(1).strip())
|
||
img = Image.open(io.BytesIO(raw))
|
||
elif url_or_data.startswith(('http://', 'https://')):
|
||
resp = requests.get(url_or_data, timeout=15)
|
||
resp.raise_for_status()
|
||
img = Image.open(io.BytesIO(resp.content))
|
||
else:
|
||
return None
|
||
if img.mode != 'RGB':
|
||
img = img.convert('RGB')
|
||
return img
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _stitch_and_upload_long_image(urls):
|
||
"""将多张图片 URL 先上传获取 file_id，再调用 workflow_run 拼接长图，返回 data.merged_image_url；失败返回 None。"""
|
||
if not urls or not isinstance(urls, (list, tuple)):
|
||
return None
|
||
from coze import upload_file as coze_upload_file, workflow_run
|
||
file_ids = []
|
||
temp_paths = []
|
||
try:
|
||
for u in urls:
|
||
img = _load_image_from_url_or_data(u)
|
||
if img is None:
|
||
continue
|
||
fd, path = tempfile.mkstemp(suffix='.png')
|
||
try:
|
||
os.close(fd)
|
||
img.save(path)
|
||
temp_paths.append(path)
|
||
resp = coze_upload_file(path)
|
||
if resp.get('code') == 0 and resp.get('data', {}).get('id'):
|
||
file_ids.append(resp['data']['id'])
|
||
except Exception:
|
||
pass
|
||
if not file_ids:
|
||
return None
|
||
parameters = {"images": [{"file_id": fid} for fid in file_ids]}
|
||
resp = workflow_run(STITCH_WORKFLOW_ID, parameters, is_async=False)
|
||
if resp.get('code') != 0:
|
||
return None
|
||
data = resp.get('data') or {}
|
||
data = json.loads(data)
|
||
merged_image_url = data.get('merged_image_url')
|
||
if merged_image_url:
|
||
return merged_image_url
|
||
output_str = data.get('output') or ''
|
||
if output_str:
|
||
try:
|
||
outer = json.loads(output_str)
|
||
inner_str = outer.get('Output', '{}')
|
||
inner = json.loads(inner_str)
|
||
data_str = inner.get('data', '[]')
|
||
inner_data = json.loads(data_str)
|
||
merged_image_url = inner_data.get('merged_image_url')
|
||
return merged_image_url
|
||
except Exception:
|
||
pass
|
||
return None
|
||
except Exception:
|
||
return None
|
||
finally:
|
||
for p in temp_paths:
|
||
try:
|
||
os.unlink(p)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _sanitize_params_for_history(params):
|
||
"""移除 base64 大字段及敏感字段，仅保留可存储的请求参数"""
|
||
exclude = ('ref_images', 'proc_images', 'layout_image')
|
||
out = {}
|
||
for k, v in (params or {}).items():
|
||
if k == 'api_key':
|
||
continue
|
||
if k in exclude:
|
||
if isinstance(v, list):
|
||
out[f'{k}_count'] = len(v)
|
||
else:
|
||
out[f'{k}_count'] = 1 if v else 0
|
||
elif isinstance(v, (str, int, float, bool, type(None))):
|
||
out[k] = v
|
||
elif isinstance(v, list) and not v:
|
||
out[k] = []
|
||
elif isinstance(v, list) and isinstance(v[0], (str, int, float, bool)):
|
||
out[k] = v
|
||
else:
|
||
out[k] = str(v)[:200] if v else None
|
||
return out
|
||
|
||
|
||
@image_bp.route('/api/generate', methods=['POST'])
|
||
@login_required
|
||
def api_generate():
|
||
"""生成图片：调用 generate_api，上传原图到 OSS，保存历史记录"""
|
||
try:
|
||
params = request.get_json() or {}
|
||
from generate_api import generate
|
||
result = generate(params)
|
||
if result.get('success') and result.get('urls'):
|
||
long_image_url = result.get("long_image_url")
|
||
result["long_image_url"] = long_image_url
|
||
import json as _json
|
||
history_id = None
|
||
try:
|
||
hid = params.get('history_id')
|
||
if hid is not None:
|
||
try:
|
||
hid = int(hid)
|
||
except (TypeError, ValueError):
|
||
hid = None
|
||
conn = get_db()
|
||
_uid = current_user_id()
|
||
with conn.cursor() as cur:
|
||
if hid is not None and hid > 0:
|
||
cur.execute(
|
||
"SELECT result_urls FROM image_history WHERE id=%s AND user_id=%s",
|
||
(hid, _uid),
|
||
)
|
||
row = cur.fetchone()
|
||
existing_urls = []
|
||
if row and row.get('result_urls'):
|
||
try:
|
||
existing_urls = _json.loads(row['result_urls'])
|
||
except Exception:
|
||
existing_urls = []
|
||
new_urls = result.get('urls') or []
|
||
new_url = new_urls[0] if new_urls else None
|
||
idx = params.get('history_index', 0)
|
||
try:
|
||
idx = int(idx)
|
||
except (TypeError, ValueError):
|
||
idx = 0
|
||
if new_url:
|
||
if not isinstance(existing_urls, list):
|
||
existing_urls = []
|
||
while len(existing_urls) <= idx:
|
||
existing_urls.append(existing_urls[-1] if existing_urls else new_url)
|
||
existing_urls[idx] = new_url
|
||
merged_result_urls = existing_urls or new_urls
|
||
cur.execute(
|
||
"""UPDATE image_history SET panel_type=%s, original_urls=%s, params=%s, result_urls=%s
|
||
WHERE id=%s AND user_id=%s""",
|
||
(
|
||
params.get('panel_type', ''),
|
||
_json.dumps(result.get('original_urls') or []),
|
||
_json.dumps(_sanitize_params_for_history(params)),
|
||
_json.dumps(merged_result_urls),
|
||
hid,
|
||
_uid,
|
||
),
|
||
)
|
||
if cur.rowcount > 0:
|
||
history_id = hid
|
||
else:
|
||
cur.execute(
|
||
"""INSERT INTO image_history (user_id, panel_type, original_urls, params, result_urls, long_image_url)
|
||
VALUES (%s, %s, %s, %s, %s, %s)""",
|
||
(
|
||
_uid,
|
||
params.get('panel_type', ''),
|
||
_json.dumps(result.get('original_urls') or []),
|
||
_json.dumps(_sanitize_params_for_history(params)),
|
||
_json.dumps(result.get('urls') or []),
|
||
long_image_url,
|
||
),
|
||
)
|
||
history_id = cur.lastrowid
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
if history_id is not None:
|
||
result['history_id'] = history_id
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
import traceback
|
||
traceback.print_exc()
|
||
return jsonify({'success': False, 'urls': [], 'error': str(e)})
|
||
|
||
|
||
@image_bp.route('/api/version')
|
||
def api_version():
|
||
"""检测更新：请求 APP_UPDATE_URL 获取最新版本信息，与当前版本比较"""
|
||
current_version = (os.environ.get('APP_VERSION', '1.0.0') or '1.0.0').strip()
|
||
update_url = (os.environ.get('APP_UPDATE_URL', '') or '').strip()
|
||
result = {
|
||
'version': current_version,
|
||
'desc': '',
|
||
'url': '',
|
||
'has_update': False,
|
||
'latest_version': current_version,
|
||
'file_url': '',
|
||
}
|
||
if not update_url:
|
||
return jsonify(result)
|
||
try:
|
||
resp = requests.get(update_url, timeout=10)
|
||
resp.raise_for_status()
|
||
data = resp.json() or {}
|
||
latest_version = (data.get('version') or '').strip()
|
||
file_url = (data.get('file_url') or '').strip()
|
||
result['latest_version'] = latest_version
|
||
result['file_url'] = file_url
|
||
result['url'] = file_url
|
||
# 版本不一致则视为有更新
|
||
if latest_version and latest_version != current_version:
|
||
result['has_update'] = True
|
||
except Exception:
|
||
pass
|
||
return jsonify(result)
|
||
|
||
|
||
def _run_update_and_exit(zip_path, target_dir):
|
||
"""在后台延迟后启动 update.exe（脱离当前进程），然后退出当前程序"""
|
||
def _do():
|
||
import time
|
||
time.sleep(1.5) # 确保 HTTP 响应已发送
|
||
# exe_dir = target_dir
|
||
# exe_dir = "D:\\pack\\nanri\\update_exe\\main.dist"
|
||
exe_dir = os.path.join(BASE_DIR,"update")
|
||
update_exe = os.path.join(exe_dir, 'update.exe')
|
||
if not os.path.isfile(update_exe):
|
||
return
|
||
try:
|
||
creationflags = 0
|
||
if sys.platform == 'win32':
|
||
creationflags = subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP
|
||
subprocess.Popen(
|
||
[update_exe, '--zip', zip_path, '--target', target_dir, '--process',client_name],
|
||
cwd=exe_dir,
|
||
creationflags=creationflags,
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.DEVNULL,
|
||
close_fds=True,
|
||
)
|
||
except Exception:
|
||
pass
|
||
os._exit(0)
|
||
t = threading.Thread(target=_do, daemon=False)
|
||
t.start()
|
||
|
||
|
||
@image_bp.route('/api/update/do', methods=['POST'])
|
||
def api_update_do():
|
||
"""执行更新：下载 zip 到 tmp，启动 update.exe 后退出程序"""
|
||
data = request.get_json() or {}
|
||
file_url = (data.get('file_url') or '').strip()
|
||
if not file_url:
|
||
return jsonify({'success': False, 'error': '缺少 file_url'}), 400
|
||
parsed = urlparse(file_url)
|
||
if parsed.scheme not in ('http', 'https'):
|
||
return jsonify({'success': False, 'error': '无效的下载地址'}), 400
|
||
tmp_dir = os.path.join(BASE_DIR, 'tmp')
|
||
try:
|
||
os.makedirs(tmp_dir, exist_ok=True)
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': f'创建 tmp 目录失败: {e}'}), 500
|
||
# 使用 URL 中的文件名或默认版本名
|
||
filename = os.path.basename(parsed.path) or 'update.zip'
|
||
zip_path = os.path.join(tmp_dir, filename)
|
||
try:
|
||
resp = requests.get(file_url, timeout=300, stream=True)
|
||
resp.raise_for_status()
|
||
with open(zip_path, 'wb') as f:
|
||
for chunk in resp.iter_content(chunk_size=65536):
|
||
if chunk:
|
||
f.write(chunk)
|
||
except requests.RequestException as e:
|
||
return jsonify({'success': False, 'error': f'下载失败: {e}'}), 502
|
||
_run_update_and_exit(zip_path, BASE_DIR)
|
||
return jsonify({'success': True, 'message': '更新已启动，程序即将退出'})
|
||
|
||
|
||
@image_bp.route('/api/download')
|
||
@login_required
|
||
def api_download():
|
||
"""代理下载图片，解决跨域 fetch 无法下载的问题"""
|
||
url = request.args.get('url', '').strip()
|
||
filename = request.args.get('filename', 'image.png')
|
||
if not url:
|
||
return jsonify({'success': False, 'error': '缺少 url 参数'}), 400
|
||
parsed = urlparse(url)
|
||
if parsed.scheme not in ('http', 'https'):
|
||
return jsonify({'success': False, 'error': '仅支持 http/https 链接'}), 400
|
||
try:
|
||
resp = requests.get(url, timeout=30, stream=True)
|
||
resp.raise_for_status()
|
||
content_type = resp.headers.get('Content-Type', 'image/png')
|
||
encoded = quote(filename, safe='')
|
||
disposition = f"attachment; filename*=UTF-8''{encoded}"
|
||
return Response(
|
||
resp.iter_content(chunk_size=8192),
|
||
mimetype=content_type,
|
||
headers={'Content-Disposition': disposition}
|
||
)
|
||
except requests.RequestException as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 502
|
||
|
||
|
||
@image_bp.route('/api/stitch/save', methods=['POST'])
|
||
@login_required
|
||
def api_stitch_save():
|
||
"""手动拼接：接收图片 URL 列表（支持 http 或 data URL），拼接并上传，返回长图 URL"""
|
||
try:
|
||
data = request.get_json() or {}
|
||
urls = data.get('urls')
|
||
if not urls or not isinstance(urls, list):
|
||
return jsonify({'success': False, 'error': '请提供 urls 数组'}), 400
|
||
urls = [u for u in urls if u and isinstance(u, str)]
|
||
if not urls:
|
||
return jsonify({'success': False, 'error': '没有有效的图片'}), 400
|
||
long_image_url = _stitch_and_upload_long_image(urls)
|
||
if not long_image_url:
|
||
return jsonify({'success': False, 'error': '拼接或上传失败'}), 500
|
||
return jsonify({'success': True, 'long_image_url': long_image_url})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@image_bp.route('/api/history')
|
||
@login_required
|
||
def api_history():
|
||
"""分页获取当前用户的历史图库，支持按 panel_type 栏目筛选"""
|
||
page = max(1, int(request.args.get('page', 1)))
|
||
page_size = min(50, max(10, int(request.args.get('page_size', 20))))
|
||
panel_type = (request.args.get('panel_type') or '').strip()
|
||
offset = (page - 1) * page_size
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
where_user = "user_id = %s"
|
||
params_where = [current_user_id()]
|
||
if panel_type:
|
||
where_user += " AND panel_type = %s"
|
||
params_where.append(panel_type)
|
||
cur.execute(
|
||
"""SELECT id, created_at, panel_type, original_urls, params, result_urls, long_image_url
|
||
FROM image_history WHERE """ + where_user + """ ORDER BY created_at DESC LIMIT %s OFFSET %s""",
|
||
params_where + [page_size, offset],
|
||
)
|
||
rows = cur.fetchall()
|
||
cur.execute("SELECT COUNT(*) as total FROM image_history WHERE " + where_user, params_where)
|
||
total = cur.fetchone()['total']
|
||
conn.close()
|
||
|
||
def _parse_json(val, default=None):
|
||
if val is None:
|
||
return default if default is not None else []
|
||
if isinstance(val, (list, dict)):
|
||
return val
|
||
try:
|
||
return json.loads(val)
|
||
except Exception:
|
||
return default if default is not None else []
|
||
|
||
items = []
|
||
for r in rows:
|
||
items.append({
|
||
'id': r['id'],
|
||
'created_at': r['created_at'].strftime('%Y-%m-%d %H:%M') if r['created_at'] else '',
|
||
'panel_type': r['panel_type'] or '',
|
||
'original_urls': _parse_json(r['original_urls'], []),
|
||
'params': _parse_json(r['params'], {}),
|
||
'result_urls': _parse_json(r['result_urls'], []),
|
||
'long_image_url': (r.get('long_image_url') or '').strip() or None,
|
||
})
|
||
return jsonify({'success': True, 'items': items, 'total': total, 'page': page, 'page_size': page_size})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)})
|
||
|
||
|
||
|
||
|
||
|