feat: cap cmhub image concurrency

app/ai.py

@@ -2,7 +2,7 @@
 
 import base64
 import ipaddress
-from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
+from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
 import copy
 import json
 import mimetypes
@@ -39,6 +39,7 @@ class CMHubError(AIError):
 CMHUB_IMAGE_MAX_BYTES = 20 * 1024 * 1024
 CMHUB_TITLE_READ_TIMEOUT_SECONDS = 600
 CMHUB_IMAGE_READ_TIMEOUT_SECONDS = 650
+CMHUB_IMAGE_CONCURRENCY_LIMIT = 5
 
 _RESOLUTION_SIZES = {
     "512": (512, 512),
@@ -124,12 +125,7 @@ def gen_cover(
 ):
     """Generate a new cover image and save it as a JPEG file."""
 
-    _notify_step(on_step, "cover_validate_input")
-    old_cover_path = os.path.abspath(str(old_cover_path))
-    if not os.path.exists(old_cover_path):
-        raise FileNotFoundError("旧封面图片不存在: %s" % old_cover_path)
-    if not out_path:
-        raise AIError("缺少新封面输出路径")
+    old_cover_path, out_path = _prepare_cover_input(old_cover_path, out_path, on_step)
 
     cfg = appconfig.load_config() if config is None else config
     ai_cfg = appconfig.ai_config(cfg)
@@ -198,6 +194,36 @@ def gen_cover(
     return _save_jpeg(image_bytes, out_path, resolution, quality)
 
 
+def cmhub_image_concurrency_plan(ai_cfg):
+    """Return protected cmhub image request/download concurrency."""
+
+    configured = _positive_int((ai_cfg or {}).get("image_concurrency", 1), 1)
+    actual = min(configured, CMHUB_IMAGE_CONCURRENCY_LIMIT)
+    return {
+        "configured_image_concurrency": configured,
+        "request_concurrency": actual,
+        "download_concurrency": actual,
+        "limit": CMHUB_IMAGE_CONCURRENCY_LIMIT,
+    }
+
+
+def _positive_int(value, default=1):
+    try:
+        return max(1, int(value))
+    except (TypeError, ValueError):
+        return max(1, int(default or 1))
+
+
+def _prepare_cover_input(old_cover_path, out_path, on_step=None):
+    _notify_step(on_step, "cover_validate_input")
+    old_cover_path = os.path.abspath(str(old_cover_path))
+    if not os.path.exists(old_cover_path):
+        raise FileNotFoundError("旧封面图片不存在: %s" % old_cover_path)
+    if not out_path:
+        raise AIError("缺少新封面输出路径")
+    return old_cover_path, out_path
+
+
 def _notify_step(callback, step):
     if callback is None:
         return
@@ -207,6 +233,22 @@ def _notify_step(callback, step):
         pass
 
 
+def _notify_step_event(callback, step, result="success", detail=None, level="info"):
+    if callback is None:
+        return
+    payload = {
+        "step": step,
+        "result": result,
+        "level": level,
+    }
+    if detail is not None:
+        payload["detail"] = detail
+    try:
+        callback(payload)
+    except Exception:
+        pass
+
+
 
 def _notify_retry(callback, step, attempt, attempts, exc, model):
     if callback is None:
@@ -471,83 +513,190 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
         task for task in cover_candidates
         if getattr(task, "id", None) 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]
-            try:
-                set_step(task, "cover_prompt_render")
-                _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "start")
-                rendered_cover_prompt = prompt_module.render_prompt(
-                    cover_prompt,
-                    _prompt_context(task, new_title, account_by_alias),
-                )
-                _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "success")
-                set_step(task, "cover_submit")
-                _emit_generation_event(on_event, task, "cover", "cover_submit", "start")
-                futures[
-                    executor.submit(
-                        gen_cover,
+
+    def prepare_cover_task(task, new_title):
+        set_step(task, "cover_prompt_render")
+        _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "start")
+        rendered_cover_prompt = prompt_module.render_prompt(
+            cover_prompt,
+            _prompt_context(task, new_title, account_by_alias),
+        )
+        _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "success")
+        set_step(task, "cover_submit")
+        _emit_generation_event(on_event, task, "cover", "cover_submit", "start")
+        return rendered_cover_prompt, _new_cover_path(task, account_by_alias, image_root)
+
+    def record_cover_failure(task, exc, fallback_step="cover_request"):
+        summary["failed"] += 1
+        summary["ok"] = False
+        step = get_step(task, fallback_step)
+        error = _mark_generate_failed(task, exc, db_path, on_task_update)
+        _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error")
+        _emit_generation_error(on_error, task, "cover", step, exc, error)
+
+    def record_cover_cancelled(task, fallback_step="cover_request"):
+        summary["cancelled"] = True
+        _emit_generation_event(
+            on_event,
+            task,
+            "cover",
+            get_step(task, fallback_step),
+            "cancelled",
+            level="warning",
+        )
+
+    def persist_cover_success(task, new_title, new_cover_path):
+        set_step(task, "db_write")
+        _emit_generation_event(on_event, task, "cover", "db_write", "start")
+        db.set_generated(task.id, new_title, new_cover_path, path=db_path)
+        summary["cover_done"] += 1
+        summary["generated_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,
+                },
+            )
+        _emit_generation_event(on_event, task, "cover", "db_write", "success", detail=new_cover_path)
+
+    def run_direct_cover_tasks():
+        with ThreadPoolExecutor(
+            max_workers=_positive_int(generation_cfg.get("image_concurrency", 1), 1)
+        ) as executor:
+            futures = {}
+            for task in cover_tasks:
+                if should_stop():
+                    summary["cancelled"] = True
+                    break
+                new_title = title_results[task.id]
+                try:
+                    rendered_cover_prompt, new_cover_path = prepare_cover_task(task, new_title)
+                    futures[
+                        executor.submit(
+                            gen_cover,
+                            rendered_cover_prompt,
+                            getattr(task, "old_cover_path", "") or "",
+                            new_cover_path,
+                            resolution=generation_cfg.get("resolution"),
+                            jpg_quality=generation_cfg.get("jpg_quality"),
+                            retry=generation_cfg.get("retry"),
+                            config=config,
+                            models_path=models_path,
+                            on_step=step_callback(task, "cover"),
+                            on_event=step_callback(task, "cover"),
+                            cmhub_config_path=cmhub_config_path,
+                        )
+                    ] = (task, new_title)
+                except Exception as exc:
+                    record_cover_failure(task, exc, fallback_step="cover_prompt_render")
+                    _emit_generation_progress(on_progress, summary)
+            for future in as_completed(futures):
+                task, new_title = futures[future]
+                if should_stop():
+                    summary["cancelled"] = True
+                    _cancel_pending(futures)
+                try:
+                    persist_cover_success(task, new_title, future.result())
+                except CancelledError:
+                    record_cover_cancelled(task)
+                except Exception as exc:
+                    record_cover_failure(task, exc)
+                _emit_generation_progress(on_progress, summary)
+
+    def run_cmhub_cover_tasks():
+        plan = cmhub_image_concurrency_plan(generation_cfg)
+        request_futures = {}
+        download_futures = {}
+        next_index = 0
+
+        def submit_next_request(request_executor):
+            nonlocal next_index
+            while next_index < len(cover_tasks):
+                if should_stop():
+                    summary["cancelled"] = True
+                    return False
+                task = cover_tasks[next_index]
+                next_index += 1
+                new_title = title_results[task.id]
+                try:
+                    rendered_cover_prompt, new_cover_path = prepare_cover_task(task, new_title)
+                    future = request_executor.submit(
+                        _request_cmhub_cover_image,
                         rendered_cover_prompt,
                         getattr(task, "old_cover_path", "") or "",
-                        _new_cover_path(task, account_by_alias, image_root),
+                        new_cover_path,
                         resolution=generation_cfg.get("resolution"),
                         jpg_quality=generation_cfg.get("jpg_quality"),
                         retry=generation_cfg.get("retry"),
                         config=config,
-                        models_path=models_path,
+                        cmhub_config_path=cmhub_config_path,
                         on_step=step_callback(task, "cover"),
                         on_event=step_callback(task, "cover"),
-                        cmhub_config_path=cmhub_config_path,
                     )
