feat(suite): confirm planned generation cost

2026-07-17 11:59:07 +08:00
parent 1ef5721a20
commit d5d991e844
6 changed files with 440 additions and 39 deletions
@@ -1839,6 +1839,9 @@ class SuiteTaskState:
    ai_price_worker: object = None
    ai_price_thread: object = None
    ai_confirmation_open: bool = False
+    generation_price_worker: object = None
+    generation_price_thread: object = None
+    generation_confirmation_open: bool = False
    download_queue: list = field(default_factory=list)
    downloads: dict = field(default_factory=dict)
    download_tokens: dict = field(default_factory=dict)
@@ -2358,14 +2361,18 @@ class ProductSuiteTab(QWidget):
        destructive=False,
        confirm_text="确认",
        cancel_text="取消",
+        default_cancel=False,
    ):
        box = QMessageBox(self)
        box.setIcon(QMessageBox.Warning if destructive else QMessageBox.Question)
        box.setWindowTitle(str(title))
        box.setText(str(message))
        confirm_button = box.addButton(str(confirm_text), QMessageBox.AcceptRole)
-        box.addButton(str(cancel_text), QMessageBox.RejectRole)
-        if not destructive:
+        cancel_button = box.addButton(str(cancel_text), QMessageBox.RejectRole)
+        if default_cancel:
+            box.setDefaultButton(cancel_button)
+            box.setEscapeButton(cancel_button)
+        elif not destructive:
            box.setDefaultButton(confirm_button)
        box.exec()
        return box.clickedButton() is confirm_button
@@ -2524,7 +2531,11 @@ class ProductSuiteTab(QWidget):
                return
            state.generation_stop_requested = True
            state.worker.cancel()
-        for worker in (state.ai_worker, state.ai_price_worker):
+        for worker in (
+            state.ai_worker,
+            state.ai_price_worker,
+            state.generation_price_worker,
+        ):
            if worker is not None:
                worker.cancel()
        if state.pull_running():
@@ -2589,6 +2600,12 @@ class ProductSuiteTab(QWidget):
        if self._loading:
            return
        if self._displayed_state is not None:
+            for worker in (
+                self._displayed_state.ai_price_worker,
+                self._displayed_state.generation_price_worker,
+            ):
+                if worker is not None:
+                    worker.cancel()
            self._save_controls_to_state(self._displayed_state)
            self._flush_prompt_save(self._displayed_state)
        state = self._state_for_index(index)
@@ -4121,10 +4138,13 @@ class ProductSuiteTab(QWidget):
        ][: ai.CMHUB_VISION_MAX_IMAGES]

    def _ai_write_catalog_params(self):
+        return self._cmhub_catalog_params("vision_alias")
+
+    def _cmhub_catalog_params(self, alias_key):
        try:
            cmhub_config = appconfig.cmhub_config(self.config)
            base_url = appconfig.normalize_cmhub_base_url(cmhub_config.get("base_url"))
-            alias = str(cmhub_config.get("vision_alias") or "").strip()
+            alias = str(cmhub_config.get(alias_key) or "").strip()
            api_key = appconfig.get_cmhub_api_key(self.cmhub_config_path)
        except Exception:
            return None
@@ -4356,6 +4376,13 @@ class ProductSuiteTab(QWidget):
        state = self._displayed_state
        if state is None:
            return
+        if state.generation_price_worker is not None:
+            state.generation_price_worker.cancel()
+            self._status("已请求取消读取套图预计扣点", "warning")
+            return
+        if state.generation_confirmation_open:
+            self._status("正在等待套图生成确认", "info")
+            return
        if state.generation_running():
            retrying = state.generation_mode == "retry"
            if state.generation_stop_requested:
@@ -4476,17 +4503,31 @@ class ProductSuiteTab(QWidget):
            for slot_index, spec in enumerate(specs):
                spec["generation_round_key"] = generation_round_key
                spec["generation_slot_index"] = slot_index
-        if confirm_batch and not self._confirm(
-            "确认生成商品套图",
-            self._generation_confirmation_message(
+        if confirm_batch:
+            return self._start_generation_price_confirmation(
                state,
                local_assets,
                specs,
-            ),
-            confirm_text="确认生成",
-            cancel_text="返回修改",
-        ):
-            return False
+                generation_round_key,
+                template_text,
+            )
+        return self._start_generation_worker(
+            state,
+            specs,
+            generation_round_key,
+            retrying=retrying,
+            retry_job_id=retry_job_id,
+        )
+
+    def _start_generation_worker(
+        self,
+        state,
+        specs,
+        generation_round_key,
+        *,
+        retrying,
+        retry_job_id,
+    ):
        self._persist_state(state)
        run_token = uuid.uuid4().hex
        worker = ProductSuiteGenerateWorker(
@@ -4540,7 +4581,213 @@ class ProductSuiteTab(QWidget):
        )
        return True

