diff --git a/core/backend/apps/ai/services.py b/core/backend/apps/ai/services.py index ae7b6fc..12292f2 100644 --- a/core/backend/apps/ai/services.py +++ b/core/backend/apps/ai/services.py @@ -819,16 +819,21 @@ _STANDALONE_TASK_TYPE = { } -def generate_standalone_image(*, team, user, prompt: str, mode: str = "image", count: int = 1) -> list[Asset]: - """不绑定项目的独立生图(图片创作 / 模特上身图 / 平台套图)。复用项目内生图链路,AITask.project=None。""" +def enqueue_standalone_images(*, team, user, prompt: str, mode: str = "image", count: int = 1) -> list[AITask]: + """独立生图(图片创作 / 模特上身图 / 平台套图)改为**异步**:本函数在 Web 请求里只做「建任务 + + 预留额度」这种秒级的活,真正 ~30s 的 ARK 出图交给 Celery worker(generate_standalone_image_task)。 + + 这样 Web 层(gunicorn)不会被慢出图请求占住 worker → 健康探针不会被饿死 → 根治"几张图就整站 502"。 + 且任务一旦提交(额度已预留),浏览器关掉 / 断网都不影响——worker 照样把图生成并落库,扣费/退费在 + worker 内闭环。返回已 RESERVED 的 AITask 列表,前端拿 id 轮询 GET /api/ai/generate-image/?ids=… 取结果。""" + from apps.ai.tasks import generate_standalone_image_task + model_config = get_default_model(ModelConfig.Capability.IMAGE) if model_config is None: raise ValueError("no active image model configured") - category = _STANDALONE_CATEGORY.get(mode, Asset.Category.UNCATEGORIZED) task_type = _STANDALONE_TASK_TYPE.get(mode, AITask.Type.PRODUCT_IMAGE) - count = max(1, min(int(count or 1), 4)) - provider = get_image_provider(model_config) - assets: list[Asset] = [] + count = max(1, min(int(count or 1), 12)) + tasks: list[AITask] = [] for index in range(count): cost = estimate_cost(model_config) task = AITask.objects.create( @@ -839,42 +844,62 @@ def generate_standalone_image(*, team, user, prompt: str, mode: str = "image", c status=AITask.Status.CREATED, model_config=model_config, idempotency_key=f"standalone-image:{team.id}:{uuid.uuid4()}", - request_payload={"model": model_config.name, "endpoint": model_config.endpoint, "prompt": prompt, "mode": mode}, + request_payload={"model": model_config.name, "endpoint": model_config.endpoint, "prompt": prompt, "mode": mode, "index": index}, estimated_cost=cost, ) + # 预留额度若余额不足会抛 ValueError,在同步的 Web 请求里立刻反馈给前端(不会先建半套任务) reserve_credit(team=team, user=user, task=task, amount=cost) task.status = AITask.Status.RESERVED task.save(update_fields=["status", "updated_at"]) - reservation = task.credit_reservation - try: - response = provider.image_generation(model=model_config.name, endpoint=model_config.endpoint, prompt=prompt) - media = provider.extract_first_media_url(response) - with transaction.atomic(): - task.status = AITask.Status.SUCCEEDED - task.response_payload = response - task.actual_cost = task.estimated_cost - task.completed_at = timezone.now() - task.save(update_fields=["status", "response_payload", "actual_cost", "completed_at", "updated_at"]) - charge_reserved_credit(reservation=reservation, actual_amount=task.actual_cost) - fileobj, content_type = VolcanoArkProvider.media_to_bytes(media) - suffix = ".jpg" if "jpeg" in content_type else (".webp" if "webp" in content_type else ".png") - asset_id = uuid.uuid4() - object_key = f"teams/{team.id}/standalone/{asset_id}{suffix}" - stored = TosStorage().upload_fileobj(fileobj=fileobj, object_key=object_key, content_type=content_type) - asset = Asset.objects.create( - id=asset_id, team=team, created_by=user, name=f"AI 生成 · {mode} · {index + 1}", - asset_type=Asset.Type.IMAGE, source=Asset.Source.AI_GENERATED, category=category, origin_task=task, - ) - AssetFile.objects.create(asset=asset, object_key=stored.object_key, bucket=stored.bucket, content_type=stored.content_type, size_bytes=stored.size_bytes, is_primary=True) - assets.append(asset) - except Exception as exc: - task.status = AITask.Status.FAILED - task.error_message = str(exc) + tasks.append(task) + # 额度都预留成功后再统一派发,避免"派发了任务但后面某张预留失败"的半成品状态 + for task in tasks: + generate_standalone_image_task.delay(str(task.id)) + return tasks + + +def run_standalone_image_task(*, task_id: str) -> None: + """Celery worker 内执行**单张**图的慢活:调 ARK → 成功落库扣费 / 失败退费。 + 幂等:只处理 RESERVED 状态的任务,重复投递(celery retry / 重启重放)不会二次出图、二次扣费。""" + task = AITask.objects.select_related("team", "created_by", "model_config").filter(id=task_id).first() + if task is None or task.status != AITask.Status.RESERVED: + return + team = task.team + user = task.created_by + payload = task.request_payload or {} + prompt = str(payload.get("prompt") or "") + mode = str(payload.get("mode") or "image") + index = int(payload.get("index") or 0) + category = _STANDALONE_CATEGORY.get(mode, Asset.Category.UNCATEGORIZED) + model_config = task.model_config + provider = get_image_provider(model_config) + reservation = task.credit_reservation + try: + response = provider.image_generation(model=model_config.name, endpoint=model_config.endpoint, prompt=prompt) + media = provider.extract_first_media_url(response) + with transaction.atomic(): + task.status = AITask.Status.SUCCEEDED + task.response_payload = response + task.actual_cost = task.estimated_cost task.completed_at = timezone.now() - task.save(update_fields=["status", "error_message", "completed_at", "updated_at"]) - release_credit(reservation=reservation, reason=str(exc)) - raise - return assets + task.save(update_fields=["status", "response_payload", "actual_cost", "completed_at", "updated_at"]) + charge_reserved_credit(reservation=reservation, actual_amount=task.actual_cost) + fileobj, content_type = VolcanoArkProvider.media_to_bytes(media) + suffix = ".jpg" if "jpeg" in content_type else (".webp" if "webp" in content_type else ".png") + asset_id = uuid.uuid4() + object_key = f"teams/{team.id}/standalone/{asset_id}{suffix}" + stored = TosStorage().upload_fileobj(fileobj=fileobj, object_key=object_key, content_type=content_type) + asset = Asset.objects.create( + id=asset_id, team=team, created_by=user, name=f"AI 生成 · {mode} · {index + 1}", + asset_type=Asset.Type.IMAGE, source=Asset.Source.AI_GENERATED, category=category, origin_task=task, + ) + AssetFile.objects.create(asset=asset, object_key=stored.object_key, bucket=stored.bucket, content_type=stored.content_type, size_bytes=stored.size_bytes, is_primary=True) + except Exception as exc: # noqa: BLE001 — 失败要退费并把错误记进 AITask 供前端轮询读取;不向上抛(避免 celery 重试二次扣费) + task.status = AITask.Status.FAILED + task.error_message = str(exc) + task.completed_at = timezone.now() + task.save(update_fields=["status", "error_message", "completed_at", "updated_at"]) + release_credit(reservation=reservation, reason=str(exc)) # ── 旁白配音(TTS):每镜旁白合成一段语音,导出时作为人声轨混在 BGM 之上 ── diff --git a/core/backend/apps/ai/tasks.py b/core/backend/apps/ai/tasks.py index f96ea7a..2ce5043 100644 --- a/core/backend/apps/ai/tasks.py +++ b/core/backend/apps/ai/tasks.py @@ -10,3 +10,13 @@ def submit_ai_task(self, task_id: str) -> str: def poll_ai_task(self, task_id: str) -> str: return task_id + +@app.task(bind=True, max_retries=0) +def generate_standalone_image_task(self, task_id: str) -> str: + """单张独立图的慢活(ARK 出图 ~30s)在 worker 内跑,Web 层不被占住。 + 幂等且失败自退费(见 run_standalone_image_task),故 max_retries=0,不向上抛重试。""" + from apps.ai.services import run_standalone_image_task + + run_standalone_image_task(task_id=task_id) + return task_id + diff --git a/core/backend/apps/ai/views.py b/core/backend/apps/ai/views.py index d0ec7db..e73be37 100644 --- a/core/backend/apps/ai/views.py +++ b/core/backend/apps/ai/views.py @@ -9,14 +9,18 @@ from apps.common.celery_health import require_worker from .models import AITask, ModelConfig from .serializers import AITaskSerializer, ModelConfigSerializer -from .services import generate_standalone_image +from .services import enqueue_standalone_images class GenerateImageView(APIView): - """POST /api/ai/generate-image/ — 独立生图(不绑项目)· 图片创作/模特图/平台套图共用。""" + """独立生图(不绑项目)· 图片创作/模特图/平台套图共用 —— **异步**。 + + POST /api/ai/generate-image/ 提交生成,秒级返回 RESERVED 任务列表(慢出图交给 worker)。 + GET /api/ai/generate-image/?ids=… 轮询这些任务的状态;成功的任务带回成图 asset。 + """ def post(self, request): - require_worker() + require_worker() # 异步出图依赖 worker 兜底执行,没 worker 直接拒绝(否则任务永远 RESERVED) prompt = str(request.data.get("prompt") or "").strip() if not prompt: return Response({"detail": "prompt 不能为空"}, status=status.HTTP_400_BAD_REQUEST) @@ -27,12 +31,34 @@ class GenerateImageView(APIView): count = 1 team = get_current_team(request.user) try: - assets = generate_standalone_image(team=team, user=request.user, prompt=prompt, mode=mode, count=count) - except ValueError as exc: + tasks = enqueue_standalone_images(team=team, user=request.user, prompt=prompt, mode=mode, count=count) + except ValueError as exc: # 无可用模型 / 余额不足等,立即反馈 return Response({"detail": str(exc)}, status=status.HTTP_400_BAD_REQUEST) - except Exception as exc: # noqa: BLE001 — 生成失败已回滚额度,返回明确错误给前端 - return Response({"detail": f"生成失败: {exc}"}, status=status.HTTP_502_BAD_GATEWAY) - return Response({"assets": AssetSerializer(assets, many=True).data}, status=status.HTTP_201_CREATED) + return Response( + {"tasks": [{"id": str(t.id), "status": t.status} for t in tasks]}, + status=status.HTTP_202_ACCEPTED, + ) + + def get(self, request): + team = get_current_team(request.user) + ids = [s for s in str(request.query_params.get("ids") or "").split(",") if s] + if not ids: + return Response({"tasks": []}) + tasks = AITask.objects.filter(team=team, id__in=ids).prefetch_related( + "generated_assets", "generated_assets__files" + ) + data = [ + { + "id": str(t.id), + "status": t.status, + "error_message": t.error_message, + "assets": AssetSerializer( + [a for a in t.generated_assets.all() if not a.is_deleted], many=True + ).data, + } + for t in tasks + ] + return Response({"tasks": data}) class AITaskViewSet(TeamScopedViewSetMixin, ReadOnlyModelViewSet): diff --git a/core/frontend/src/App.tsx b/core/frontend/src/App.tsx index 2b3e133..1eff149 100644 --- a/core/frontend/src/App.tsx +++ b/core/frontend/src/App.tsx @@ -355,29 +355,30 @@ export function App() { } function generateImages(payload: { prompt: string; mode?: "image" | "model" | "cover"; count?: number }) { - // 每张图一条独立 HTTP 请求(不再让一条请求串行出 N 张)——单张约 30s 不会触发后端请求超时。 - // 但不全量并行:同时最多 CONCURRENCY 条,避免 N 条一起砸向单个后端 pod 致内存峰值翻倍 OOM→502。 - // 任一张失败只丢那一张(各自独立扣费/回滚),其余正常返回。 - const n = Math.max(1, Math.min(Number(payload.count) || 1, 12)); - const CONCURRENCY = 2; + // 异步生图:提交后立刻拿到任务列表,前端轮询直到出图。慢的 ARK 出图在 Celery worker 里跑—— + // Web 层不被 ~30s 请求占住 → 健康探针不饿死 → 根治"几张图整站 502";且提交成功后浏览器关掉/断网, + // worker 仍会把图生成并落库(扣费/退费在 worker 内闭环),重开素材库即可见。 return action(async () => { - const results: Array<{ ok: true; assets: Asset[] } | { ok: false; reason: unknown }> = new Array(n); - let next = 0; - async function worker() { - for (let i = next++; i < n; i = next++) { - try { - const res = await api.generateImage({ prompt: payload.prompt, mode: payload.mode, count: 1 }); - results[i] = { ok: true, assets: res.assets }; - } catch (reason) { - results[i] = { ok: false, reason }; - } + const { tasks } = await api.submitGenerateImage(payload); + const pending = new Set(tasks.map((t) => t.id)); + if (pending.size === 0) throw new Error("未能提交生成任务"); + const TERMINAL = new Set(["succeeded", "failed", "cancelled", "compensating"]); + const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + const assets: Asset[] = []; + let lastErr = ""; + const deadline = Date.now() + 6 * 60 * 1000; // 兜底:6 分钟还没全好就停止轮询(图仍会在后台出完) + while (pending.size > 0 && Date.now() < deadline) { + await sleep(2500); + const res = await api.generateImageStatus([...pending]); + for (const t of res.tasks) { + if (!TERMINAL.has(t.status)) continue; + pending.delete(t.id); + if (t.status === "succeeded") assets.push(...(t.assets || [])); + else if (t.error_message) lastErr = t.error_message; } } - await Promise.all(Array.from({ length: Math.min(CONCURRENCY, n) }, worker)); - const assets = results.flatMap((r) => (r.ok ? r.assets : [])); if (assets.length === 0) { - const firstErr = results.find((r) => !r.ok) as { ok: false; reason: unknown } | undefined; - throw new Error(firstErr ? String((firstErr.reason as Error)?.message || firstErr.reason || "生成失败") : "未生成任何图片"); + throw new Error(lastErr || (pending.size > 0 ? "生成超时,图片仍在后台生成,稍后可在素材库查看" : "未生成任何图片")); } return { assets }; }, "图片已生成"); diff --git a/core/frontend/src/api.ts b/core/frontend/src/api.ts index 0974bab..be50f45 100644 --- a/core/frontend/src/api.ts +++ b/core/frontend/src/api.ts @@ -296,8 +296,12 @@ export const api = { aiTasks() { return request>("/api/ai/tasks/"); }, - generateImage(payload: { prompt: string; mode?: "image" | "model" | "cover"; count?: number }) { - return request<{ assets: Asset[] }>("/api/ai/generate-image/", { method: "POST", body: JSON.stringify(payload) }); + // 异步生图:提交后秒级返回任务列表(慢出图在 worker 跑),再用 generateImageStatus 轮询取结果 + submitGenerateImage(payload: { prompt: string; mode?: "image" | "model" | "cover"; count?: number }) { + return request<{ tasks: { id: string; status: string }[] }>("/api/ai/generate-image/", { method: "POST", body: JSON.stringify(payload) }); + }, + generateImageStatus(ids: string[]) { + return request<{ tasks: { id: string; status: string; error_message: string; assets: Asset[] }[] }>(`/api/ai/generate-image/?ids=${encodeURIComponent(ids.join(","))}`); }, recharge(payload: { amount: number | string; bonus?: number | string; channel?: string }) { return request("/api/billing/recharge/", { method: "POST", body: JSON.stringify(payload) });