889afecde6
- 新增脚本级默认超时常量，统一使用 180 秒\n- 修改 --timeout 默认值为 180 秒\n- 调整图片后端 timeout_seconds 的脚本兜底值为 180 秒\n- 补充中文注释，说明超时策略与行为
618 lines
22 KiB
Python
618 lines
22 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
批量生成 Dota2 英雄抖音图片脚本。
|
||
|
||
设计目标：
|
||
1. 这是一个一次性本地脚本，不依赖机器人运行时，直接走 HTTP 调用 OpenAI 兼容图片接口。
|
||
2. 脚本会自动从 OpenDota 拉取全部英雄列表，避免手工维护英雄名称。
|
||
3. 每个英雄默认生成 4 张图片，其中 2 张中文排版、2 张日文排版，统一使用相同的画面结构与风格模板，尽量保证成片风格一致。
|
||
4. 脚本支持断点续跑：如果目标文件已经存在，则默认跳过，避免重复计费。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import base64
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
import requests
|
||
import yaml
|
||
|
||
|
||
# 这里固定使用 OpenDota 的公开英雄接口，避免在脚本里硬编码整份英雄名单。
|
||
OPENDOTA_HERO_STATS_URL = "https://api.opendota.com/api/heroStats"
|
||
|
||
# 这里给出一个稳定的 9:16 尺寸，尽量兼容常见 OpenAI 兼容图片网关。
|
||
DEFAULT_IMAGE_SIZE = "1024x1792"
|
||
|
||
# 这里统一定义输出根目录，方便后续在一个目录里筛图、剪辑、上传抖音。
|
||
DEFAULT_OUTPUT_DIR = Path("temp") / "dota2_douyin_images"
|
||
|
||
# 这里固定每个英雄的 4 张图的语言分布，避免每次运行时还要手动指定。
|
||
DEFAULT_LANGUAGE_VARIANTS = ["zh", "zh", "ja", "ja"]
|
||
|
||
# 这里把默认并发数固定为 4，满足你“开 4 个线程跑”的诉求。
|
||
DEFAULT_MAX_WORKERS = 4
|
||
|
||
# 这里把图片生成默认超时统一收敛到 180 秒：
|
||
# 1. 300 秒过长时，失败反馈会比较慢，不利于批量任务及时发现问题；
|
||
# 2. 60 秒又可能不足以覆盖部分高峰期图片生成耗时；
|
||
# 3. 180 秒作为脚本级默认值，更适合当前这条 Dota2 抖音批量生图链路。
|
||
DEFAULT_REQUEST_TIMEOUT_SECONDS = 180
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
"""解析命令行参数。"""
|
||
parser = argparse.ArgumentParser(
|
||
description="批量为 Dota2 全英雄生成抖音竖版图片。"
|
||
)
|
||
parser.add_argument(
|
||
"--config",
|
||
default="config.yaml",
|
||
help="项目根目录下的配置文件路径，默认读取 config.yaml。",
|
||
)
|
||
parser.add_argument(
|
||
"--scene",
|
||
default="image.generate",
|
||
help="LLM 场景名，默认使用 config.yaml 中的 image.generate。",
|
||
)
|
||
parser.add_argument(
|
||
"--output-dir",
|
||
default=str(DEFAULT_OUTPUT_DIR),
|
||
help="图片输出目录。",
|
||
)
|
||
parser.add_argument(
|
||
"--count-per-hero",
|
||
type=int,
|
||
default=4,
|
||
help="每个英雄生成的图片数量，默认 4 张（2 张中文、2 张日文）。",
|
||
)
|
||
parser.add_argument(
|
||
"--size",
|
||
default=DEFAULT_IMAGE_SIZE,
|
||
help="图片尺寸，默认 1024x1792（9:16）。",
|
||
)
|
||
parser.add_argument(
|
||
"--quality",
|
||
default="high",
|
||
help="图片质量参数，默认 high。",
|
||
)
|
||
parser.add_argument(
|
||
"--timeout",
|
||
type=int,
|
||
default=DEFAULT_REQUEST_TIMEOUT_SECONDS,
|
||
help="单次请求超时时间（秒），默认 180 秒。",
|
||
)
|
||
parser.add_argument(
|
||
"--delay",
|
||
type=float,
|
||
default=1.5,
|
||
help="每次成功生成后的等待时间（秒），默认 1.5 秒，避免打满网关。",
|
||
)
|
||
parser.add_argument(
|
||
"--max-retries",
|
||
type=int,
|
||
default=3,
|
||
help="单张图片失败后的最大重试次数，默认 3 次。",
|
||
)
|
||
parser.add_argument(
|
||
"--hero-limit",
|
||
type=int,
|
||
default=0,
|
||
help="仅生成前 N 个英雄，0 表示全部生成，便于先小范围试跑。",
|
||
)
|
||
parser.add_argument(
|
||
"--hero-filter",
|
||
default="",
|
||
help="只生成英雄名中包含该关键字的英雄，便于单独补图。",
|
||
)
|
||
parser.add_argument(
|
||
"--max-workers",
|
||
type=int,
|
||
default=DEFAULT_MAX_WORKERS,
|
||
help="并发线程数，默认 4。",
|
||
)
|
||
parser.add_argument(
|
||
"--force",
|
||
action="store_true",
|
||
help="即使目标文件已存在，也强制重新生成。",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def load_yaml_config(config_path: str) -> Dict[str, Any]:
|
||
"""读取 YAML 配置文件。"""
|
||
with open(config_path, "r", encoding="utf-8") as file_obj:
|
||
return yaml.safe_load(file_obj) or {}
|
||
|
||
|
||
def resolve_image_backend(config_data: Dict[str, Any], scene_name: str) -> Dict[str, Any]:
|
||
"""
|
||
根据 scene 解析图片后端配置。
|
||
|
||
这里故意只实现本脚本需要的最小能力：
|
||
1. 先用 llm.scenes 把 scene 映射到 backend 名；
|
||
2. 再从 llm.backends 里取出接口配置；
|
||
3. 保持脚本简单直接，不引入项目运行时数据库逻辑。
|
||
"""
|
||
llm_config = config_data.get("llm", {}) or {}
|
||
scenes = llm_config.get("scenes", {}) or {}
|
||
backends = llm_config.get("backends", {}) or {}
|
||
|
||
backend_name = str(scenes.get(scene_name) or "").strip()
|
||
if not backend_name:
|
||
raise ValueError(f"未在 config.yaml 的 llm.scenes 中找到场景: {scene_name}")
|
||
|
||
backend_config = backends.get(backend_name, {}) or {}
|
||
if not backend_config:
|
||
raise ValueError(f"未在 config.yaml 的 llm.backends 中找到后端: {backend_name}")
|
||
|
||
return {
|
||
"backend_name": backend_name,
|
||
"provider": str(backend_config.get("provider") or "").strip(),
|
||
"api_base_url": str(backend_config.get("api_base_url") or backend_config.get("base_url") or "").strip(),
|
||
"api_key": str(backend_config.get("api_key") or "").strip(),
|
||
"model": str(backend_config.get("model") or "gpt-image-1").strip(),
|
||
"endpoint": "images/generations",
|
||
# 这里给脚本自己的后端超时兜底值也同步改成 180 秒，
|
||
# 避免配置文件里没写 timeout_seconds 时又悄悄回退到旧的 300 秒。
|
||
"timeout_seconds": int(backend_config.get("timeout_seconds") or DEFAULT_REQUEST_TIMEOUT_SECONDS),
|
||
}
|
||
|
||
|
||
def build_request_url(api_base_url: str, endpoint: str) -> str:
|
||
"""拼接图片接口 URL。"""
|
||
return f"{api_base_url.rstrip('/')}/{endpoint.lstrip('/')}"
|
||
|
||
|
||
def build_auth_header(api_key: str) -> str:
|
||
"""生成 Bearer 鉴权头。"""
|
||
normalized_api_key = str(api_key or "").strip()
|
||
if normalized_api_key.lower().startswith("bearer "):
|
||
return normalized_api_key
|
||
return f"Bearer {normalized_api_key}"
|
||
|
||
|
||
def sanitize_filename(value: str) -> str:
|
||
"""
|
||
清理文件名中的非法字符。
|
||
|
||
这里保留中英文、数字、下划线、连字符，避免 Windows 路径报错。
|
||
"""
|
||
cleaned = re.sub(r"[\\/:*?\"<>|]+", "_", value.strip())
|
||
cleaned = re.sub(r"\s+", "_", cleaned)
|
||
return cleaned or "unknown"
|
||
|
||
|
||
def fetch_dota2_heroes() -> List[Dict[str, str]]:
|
||
"""
|
||
从 OpenDota 拉取英雄信息。
|
||
|
||
返回字段说明：
|
||
1. localized_name：更适合放进中文提示词里；
|
||
2. english_name：更适合做英文辅助描述和文件夹命名；
|
||
3. hero_id：方便后续写入清单或排查问题。
|
||
"""
|
||
response = requests.get(OPENDOTA_HERO_STATS_URL, timeout=60)
|
||
response.raise_for_status()
|
||
|
||
hero_rows = response.json() or []
|
||
heroes: List[Dict[str, str]] = []
|
||
for hero_row in hero_rows:
|
||
localized_name = str(hero_row.get("localized_name") or "").strip()
|
||
internal_name = str(hero_row.get("name") or "").strip()
|
||
english_name = internal_name.replace("npc_dota_hero_", "").replace("_", " ").title()
|
||
|
||
if not localized_name:
|
||
continue
|
||
|
||
heroes.append(
|
||
{
|
||
"hero_id": str(hero_row.get("id") or "").strip(),
|
||
"localized_name": localized_name,
|
||
"english_name": english_name,
|
||
}
|
||
)
|
||
|
||
# 这里按照英雄英文名排序，保证多次运行时输出顺序稳定。
|
||
heroes.sort(key=lambda item: item["english_name"])
|
||
return heroes
|
||
|
||
|
||
def resolve_text_language(image_index: int) -> str:
|
||
"""
|
||
根据图片序号确定当前文案语言。
|
||
|
||
约定规则：
|
||
1. 第 1、2 张固定走中文排版；
|
||
2. 第 3、4 张固定走日文排版；
|
||
3. 如果用户把 count 调大，则从头循环复用这套语言分布。
|
||
"""
|
||
variant_index = (image_index - 1) % len(DEFAULT_LANGUAGE_VARIANTS)
|
||
return DEFAULT_LANGUAGE_VARIANTS[variant_index]
|
||
|
||
|
||
def build_consistent_prompt(hero: Dict[str, str], image_index: int) -> str:
|
||
"""
|
||
构造统一风格的提示词。
|
||
|
||
提示词策略：
|
||
1. 固定所有英雄共用的版式、镜头语言、色彩、文字排版、雷达图要求；
|
||
2. 只替换英雄身份信息，尽量让最终成片拥有统一系列感；
|
||
3. 用“偏 JOJO 气质、夸张漫画表现”来强化目标风格。
|
||
"""
|
||
hero_name_cn = hero["localized_name"]
|
||
hero_name_en = hero["english_name"]
|
||
text_language = resolve_text_language(image_index)
|
||
if text_language == "zh":
|
||
text_language_desc = "画面中的标题、副标题、能力说明文字统一使用中文排版，字体要有热血漫画海报感，禁止出现日文。"
|
||
text_language_label = "中文"
|
||
else:
|
||
text_language_desc = "画面中的标题、副标题、能力说明文字统一使用日文排版，字体要有热血漫画海报感，禁止出现中文。"
|
||
text_language_label = "日文"
|
||
|
||
return f"""
|
||
请为短视频封面创作一张高完成度竖版插画，主体是 Dota2 英雄 {hero_name_cn}（{hero_name_en}）。
|
||
|
||
核心要求：
|
||
1. 角色设定明确为 Dota2 的风格体系下的“至宝级华丽皮肤质感”，但角色身份必须是 {hero_name_cn} 本人，不要画成别的英雄。
|
||
2. 画面整体要强烈偏向 JOJO 气质：夸张肌肉与体块、强烈明暗对比、戏剧化姿势、锐利线条、张力十足的漫画分镜感、厚重阴影、速度线、压迫感构图。
|
||
3. 需要比普通日漫更偏 JOJO 风，风格统一、成熟、硬朗、华丽，视觉冲击力强。
|
||
4. 画面左下角固定放一个“能力雷达图”，用日式游戏 UI 风格表现，半透明发光面板，结构清晰。
|
||
5. 画面中加入醒目的文字排版，像热血漫画标题与角色名字幕，排版要高级，不能乱码。
|
||
6. 构图固定为 9:16 竖版海报，适合抖音封面，角色居中偏上，保留底部与左下角的信息区。
|
||
7. 背景使用史诗感能量、替身感氛围、漫画速度线、粒子、光效，但不要遮挡主体脸和武器。
|
||
8. 质感统一为高细节、高完成度、商业海报、收藏级插画。
|
||
9. {text_language_desc}
|
||
|
||
稳定性要求：
|
||
1. 全系列都保持相同的版式语言、相同的信息层级、相同的雷达图位置、相同的标题风格。
|
||
2. 当前是同一英雄的第 {image_index} 张候选图，本张必须输出{text_language_label}版本；请只在姿势、镜头角度、背景能量流向上做有限变化，不要改变整体系列风格。
|
||
3. 不要出现水印、签名、Logo、拼贴、多角色、手部崩坏、脸部畸形、文字糊成乱码。
|
||
""".strip()
|
||
|
||
|
||
def extract_image_bytes(response_json: Dict[str, Any], timeout_seconds: int) -> bytes:
|
||
"""
|
||
从 OpenAI 兼容响应中提取图片字节。
|
||
|
||
兼容两种常见返回格式：
|
||
1. b64_json：直接解码；
|
||
2. url：再补一次下载。
|
||
"""
|
||
data_list = response_json.get("data") or []
|
||
if not data_list:
|
||
raise ValueError(f"接口返回里没有 data 字段: {json.dumps(response_json, ensure_ascii=False)[:500]}")
|
||
|
||
first_item = data_list[0] or {}
|
||
b64_content = (
|
||
first_item.get("b64_json")
|
||
or first_item.get("image_base64")
|
||
or first_item.get("base64")
|
||
or ""
|
||
)
|
||
if b64_content:
|
||
return base64.b64decode(b64_content)
|
||
|
||
image_url = str(first_item.get("url") or first_item.get("image_url") or "").strip()
|
||
if image_url:
|
||
download_response = requests.get(image_url, timeout=timeout_seconds)
|
||
download_response.raise_for_status()
|
||
return download_response.content
|
||
|
||
raise ValueError(f"无法从响应中提取图片内容: {json.dumps(first_item, ensure_ascii=False)[:500]}")
|
||
|
||
|
||
def generate_one_image(
|
||
request_url: str,
|
||
api_key: str,
|
||
model: str,
|
||
prompt: str,
|
||
image_size: str,
|
||
image_quality: str,
|
||
timeout_seconds: int,
|
||
) -> bytes:
|
||
"""调用 OpenAI 兼容图片接口生成单张图片。"""
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": build_auth_header(api_key),
|
||
}
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"size": image_size,
|
||
"quality": image_quality,
|
||
"response_format": "b64_json",
|
||
"user": "dota2_douyin_batch_generator",
|
||
}
|
||
|
||
response = requests.post(
|
||
request_url,
|
||
headers=headers,
|
||
json=payload,
|
||
timeout=timeout_seconds,
|
||
)
|
||
response.raise_for_status()
|
||
response_json = response.json() or {}
|
||
return extract_image_bytes(response_json, timeout_seconds)
|
||
|
||
|
||
def append_manifest_row(manifest_path: Path, row: Dict[str, Any]) -> None:
|
||
"""
|
||
以 JSONL 方式追加生成记录。
|
||
|
||
这样做的好处是：
|
||
1. 即使脚本中途停止，前面已成功的记录也不会丢；
|
||
2. 方便后续按英雄筛选、统计或补图。
|
||
"""
|
||
with manifest_path.open("a", encoding="utf-8") as file_obj:
|
||
file_obj.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||
|
||
|
||
def ensure_output_dir(output_dir: Path) -> None:
|
||
"""确保输出目录存在。"""
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
|
||
def build_generation_tasks(
|
||
heroes: List[Dict[str, str]],
|
||
output_dir: Path,
|
||
count_per_hero: int,
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
预先展开所有生图任务。
|
||
|
||
这样做的目的：
|
||
1. 先把“英雄 x 第几张图”拍平成统一任务列表，便于线程池直接消费；
|
||
2. 任务对象中提前算好输出目录、文件名、提示词，线程里只负责执行；
|
||
3. 任务顺序保持稳定，后续日志更容易排查。
|
||
"""
|
||
tasks: List[Dict[str, Any]] = []
|
||
total_heroes = len(heroes)
|
||
|
||
for hero_index, hero in enumerate(heroes, start=1):
|
||
hero_slug = sanitize_filename(hero["english_name"].lower().replace(" ", "_"))
|
||
hero_dir = output_dir / f"{hero_slug}_{sanitize_filename(hero['localized_name'])}"
|
||
ensure_output_dir(hero_dir)
|
||
|
||
for image_index in range(1, count_per_hero + 1):
|
||
file_name = f"{hero_slug}_{image_index:02d}.png"
|
||
image_path = hero_dir / file_name
|
||
tasks.append(
|
||
{
|
||
"hero": hero,
|
||
"hero_index": hero_index,
|
||
"total_heroes": total_heroes,
|
||
"hero_slug": hero_slug,
|
||
"hero_dir": hero_dir,
|
||
"image_index": image_index,
|
||
"image_path": image_path,
|
||
"prompt": build_consistent_prompt(hero, image_index),
|
||
}
|
||
)
|
||
|
||
return tasks
|
||
|
||
|
||
def run_single_generation_task(
|
||
task: Dict[str, Any],
|
||
request_url: str,
|
||
api_key: str,
|
||
model: str,
|
||
image_size: str,
|
||
image_quality: str,
|
||
timeout_seconds: int,
|
||
max_retries: int,
|
||
delay_seconds: float,
|
||
force: bool,
|
||
manifest_path: Path,
|
||
manifest_lock: threading.Lock,
|
||
print_lock: threading.Lock,
|
||
) -> Tuple[str, Dict[str, Any]]:
|
||
"""
|
||
在线程池中执行单个图片生成任务。
|
||
|
||
返回值约定：
|
||
1. status 为 success / skipped / failed 三种之一；
|
||
2. payload 会带上日志和清单记录所需的数据，主线程只负责汇总结果；
|
||
3. manifest 写入放在线程内完成，但通过锁保证同一时刻只有一个线程落盘。
|
||
"""
|
||
hero = task["hero"]
|
||
image_index = task["image_index"]
|
||
image_path: Path = task["image_path"]
|
||
prompt = task["prompt"]
|
||
|
||
with print_lock:
|
||
print(
|
||
f"\n[{task['hero_index']}/{task['total_heroes']}] "
|
||
f"处理英雄: {hero['localized_name']} ({hero['english_name']}) "
|
||
f"- 第 {image_index} 张"
|
||
)
|
||
|
||
if image_path.exists() and not force:
|
||
with print_lock:
|
||
print(f" - 已存在，跳过: {image_path.name}")
|
||
return "skipped", {
|
||
"hero_id": hero["hero_id"],
|
||
"localized_name": hero["localized_name"],
|
||
"english_name": hero["english_name"],
|
||
"image_index": image_index,
|
||
"image_path": str(image_path.as_posix()),
|
||
}
|
||
|
||
last_error: Optional[str] = None
|
||
for retry_index in range(1, max_retries + 1):
|
||
try:
|
||
with print_lock:
|
||
print(f" - 生成第 {image_index} 张，尝试 {retry_index}/{max_retries}")
|
||
|
||
image_bytes = generate_one_image(
|
||
request_url=request_url,
|
||
api_key=api_key,
|
||
model=model,
|
||
prompt=prompt,
|
||
image_size=image_size,
|
||
image_quality=image_quality,
|
||
timeout_seconds=timeout_seconds,
|
||
)
|
||
|
||
with image_path.open("wb") as file_obj:
|
||
file_obj.write(image_bytes)
|
||
|
||
manifest_row = {
|
||
"hero_id": hero["hero_id"],
|
||
"localized_name": hero["localized_name"],
|
||
"english_name": hero["english_name"],
|
||
"image_index": image_index,
|
||
"image_path": str(image_path.as_posix()),
|
||
"size": image_size,
|
||
"quality": image_quality,
|
||
"model": model,
|
||
"request_url": request_url,
|
||
"generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
"prompt": prompt,
|
||
}
|
||
|
||
# 这里用锁保护清单写入，避免多个线程同时写 JSONL 时内容互相穿插。
|
||
with manifest_lock:
|
||
append_manifest_row(manifest_path, manifest_row)
|
||
|
||
with print_lock:
|
||
print(f" - 生成成功: {image_path.name}")
|
||
|
||
time.sleep(delay_seconds)
|
||
return "success", manifest_row
|
||
except Exception as exc:
|
||
last_error = str(exc)
|
||
with print_lock:
|
||
print(f" - 生成失败: {last_error}")
|
||
if retry_index < max_retries:
|
||
# 这里做一个简短退避，降低临时网络波动或网关限流的影响。
|
||
time.sleep(min(5, retry_index * 2))
|
||
|
||
failed_row = {
|
||
"hero_id": hero["hero_id"],
|
||
"localized_name": hero["localized_name"],
|
||
"english_name": hero["english_name"],
|
||
"image_index": image_index,
|
||
"image_path": str(image_path.as_posix()),
|
||
"size": image_size,
|
||
"quality": image_quality,
|
||
"model": model,
|
||
"request_url": request_url,
|
||
"generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
"error": last_error or "未知错误",
|
||
}
|
||
with manifest_lock:
|
||
append_manifest_row(manifest_path, failed_row)
|
||
return "failed", failed_row
|
||
|
||
|
||
def main() -> int:
|
||
"""脚本入口。"""
|
||
args = parse_args()
|
||
config_data = load_yaml_config(args.config)
|
||
backend = resolve_image_backend(config_data, args.scene)
|
||
|
||
if backend["provider"] != "openai_compatible":
|
||
raise ValueError(
|
||
f"场景 {args.scene} 对应的 provider 不是 openai_compatible，而是 {backend['provider']}"
|
||
)
|
||
if not backend["api_base_url"]:
|
||
raise ValueError("图片后端缺少 api_base_url/base_url 配置")
|
||
if not backend["api_key"]:
|
||
raise ValueError("图片后端缺少 api_key 配置")
|
||
|
||
output_dir = Path(args.output_dir)
|
||
ensure_output_dir(output_dir)
|
||
|
||
manifest_path = output_dir / "generation_manifest.jsonl"
|
||
request_url = build_request_url(backend["api_base_url"], backend["endpoint"])
|
||
# 这里优先使用命令行显式传入的超时值；
|
||
# 若用户未额外指定，则沿用 argparse 默认值 180 秒。
|
||
# 这样这个脚本的行为是稳定可预期的，不会再因为历史默认值导致请求挂太久。
|
||
timeout_seconds = int(args.timeout or backend["timeout_seconds"])
|
||
|
||
heroes = fetch_dota2_heroes()
|
||
if args.hero_filter:
|
||
keyword = args.hero_filter.lower().strip()
|
||
heroes = [
|
||
hero for hero in heroes
|
||
if keyword in hero["localized_name"].lower() or keyword in hero["english_name"].lower()
|
||
]
|
||
if args.hero_limit and args.hero_limit > 0:
|
||
heroes = heroes[:args.hero_limit]
|
||
|
||
if not heroes:
|
||
print("没有匹配到任何英雄，请检查 --hero-filter 或网络状态。", file=sys.stderr)
|
||
return 1
|
||
|
||
print(f"共准备生成 {len(heroes)} 个英雄，每个英雄 {args.count_per_hero} 张。")
|
||
print(f"图片接口: {request_url}")
|
||
print(f"输出目录: {output_dir.resolve()}")
|
||
print(f"并发线程数: {args.max_workers}")
|
||
|
||
total_success = 0
|
||
total_skipped = 0
|
||
total_failed = 0
|
||
manifest_lock = threading.Lock()
|
||
print_lock = threading.Lock()
|
||
tasks = build_generation_tasks(
|
||
heroes=heroes,
|
||
output_dir=output_dir,
|
||
count_per_hero=args.count_per_hero,
|
||
)
|
||
|
||
# 这里将所有任务交给线程池统一调度，让脚本能够同时发起 4 个图片请求。
|
||
with ThreadPoolExecutor(max_workers=max(1, int(args.max_workers))) as executor:
|
||
future_to_task = {
|
||
executor.submit(
|
||
run_single_generation_task,
|
||
task,
|
||
request_url,
|
||
backend["api_key"],
|
||
backend["model"],
|
||
args.size,
|
||
args.quality,
|
||
timeout_seconds,
|
||
args.max_retries,
|
||
args.delay,
|
||
args.force,
|
||
manifest_path,
|
||
manifest_lock,
|
||
print_lock,
|
||
): task
|
||
for task in tasks
|
||
}
|
||
|
||
for future in as_completed(future_to_task):
|
||
status, _ = future.result()
|
||
if status == "success":
|
||
total_success += 1
|
||
elif status == "skipped":
|
||
total_skipped += 1
|
||
else:
|
||
total_failed += 1
|
||
|
||
print("\n生成完成。")
|
||
print(f"成功: {total_success}")
|
||
print(f"跳过: {total_skipped}")
|
||
print(f"失败: {total_failed}")
|
||
print(f"清单文件: {manifest_path.resolve()}")
|
||
return 0 if total_failed == 0 else 2
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|