-    def _generation_confirmation_message(self, state, local_assets, specs):
+    def _start_generation_price_confirmation(
+        self,
+        state,
+        local_assets,
+        specs,
+        generation_round_key,
+        template_text,
+    ):
+        if state.generation_price_worker is not None or state.generation_confirmation_open:
+            self._status("正在读取套图预计扣点", "info")
+            return False
+        snapshot = self._generation_plan_snapshot(state, local_assets, specs, template_text)
+        params = self._cmhub_catalog_params("image_alias")
+        cached_models = (
+            cmhub_models.cached_model_catalog(params["base_url"], params["alias"])
+            if params is not None
+            else None
+        )
+        if cached_models is not None:
+            return self._confirm_generation_price_request(
+                state,
+                local_assets,
+                specs,
+                generation_round_key,
+                snapshot,
+                self._generation_points_estimate(cached_models, params["alias"], specs),
+            )
+        if params is None:
+            return self._confirm_generation_price_request(
+                state,
+                local_assets,
+                specs,
+                generation_round_key,
+                snapshot,
+                None,
+            )
+        worker = CMHubModelCatalogWorker(
+            params["base_url"],
+            params["api_key"],
+            connect_timeout=params["connect_timeout"],
+            use_system_proxy=params["use_system_proxy"],
+        )
+        state.generation_price_worker = worker
+        worker.finished.connect(
+            lambda result, state=state, assets=list(local_assets), specs=list(specs), key=generation_round_key, plan=snapshot, alias=params["alias"]:
+            self._on_generation_catalog_finished(
+                state,
+                assets,
+                specs,
+                key,
+                plan,
+                alias,
+                result,
+            )
+        )
+        worker.cancelled.connect(
+            lambda result, state=state: self._on_generation_catalog_cancelled(state, result)
+        )
+        state.generation_price_thread = self._start_thread(worker, "商品套图读取生图扣点")
+        if state is self._displayed_state:
+            self._apply_running_state(state)
+        self._status("正在读取套图预计扣点", "info")
+        return False
+
+    @staticmethod
+    def _generation_points_estimate(models, alias, specs):
+        unit_cost = cmhub_models.unambiguous_points_cost(
+            models,
+            alias,
+            "image",
+            requires_image=True,
+        )
+        if unit_cost is None:
+            return None
+        return (unit_cost, unit_cost * len(specs))
+
+    @staticmethod
+    def _generation_specs_signature(specs):
+        return tuple(
+            (
+                int(spec.get("source_asset_id") or 0),
+                str(spec.get("category") or spec.get("job_type") or ""),
+                int(spec.get("category_index") or 0),
+                int(spec.get("source_index") or 0),
+                str(spec.get("prompt") or ""),
+            )
+            for spec in specs
+        )
+
+    def _generation_plan_snapshot(self, state, local_assets, specs, template_text):
+        return {
+            "asset_ids": tuple(int(asset.id) for asset in local_assets),
+            "settings": product_suite.normalize_suite_settings(state.settings),
+            "prompt": str(state.prompt or ""),
+            "item_id": str(state.item_id or ""),
+            "template_text": str(template_text or ""),
+            "specs": self._generation_specs_signature(specs),
+        }
+
+    def _generation_plan_is_current(self, state, snapshot):
+        if not self._is_open_suite_state(state) or state.generation_running():
+            return False
+        try:
+            template_text = prompts.load_product_suite_prompt(self.product_suite_prompt_path)
+        except prompts.PromptError:
+            return False
+        local_assets = [asset for asset in self._original_assets(state) if _asset_usable(asset)]
+        current_specs = product_suite.build_job_specs(
+            local_assets,
+            state.prompt,
+            state.settings,
+            state.item_id or "未绑定商品",
+            template_text=template_text,
+        )
+        return snapshot == self._generation_plan_snapshot(
+            state,
+            local_assets,
+            current_specs,
+            template_text,
+        )
+
+    def _on_generation_catalog_finished(
+        self,
+        state,
+        local_assets,
+        specs,
+        generation_round_key,
+        snapshot,
+        alias,
+        result,
+    ):
+        state.generation_price_worker = None
+        state.generation_price_thread = None
+        if state is not self._displayed_state or not self._is_open_suite_state(state):
+            return
+        estimate = None
+        if result.get("ok") is False:
+            self._status("暂时无法取得套图预计扣点，实际以网关返回为准", "warning")
+        elif result.get("cancelled"):
+            self._apply_running_state(state)
+            return
+        else:
+            estimate = self._generation_points_estimate(result.get("models") or [], alias, specs)
+        self._confirm_generation_price_request(
+            state,
+            local_assets,
+            specs,
+            generation_round_key,
+            snapshot,
+            estimate,
+        )
+        if state is self._displayed_state:
+            self._apply_running_state(state)
+
+    def _on_generation_catalog_cancelled(self, state, result):
+        state.generation_price_worker = None
+        state.generation_price_thread = None
+        if state is self._displayed_state and self._is_open_suite_state(state):
+            self._status("已取消读取套图预计扣点", "warning")
+            self._apply_running_state(state)
+
+    def _confirm_generation_price_request(
+        self,
+        state,
+        local_assets,
+        specs,
+        generation_round_key,
+        snapshot,
+        estimate,
+    ):
+        if state is not self._displayed_state or not self._is_open_suite_state(state):
+            return False
+        if not self._generation_plan_is_current(state, snapshot):
+            self._status("商品原图或生成设置已变化，请重新点击生成套图", "warning")
+            return False
+        state.generation_confirmation_open = True
+        try:
+            confirmed = self._confirm(
+                "确认生成商品套图",
+                self._generation_confirmation_message(
+                    state,
+                    local_assets,
+                    specs,
+                    estimate=estimate,
+                ),
+                confirm_text="确认生成",
+                cancel_text="返回修改",
+                default_cancel=True,
+            )
+        finally:
+            state.generation_confirmation_open = False
+        if not confirmed:
+            if state is self._displayed_state:
+                self._apply_running_state(state)
+            return False
+        if not self._generation_plan_is_current(state, snapshot):
+            self._status("商品原图或生成设置已变化，请重新点击生成套图", "warning")
+            return False
+        return self._start_generation_worker(
+            state,
+            specs,
+            generation_round_key,
+            retrying=False,
+            retry_job_id=None,
+        )
+
+    def _generation_confirmation_message(self, state, local_assets, specs, *, estimate=None):
        counts = {}
        for spec in specs:
            category = str(
@@ -4556,26 +4803,41 @@ class ProductSuiteTab(QWidget):
            for name in ordered
            if counts.get(name, 0) > 0
        ]