-                ] = (task, new_title)
-            except Exception as exc:
-                summary["failed"] += 1
-                summary["ok"] = False
-                step = get_step(task, "cover_prompt_render")
-                error = _mark_generate_failed(task, exc, db_path, on_task_update)
-                _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error")
-                _emit_generation_error(on_error, task, "cover", step, exc, error)
-                _emit_generation_progress(on_progress, summary)
-        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()
-                set_step(task, "db_write")
-                _emit_generation_event(on_event, task, "cover", "db_write", "start")
-                db.set_generated(task.id, new_title, new_cover_path, path=db_path)
-                summary["cover_done"] += 1
-                summary["generated_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,
-                        },
-                    )
-                _emit_generation_event(on_event, task, "cover", "db_write", "success", detail=new_cover_path)
-            except CancelledError:
-                summary["cancelled"] = True
-                _emit_generation_event(on_event, task, "cover", get_step(task, "cover_request"), "cancelled", level="warning")
-            except Exception as exc:
-                summary["failed"] += 1
-                summary["ok"] = False
-                step = get_step(task, "cover_request")
-                error = _mark_generate_failed(task, exc, db_path, on_task_update)
-                _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error")
-                _emit_generation_error(on_error, task, "cover", step, exc, error)
-            _emit_generation_progress(on_progress, summary)
+                    request_futures[future] = (task, new_title)
+                    return True
+                except Exception as exc:
+                    record_cover_failure(task, exc, fallback_step="cover_prompt_render")
+                    _emit_generation_progress(on_progress, summary)
+            return False
+
+        with ThreadPoolExecutor(max_workers=plan["request_concurrency"]) as request_executor, \
+                ThreadPoolExecutor(max_workers=plan["download_concurrency"]) as download_executor:
+            for _ in range(plan["request_concurrency"]):
+                if not submit_next_request(request_executor):
+                    break
+
+            while request_futures or download_futures:
+                if should_stop():
+                    summary["cancelled"] = True
+                    _cancel_pending(request_futures)
+                done, _ = wait(
+                    set(request_futures.keys()) | set(download_futures.keys()),
+                    return_when=FIRST_COMPLETED,
+                )
+                for future in done:
+                    if future in request_futures:
+                        task, new_title = request_futures.pop(future)
+                        try:
+                            request_result = future.result()
+                            download_future = download_executor.submit(
+                                _download_and_save_cmhub_cover,
+                                request_result,
+                                on_step=step_callback(task, "cover"),
+                            )
+                            download_futures[download_future] = (task, new_title)
+                        except CancelledError:
+                            record_cover_cancelled(task)
+                            _emit_generation_progress(on_progress, summary)
+                        except Exception as exc:
+                            record_cover_failure(task, exc)
+                            _emit_generation_progress(on_progress, summary)
+                    else:
+                        task, new_title = download_futures.pop(future)
+                        try:
+                            persist_cover_success(task, new_title, future.result())
+                        except CancelledError:
+                            record_cover_cancelled(task, fallback_step="cover_download")
+                        except Exception as exc:
+                            record_cover_failure(task, exc, fallback_step="cover_download")
+                        _emit_generation_progress(on_progress, summary)
+                while (
+                    not summary["cancelled"]
+                    and not should_stop()
+                    and len(request_futures) < plan["request_concurrency"]
+                    and next_index < len(cover_tasks)
+                ):
+                    if not submit_next_request(request_executor):
+                        break
+
+    if _ai_backend(generation_cfg) == "cmhub":
+        run_cmhub_cover_tasks()
+    else:
+        run_direct_cover_tasks()
 
     if summary["cancelled"]:
         summary["ok"] = False
