234 lines
8.6 KiB
Python
234 lines
8.6 KiB
Python
"""
|
||
通义千问文生图（DashScope 多模态生成 API）。
|
||
默认模型：qwen-image-2.0-pro（与阿里云 Model Studio 文档一致）。
|
||
文档：https://www.alibabacloud.com/help/en/model-studio/qwen-image-api
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import re
|
||
from typing import Any, Optional
|
||
|
||
import requests
|
||
|
||
import config
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 北京地域同步接口（与 QWEN_BASE_URL 常用配置同账号）
|
||
DEFAULT_MULTIMODAL_BASE = 'https://dashscope.aliyuncs.com/api/v1'
|
||
MULTIMODAL_GENERATION_PATH = '/services/aigc/multimodal-generation/generation'
|
||
|
||
_NEG_DEFAULT = (
|
||
'低分辨率, 模糊, 畸形手指, 过度饱和, 蜡像感, 杂乱构图, 扭曲文字, '
|
||
'公司商标, 企业LOGO, 投标人名称, 水印, 证件照, '
|
||
'塔吊, 塔式起重机, 门式起重机, 履带吊, 汽车吊, 起重机, 吊车, 挖掘机, 挖机, 装载机, 铲车, '
|
||
'压路机, 推土机, 平地机, 摊铺机, 泵车, 混凝土泵车, 搅拌车, 罐车, 自卸车, 渣土车, '
|
||
'施工机械, 工程车辆, 重型机械, 桩机, 钻机, crane, excavator, bulldozer, tower crane'
|
||
)
|
||
|
||
_PLACEHOLDER_KEY_RE = re.compile(r'^sk-your|^sk-xxxx', re.I)
|
||
|
||
|
||
def _effective_multimodal_base() -> str:
|
||
raw = (getattr(config, 'QWEN_MULTIMODAL_BASE', '') or '').strip().rstrip('/')
|
||
if raw:
|
||
return raw
|
||
return DEFAULT_MULTIMODAL_BASE
|
||
|
||
|
||
def _generation_url() -> str:
|
||
return f'{_effective_multimodal_base()}{MULTIMODAL_GENERATION_PATH}'
|
||
|
||
|
||
def _qwen_api_key() -> str:
|
||
key = (getattr(config, 'QWEN_API_KEY', '') or '').strip()
|
||
if not key or _PLACEHOLDER_KEY_RE.match(key):
|
||
return ''
|
||
return key
|
||
|
||
|
||
def _extract_image_url(payload: dict[str, Any]) -> Optional[str]:
|
||
try:
|
||
out = payload.get('output') or {}
|
||
choices = out.get('choices') or []
|
||
if not choices:
|
||
return None
|
||
msg = (choices[0] or {}).get('message') or {}
|
||
content = msg.get('content')
|
||
if isinstance(content, list):
|
||
for part in content:
|
||
if isinstance(part, dict):
|
||
url = part.get('image')
|
||
if isinstance(url, str) and url.startswith('http'):
|
||
return url
|
||
if isinstance(content, dict):
|
||
url = content.get('image')
|
||
if isinstance(url, str) and url.startswith('http'):
|
||
return url
|
||
except Exception:
|
||
return None
|
||
return None
|
||
|
||
|
||
def _truncate_prompt(text: str, max_chars: int = 780) -> str:
|
||
t = (text or '').strip().replace('\r\n', '\n').replace('\r', '\n')
|
||
if len(t) <= max_chars:
|
||
return t
|
||
return t[: max_chars - 3] + '...'
|
||
|
||
|
||
def build_attachment_figure_prompt(title: str, body: str, *, grayscale: bool = False) -> str:
|
||
"""投标附件用图：专业工程示意风格；图中禁止任何施工机械与塔吊。"""
|
||
core = _truncate_prompt(f'{title}\n{body}', 680)
|
||
gray_clause = ''
|
||
if grayscale:
|
||
gray_clause = (
|
||
'【暗标·强制】整幅插图必须为黑白灰单色风格：仅使用黑、白及灰色阶，不得出现任何彩色、'
|
||
'饱和度色块、彩色箭头或彩色背景；分区示意仅用灰度填充区分。'
|
||
)
|
||
head = (
|
||
'建筑施工投标技术标附件插图：清晰线稿或浅色块示意，横版；可含简短中文标注。'
|
||
'禁止公司名、商标、LOGO、投标人信息。'
|
||
'【强制】画面中不得出现塔吊、门吊、汽车吊、履带吊等任何起重设备，不得出现挖掘机、装载机、压路机、'
|
||
'泵车、搅拌车、推土机等一切施工机械与工程车辆（含剪影、线稿、远影、图标）；仅允许建筑轮廓、'
|
||
'临建分区色块、道路与大门示意、流程框与箭头。'
|
||
)
|
||
if gray_clause:
|
||
head = gray_clause + '\n' + head
|
||
sched_keys = ('进度', '网络图', '横道', '开工', '完工', '工期', '里程碑')
|
||
site_keys = ('总平面', '平面布置', '布置图')
|
||
t = f'{title}\n{body}'
|
||
extra = ''
|
||
if any(k in t for k in sched_keys):
|
||
extra += (
|
||
'【进度图】仅表现文字中已给出的开工完工逻辑、清单节点与工期关系；'
|
||
'不得添加正文中未出现的日历日、工序名称或里程碑；无具体日期时不写具体年月日数字；'
|
||
'禁止画机械或塔吊。'
|
||
)
|
||
if any(k in t for k in site_keys):
|
||
extra += (
|
||
'【总平面图】临时设施分区与面积须与文字中的临时用地表一致；'
|
||
'仅绘制分区块面、道路与出入口关系，严禁塔吊与任何施工机械形象。'
|
||
)
|
||
extra = _truncate_prompt(extra, 220) if extra else ''
|
||
parts = [head]
|
||
if extra:
|
||
parts.append(extra)
|
||
parts.append('内容要点（按文字布局即可，不要画机械）：\n' + core)
|
||
return '\n'.join(parts)
|
||
|
||
|
||
def generate_qwen_image_png(
|
||
prompt: str,
|
||
*,
|
||
size: Optional[str] = None,
|
||
timeout: Optional[int] = None,
|
||
) -> tuple[Optional[bytes], Optional[str]]:
|
||
"""
|
||
调用 qwen-image-2.0-pro 同步文生图，返回 PNG 字节。
|
||
失败时 (None, error_message)。
|
||
"""
|
||
api_key = _qwen_api_key()
|
||
if not api_key:
|
||
return None, '未配置有效的通义千问 API Key（文生图与文本模型共用 DashScope Key）'
|
||
|
||
model = getattr(config, 'QWEN_IMAGE_MODEL', 'qwen-image-2.0-pro') or 'qwen-image-2.0-pro'
|
||
sz = (size or getattr(config, 'QWEN_IMAGE_SIZE', '1536*1024') or '1536*1024').strip()
|
||
prompt_clean = _truncate_prompt(prompt, 800)
|
||
url = _generation_url()
|
||
|
||
body = {
|
||
'model': model,
|
||
'input': {
|
||
'messages': [
|
||
{
|
||
'role': 'user',
|
||
'content': [{'text': prompt_clean}],
|
||
}
|
||
]
|
||
},
|
||
'parameters': {
|
||
'negative_prompt': getattr(config, 'QWEN_IMAGE_NEGATIVE_PROMPT', _NEG_DEFAULT),
|
||
'prompt_extend': bool(getattr(config, 'QWEN_IMAGE_PROMPT_EXTEND', True)),
|
||
'watermark': bool(getattr(config, 'QWEN_IMAGE_WATERMARK', False)),
|
||
'size': sz,
|
||
},
|
||
}
|
||
|
||
to = timeout if timeout is not None else max(120, int(getattr(config, 'REQUEST_TIMEOUT', 180) or 180))
|
||
try:
|
||
resp = requests.post(
|
||
url,
|
||
headers={
|
||
'Authorization': f'Bearer {api_key}',
|
||
'Content-Type': 'application/json',
|
||
},
|
||
data=json.dumps(body, ensure_ascii=False).encode('utf-8'),
|
||
timeout=to,
|
||
)
|
||
except requests.RequestException as e:
|
||
logger.warning('qwen-image 请求异常: %s', e)
|
||
return None, str(e)
|
||
|
||
try:
|
||
data = resp.json()
|
||
except Exception:
|
||
return None, f'文生图接口返回非 JSON，HTTP {resp.status_code}'
|
||
|
||
if resp.status_code != 200:
|
||
msg = data.get('message') or data.get('code') or resp.text[:500]
|
||
logger.warning('qwen-image HTTP %s: %s', resp.status_code, msg)
|
||
return None, f'文生图失败 HTTP {resp.status_code}: {msg}'
|
||
|
||
code = data.get('code')
|
||
if code and str(code).upper() not in ('OK', 'SUCCESS', '200', ''):
|
||
msg = data.get('message') or str(code)
|
||
logger.warning('qwen-image 业务错误: %s', msg)
|
||
return None, f'文生图接口错误: {msg}'
|
||
|
||
img_url = _extract_image_url(data)
|
||
if not img_url:
|
||
logger.warning('qwen-image 响应无图片 URL: %s', str(data)[:800])
|
||
return None, '文生图响应中未找到图片 URL'
|
||
|
||
try:
|
||
ir = requests.get(img_url, timeout=min(to, 120))
|
||
ir.raise_for_status()
|
||
return ir.content, None
|
||
except requests.RequestException as e:
|
||
logger.warning('下载生成图片失败: %s', e)
|
||
return None, f'下载图片失败: {e}'
|
||
|
||
|
||
def _png_bytes_to_grayscale(png_bytes: bytes) -> bytes:
|
||
"""将 PNG 转为灰度 RGB（暗标附图不得含彩色）。"""
|
||
try:
|
||
from io import BytesIO
|
||
|
||
from PIL import Image
|
||
|
||
im = Image.open(BytesIO(png_bytes)).convert('L').convert('RGB')
|
||
out = BytesIO()
|
||
im.save(out, format='PNG')
|
||
return out.getvalue()
|
||
except Exception:
|
||
return png_bytes
|
||
|
||
|
||
def generate_attachment_figure_png(
|
||
title: str,
|
||
content: str,
|
||
*,
|
||
grayscale: bool = False,
|
||
) -> tuple[Optional[bytes], Optional[str]]:
|
||
"""附件 [FIGURE] 专用：拼装提示词并生成 PNG。"""
|
||
prompt = build_attachment_figure_prompt(
|
||
title or '附图', content or '', grayscale=grayscale
|
||
)
|
||
png, err = generate_qwen_image_png(prompt)
|
||
if png and grayscale:
|
||
png = _png_bytes_to_grayscale(png)
|
||
return png, err
|