-        return (
-            "店铺：%s\n"
-            "商品ID：%s\n"
-            "可用商品原图：%d张\n"
-            "每张上传图分别作为主图生成：%s\n"
-            "%s\n"
-            "图片比例：%s\n"
-            "生成总数：%d张\n"
-            "商品卖点：已填写\n\n"
-            "本次生成会消耗 cmhub 点数。"
-            % (
-                self._account_context_label(state),
-                state.item_id or "未绑定商品",
-                len(local_assets),
-                "是" if state.settings.get("per_image_primary") else "否",
-                "\n".join(category_lines),
-                state.settings.get("ratio") or "1:1",
-                len(specs),
-            )
+        per_image_primary = bool(state.settings.get("per_image_primary"))
+        lines = [
+            "店铺：%s" % self._account_context_label(state),
+            "商品ID：%s" % (state.item_id or "未绑定商品"),
+            "可用商品原图：%d张" % len(local_assets),
+            "",
+            "逐图主图：%s" % ("已开启" if per_image_primary else "未开启"),
+            (
+                "说明：白底图只使用第一张原图；场景图、卖点图和自定义分类会按每张原图分别生成。"
+                if per_image_primary
+                else "说明：所有分类都只使用第一张可用原图生成。"
+            ),
+            "",
+        ]
+        lines.extend(category_lines)
+        lines.extend(
+            [
+                "图片比例：%s" % (state.settings.get("ratio") or "1:1"),
+                "本次生成总数：%d张" % len(specs),
+                "商品卖点：已填写",
+                "",
+            ]
        )
+        if estimate is None:
+            lines.append("本次会消耗 cmhub 点数，暂时无法取得预计扣点，实际以网关返回为准。")
+        else:
+            unit_cost, total_cost = estimate
+            lines.extend(
+                [
+                    "预计单张扣点：%s 点" % cmhub_models.format_points_cost(unit_cost),
+                    "预计总扣点：%s 点" % cmhub_models.format_points_cost(total_cost),
+                    "实际扣点以 cmhub 返回为准。",
+                ]
+            )
+        return "\n".join(lines)

    def _generation_signal_token(self, payload=None):
        token = str((payload or {}).get("run_token") or "")
@@ -5066,6 +5328,7 @@ class ProductSuiteTab(QWidget):

    def _apply_running_state(self, state):
        generation_running = state.generation_running()
+        generation_price_pending = state.generation_price_worker is not None
        pull_running = state.pull_running()
        self.pull_button.setText(
            "正在停止..."
@@ -5105,6 +5368,12 @@ class ProductSuiteTab(QWidget):
                "QPushButton { background: #cf222e; color: white; border-color: #a40e26; font-weight: 600; }"
                "QPushButton:hover { background: #a40e26; }"
            )
+        elif generation_price_pending:
+            self.generate_button.setText("取消读取扣点")
+            self.generate_button.setStyleSheet(
+                "QPushButton { background: #cf222e; color: white; border-color: #a40e26; font-weight: 600; }"
+                "QPushButton:hover { background: #a40e26; }"
+            )
        else:
            self.generate_button.setStyleSheet(
                "QPushButton { background: #2f6fed; color: white; border-color: #2459c4; font-weight: 600; }"
@@ -5423,6 +5692,7 @@ class ProductSuiteTab(QWidget):
                state.import_worker,
                state.ai_worker,
                state.ai_price_worker,
+                state.generation_price_worker,
            ):
                if worker is not None and hasattr(worker, "cancel"):
                    worker.cancel()