@@ -670,6 +819,37 @@ def _gen_cover_cmhub(
     on_step=None,
     on_event=None,
 ):
+    request_result = _request_cmhub_cover_image(
+        cover_prompt,
+        old_cover_path,
+        out_path,
+        resolution=resolution,
+        jpg_quality=jpg_quality,
+        retry=retry,
+        config=config,
+        cmhub_config_path=cmhub_config_path,
+        on_step=on_step,
+        on_event=on_event,
+        validate_input=False,
+    )
+    return _download_and_save_cmhub_cover(request_result, on_step=on_step)
+
+
+def _request_cmhub_cover_image(
+    cover_prompt,
+    old_cover_path,
+    out_path,
+    resolution,
+    jpg_quality,
+    retry,
+    config,
+    cmhub_config_path,
+    on_step=None,
+    on_event=None,
+    validate_input=True,
+):
+    if validate_input:
+        old_cover_path, out_path = _prepare_cover_input(old_cover_path, out_path, on_step)
     ai_cfg = appconfig.ai_config(config)
     _notify_step(on_step, "load_image_model")
     runtime = _cmhub_runtime(config, "image", cmhub_config_path)
@@ -686,6 +866,7 @@ def _gen_cover_cmhub(
     attempts = _attempt_count(ai_cfg, retry)
     read_timeout = _cmhub_read_timeout(config, resolution)
     _notify_step(on_step, "cover_request")
+    request_started = time.perf_counter()
     data = _cmhub_call_with_retry(
         "POST",
         appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image"),
@@ -702,18 +883,59 @@ def _gen_cover_cmhub(
             exc,
         ),
     )
+    request_elapsed = time.perf_counter() - request_started
     _emit_cmhub_metadata(on_event, data, "cover_request")
     _notify_step(on_step, "cover_parse_response")
     image_url = str(data.get("image_url") or "").strip()
     if not image_url:
         raise AIError("AI 返回中没有图片数据")
+    _notify_step_event(
+        on_step,
+        "cover_request",
+        detail="cmhub 已返回 image_url，耗时 %s" % _format_seconds(request_elapsed),
+    )
+    return {
+        "image_url": image_url,
+        "connect_timeout": runtime["connect_timeout"],
+        "read_timeout": read_timeout,
+        "out_path": out_path,
+        "resolution": resolution,
+        "quality": quality,
+    }
+
+
+def _download_and_save_cmhub_cover(request_result, on_step=None):
+    image_url = request_result["image_url"]
+    connect_timeout = request_result["connect_timeout"]
+    read_timeout = request_result["read_timeout"]
+    out_path = request_result["out_path"]
+    resolution = request_result["resolution"]
+    quality = request_result["quality"]
+    _notify_step(on_step, "cover_download")
+    download_started = time.perf_counter()
     image_bytes = _download_cmhub_image(
         image_url,
-        connect_timeout=runtime["connect_timeout"],
+        connect_timeout=connect_timeout,
         read_timeout=read_timeout,
     )
+    download_elapsed = time.perf_counter() - download_started
+    _notify_step_event(
+        on_step,
+        "cover_download",
+        detail="下载完成，%s，耗时 %s"
+        % (_format_bytes(len(image_bytes)), _format_seconds(download_elapsed)),
+    )
     _notify_step(on_step, "cover_save")
-    return _save_jpeg(image_bytes, out_path, resolution, quality)
+    save_started = time.perf_counter()
+    saved_path = _save_jpeg(image_bytes, out_path, resolution, quality)
+    save_elapsed = time.perf_counter() - save_started
+    detail = "JPEG 已保存，耗时 %s" % _format_seconds(save_elapsed)
+    try:
+        detail += "，文件 %s" % _format_bytes(os.path.getsize(saved_path))
+    except OSError:
+        pass
+    _notify_step_event(on_step, "cover_save", detail=detail)
+    return saved_path
 
 
 def _cmhub_runtime(config, operation, cmhub_config_path):
@@ -965,6 +1187,24 @@ def _emit_cmhub_metadata(callback, data, step):
         pass
 
 
+def _format_seconds(seconds):
+    value = max(0.0, float(seconds or 0.0))
+    return "%.1f秒" % value
+
+
+def _format_bytes(size):
+    value = float(max(0, int(size or 0)))
+    units = ["B", "KB", "MB", "GB"]
+    unit = units[0]
+    for unit in units:
+        if value < 1024 or unit == units[-1]:
+            break
+        value /= 1024
+    if unit == "B":
+        return "%d%s" % (int(value), unit)
+    return "%.1f%s" % (value, unit)
+
+
 def _download_cmhub_image(url, connect_timeout, read_timeout, max_bytes=CMHUB_IMAGE_MAX_BYTES):
     _assert_public_http_url(url)
     try: