feat: 完成T-303 AI批量生成
- 新增 generate_batch，先并发生成标题再并发生成封面，成功逐条 set_generated 落库 - Tab② 接入开始生成、停止、进度展示和双击新旧封面预览 - 新增 GenerateWorker，通过 worker signal 回传进度与行刷新 - 补充批量生成成功、失败、停止取消和 GUI worker 单元测试 - 同步任务看板、API、路由、当前状态与 progress 文档
This commit is contained in:
chengma
@@ -1,6 +1,7 @@
|
||||
"""AI generation helpers backed by configurable HTTP model endpoints."""
|
||||
|
||||
import base64
|
||||
from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
|
||||
import copy
|
||||
import json
|
||||
import mimetypes
|
||||
@@ -10,7 +11,9 @@ import urllib.error
|
||||
import urllib.request
|
||||
import uuid
|
||||
|
||||
from . import appconfig
|
||||
from . import appconfig, db
|
||||
from . import prompts as prompt_module
|
||||
from .config import make_slug
|
||||
|
||||
|
||||
class AIError(RuntimeError):
|
||||
@@ -105,6 +108,144 @@ def gen_cover(
|
||||
return _save_jpeg(image_bytes, out_path, resolution, quality)
|
||||
|
||||
|
||||
def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=None):
|
||||
"""Generate titles first, then covers, and persist each successful task."""
|
||||
|
||||
runtime = dict(ai_cfg or {})
|
||||
config = _runtime_config(runtime)
|
||||
generation_cfg = appconfig.ai_config(config)
|
||||
generation_cfg.update(
|
||||
{
|
||||
key: value
|
||||
for key, value in runtime.items()
|
||||
if key in {
|
||||
"title_concurrency",
|
||||
"image_concurrency",
|
||||
"retry",
|
||||
"jpg_quality",
|
||||
"resolution",
|
||||
}
|
||||
}
|
||||
)
|
||||
db_path = runtime.get("db_path")
|
||||
models_path = runtime.get("models_path", appconfig.AI_MODELS_PATH)
|
||||
image_root = runtime.get("image_dir") or appconfig.image_dir(config)
|
||||
account_by_alias = runtime.get("account_by_alias") or {}
|
||||
on_task_update = runtime.get("on_task_update")
|
||||
title_prompt = _prompt_value(prompts, "title")
|
||||
cover_prompt = _prompt_value(prompts, "cover")
|
||||
should_stop = should_stop or (lambda: False)
|
||||
eligible = [
|
||||
task for task in list(tasks)
|
||||
if getattr(task, "stage", None) == "collected"
|
||||
]
|
||||
summary = {
|
||||
"ok": True,
|
||||
"total": len(eligible),
|
||||
"title_done": 0,
|
||||
"cover_done": 0,
|
||||
"failed": 0,
|
||||
"cancelled": False,
|
||||
}
|
||||
_emit_generation_progress(on_progress, summary)
|
||||
title_results = {}
|
||||
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max(1, int(generation_cfg.get("title_concurrency", 1)))
|
||||
) as executor:
|
||||
futures = {}
|
||||
for task in eligible:
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
break
|
||||
futures[
|
||||
executor.submit(
|
||||
gen_title,
|
||||
title_prompt,
|
||||
getattr(task, "old_title", "") or "",
|
||||
retry=generation_cfg.get("retry"),
|
||||
config=config,
|
||||
models_path=models_path,
|
||||
)
|
||||
] = task
|
||||
for future in as_completed(futures):
|
||||
task = futures[future]
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
_cancel_pending(futures)
|
||||
try:
|
||||
title_results[task.id] = future.result()
|
||||
summary["title_done"] += 1
|
||||
except CancelledError:
|
||||
summary["cancelled"] = True
|
||||
except Exception as exc:
|
||||
summary["failed"] += 1
|
||||
summary["ok"] = False
|
||||
_mark_generate_failed(task, exc, db_path, on_task_update)
|
||||
_emit_generation_progress(on_progress, summary)
|
||||
|
||||
cover_tasks = [
|
||||
task for task in eligible
|
||||
if task.id in title_results
|
||||
]
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max(1, int(generation_cfg.get("image_concurrency", 1)))
|
||||
) as executor:
|
||||
futures = {}
|
||||
for task in cover_tasks:
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
break
|
||||
new_title = title_results[task.id]
|
||||
rendered_cover_prompt = prompt_module.render_prompt(
|
||||
cover_prompt,
|
||||
_prompt_context(task, new_title, account_by_alias),
|
||||
)
|
||||
futures[
|
||||
executor.submit(
|
||||
gen_cover,
|
||||
rendered_cover_prompt,
|
||||
getattr(task, "old_cover_path", "") or "",
|
||||
_new_cover_path(task, account_by_alias, image_root),
|
||||
resolution=generation_cfg.get("resolution"),
|
||||
jpg_quality=generation_cfg.get("jpg_quality"),
|
||||
retry=generation_cfg.get("retry"),
|
||||
config=config,
|
||||
models_path=models_path,
|
||||
)
|
||||
] = (task, new_title)
|
||||
for future in as_completed(futures):
|
||||
task, new_title = futures[future]
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
_cancel_pending(futures)
|
||||
try:
|
||||
new_cover_path = future.result()
|
||||
db.set_generated(task.id, new_title, new_cover_path, path=db_path)
|
||||
summary["cover_done"] += 1
|
||||
if on_task_update is not None:
|
||||
on_task_update(
|
||||
task.id,
|
||||
{
|
||||
"stage": "generated",
|
||||
"status": "success",
|
||||
"new_title": new_title,
|
||||
"new_cover_path": new_cover_path,
|
||||
},
|
||||
)
|
||||
except CancelledError:
|
||||
summary["cancelled"] = True
|
||||
except Exception as exc:
|
||||
summary["failed"] += 1
|
||||
summary["ok"] = False
|
||||
_mark_generate_failed(task, exc, db_path, on_task_update)
|
||||
_emit_generation_progress(on_progress, summary)
|
||||
|
||||
if summary["cancelled"]:
|
||||
summary["ok"] = False
|
||||
return summary
|
||||
|
||||
|
||||
def _role_model(category, name, models_path):
|
||||
if not name:
|
||||
raise AIError("未配置默认 %s 模型" % category)
|
||||
@@ -123,6 +264,100 @@ def _role_model(category, name, models_path):
|
||||
return model
|
||||
|
||||
|
||||
def _runtime_config(runtime):
|
||||
if runtime.get("config") is not None:
|
||||
return runtime["config"]
|
||||
config = appconfig.load_config()
|
||||
ai_updates = {
|
||||
key: value
|
||||
for key, value in runtime.items()
|
||||
if key in {
|
||||
"default_text_model",
|
||||
"default_image_model",
|
||||
"title_concurrency",
|
||||
"image_concurrency",
|
||||
"retry",
|
||||
"jpg_quality",
|
||||
"resolution",
|
||||
"resolution_timeouts",
|
||||
}
|
||||
}
|
||||
if ai_updates:
|
||||
config = copy.deepcopy(config)
|
||||
config.setdefault("ai", {}).update(ai_updates)
|
||||
return config
|
||||
|
||||
|
||||
def _prompt_value(prompt_values, name):
|
||||
if isinstance(prompt_values, dict):
|
||||
return str(
|
||||
prompt_values.get(name)
|
||||
or prompt_values.get(f"{name}_prompt")
|
||||
or ""
|
||||
)
|
||||
return ""
|
||||
|
||||
|
||||
def _prompt_context(task, new_title, account_by_alias):
|
||||
return {
|
||||
"old_title": getattr(task, "old_title", ""),
|
||||
"new_title": new_title,
|
||||
"item_id": getattr(task, "item_id", ""),
|
||||
"account_name": _account_name(task, account_by_alias),
|
||||
"alias": getattr(task, "alias", ""),
|
||||
}
|
||||
|
||||
|
||||
def _account_name(task, account_by_alias):
|
||||
alias = str(getattr(task, "alias", "") or "").strip()
|
||||
account = account_by_alias.get(alias)
|
||||
if account is not None:
|
||||
return getattr(account, "account_name", "") or alias
|
||||
return getattr(task, "account_name", "") or alias
|
||||
|
||||
|
||||
def _new_cover_path(task, account_by_alias, image_root):
|
||||
alias = str(getattr(task, "alias", "") or "").strip()
|
||||
account = account_by_alias.get(alias)
|
||||
slug = getattr(account, "slug", None) if account is not None else None
|
||||
if not slug:
|
||||
slug = make_slug(alias or getattr(task, "account_name", "") or "unknown")
|
||||
return os.path.abspath(
|
||||
os.path.join(
|
||||
image_root,
|
||||
slug,
|
||||
"%s_new.jpg" % getattr(task, "item_id", ""),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _mark_generate_failed(task, exc, db_path, on_task_update):
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
db.mark_failed(task.id, "generate", error, path=db_path)
|
||||
if on_task_update is not None:
|
||||
on_task_update(task.id, {"status": "failed", "last_error": error})
|
||||
|
||||
|
||||
def _cancel_pending(futures):
|
||||
for future in futures:
|
||||
if not future.done():
|
||||
future.cancel()
|
||||
|
||||
|
||||
def _emit_generation_progress(on_progress, summary):
|
||||
if on_progress is None:
|
||||
return
|
||||
payload = dict(summary)
|
||||
try:
|
||||
on_progress(payload)
|
||||
except TypeError:
|
||||
on_progress(
|
||||
payload.get("title_done", 0),
|
||||
payload.get("cover_done", 0),
|
||||
payload.get("failed", 0),
|
||||
)
|
||||
|
||||
|
||||
def _attempt_count(ai_cfg, retry):
|
||||
retry_count = ai_cfg.get("retry", 2) if retry is None else retry
|
||||
return max(1, int(retry_count) + 1)
|
||||
|
||||
Reference in New Issue
Block a user