优化视频提炼

2026-08-27 15:21:48 +08:00
parent 51620bf25b
commit 3fbc2f2dab
7 changed files with 438 additions and 122 deletions
@@ -19,6 +19,7 @@ import subprocess
 import tempfile
 import uuid
 from dataclasses import dataclass
+from datetime import timedelta
 from functools import lru_cache
 from pathlib import Path

@@ -760,6 +761,8 @@ def serialize_digest_job(task) -> dict:
        job_status = "processing"
    elif task.status == AITask.Status.SUCCEEDED:
        job_status = "succeeded"
+    elif task.status == AITask.Status.CANCELLED:
+        job_status = "cancelled"
    else:
        job_status = "failed"
    duration = float(req.get("duration_seconds") or 0)
@@ -805,6 +808,81 @@ def get_team_digest_job(*, team, task_id) -> dict | None:
    return serialize_digest_job(task)


+_STALE_DIGEST = timedelta(minutes=15)
+
+
+def _task_reservation(task):
+    try:
+        return task.credit_reservation
+    except Exception:  # noqa: BLE001 — 无预留行
+        return None
+
+
+def cancel_team_digest(*, team, task_id, reason: str = "用户取消") -> dict | None:
+    """取消进行中的提炼:停 UI 轮询、退预留积分。worker 跑到一半会看见 CANCELLED 不再落成功。"""
+    from django.db import transaction
+    from django.utils import timezone
+
+    from apps.ai.models import AITask
+    from apps.billing.services.ledger import release_credit
+
+    with transaction.atomic():
+        task = (
+            AITask.objects.select_for_update()
+            .filter(
+                id=task_id,
+                team=team,
+                task_type=AITask.Type.VIDEO_DIGEST,
+                project__isnull=True,
+                is_deleted=False,
+                purged_at__isnull=True,
+            )
+            .first()
+        )
+        if task is None:
+            return None
+        if task.status == AITask.Status.SUCCEEDED:
+            return serialize_digest_job(task)
+        if task.status in {AITask.Status.FAILED, AITask.Status.CANCELLED}:
+            return serialize_digest_job(task)
+        reservation = _task_reservation(task)
+        task.status = AITask.Status.CANCELLED
+        task.error_message = (reason or "用户取消")[:2000]
+        task.completed_at = timezone.now()
+        task.save(update_fields=["status", "error_message", "completed_at", "updated_at"])
+    if reservation is not None:
+        try:
+            release_credit(reservation=reservation, reason=(reason or "用户取消")[:200])
+        except Exception:  # noqa: BLE001
+            logger.warning("cancel digest release_credit failed for %s", task_id, exc_info=True)
+    return serialize_digest_job(task)
+
+
+def expire_stale_team_digests(*, team) -> None:
+    """卡住超过 15 分钟的提炼自动取消并退费,避免进页永远「正在拆解」。"""
+    from django.utils import timezone
+
+    from apps.ai.models import AITask
+
+    stale = AITask.objects.filter(
+        team=team,
+        task_type=AITask.Type.VIDEO_DIGEST,
+        project__isnull=True,
+        is_deleted=False,
+        purged_at__isnull=True,
+        status__in={
+            AITask.Status.CREATED,
+            AITask.Status.RESERVED,
+            AITask.Status.SUBMITTED,
+            AITask.Status.POLLING,
+            AITask.Status.POSTPROCESSING,
+        },
+        created_at__lt=timezone.now() - _STALE_DIGEST,
+    )
+    for task in stale:
+        cancel_team_digest(team=team, task_id=task.id, reason="提炼超时已自动取消")
+
+
 def get_inflight_team_digest(*, team) -> dict | None:
    from apps.ai.models import AITask

@@ -831,8 +909,11 @@ def get_inflight_team_digest(*, team) -> dict | None:
    return serialize_digest_job(task)


-def submit_team_digest(*, team, user, upload, model_config_id=None) -> dict:
-    """秒回任务 id。慢活(抽帧 + Gemini)交给 worker,离开页面也不中断。"""
+def submit_team_digest(*, team, user, upload=None, reuse_task_id=None, model_config_id=None) -> dict:
+    """秒回任务 id。慢活(抽帧 + Gemini)交给 worker,离开页面也不中断。
+
+    取消后再点生成时,可带 reuse_task_id 复用 TOS 上已有的参考视频,不用重新上传。
+    """
    from django.db import transaction

    from apps.ai.models import AITask
@@ -844,39 +925,46 @@ def submit_team_digest(*, team, user, upload, model_config_id=None) -> dict:
    if model_config is None:
        raise VideoDigestError("视频提炼需要 Gemini 3.1 Pro(会看图),当前没有启用,请联系管理员")

-    path, suffix, size, duration = _materialize_upload(upload)
-    file_name = Path(getattr(upload, "name", "") or "参考视频.mp4").name or "参考视频.mp4"
-    width = height = 0
-    source_key = source_url = cover_key = cover_url = ""
-    try:
-        width, height = probe_video_size(path)
-        cover_key, cover_url = _store_digest_cover(team=team, jpeg=extract_cover_jpeg(path, duration))
-        source_key, source_url = _store_raw_source(team=team, path=path, suffix=suffix)
-    finally:
-        Path(path).unlink(missing_ok=True)
-    if not source_key:
-        raise VideoDigestError("参考视频上传失败，请重试")
+    if upload is not None:
+        path, suffix, size, duration = _materialize_upload(upload)
+        file_name = Path(getattr(upload, "name", "") or "参考视频.mp4").name or "参考视频.mp4"
+        width = height = 0
+        source_key = source_url = cover_key = cover_url = ""
+        try:
+            width, height = probe_video_size(path)
+            cover_key, cover_url = _store_digest_cover(team=team, jpeg=extract_cover_jpeg(path, duration))
+            source_key, source_url = _store_raw_source(team=team, path=path, suffix=suffix)
+        finally:
+            Path(path).unlink(missing_ok=True)
+        if not source_key:
+            raise VideoDigestError("参考视频上传失败，请重试")
+        request_payload = {
+            "model": model_config.name,
+            "endpoint": model_config.endpoint,
+            "feature": "video_remix",
+            "duration_seconds": round(duration, 2),
+            "file_name": file_name,
+            "title": title_from_filename(file_name),
+            "file_size": size,
+            "width": width,
+            "height": height,
+            "ratio": ratio_label(width, height),
+            "suffix": suffix,
+            "source_key": source_key,
+            "source_url": source_url,
+            "cover_key": cover_key,
+            "cover_url": cover_url,
+        }
+    else:
+        request_payload = _payload_from_existing_source(
+            team=team,
+            reuse_task_id=reuse_task_id,
+            model_config=model_config,
+        )

    quote = quote_video_digest(team=team, model_config=model_config)
-    request_payload = {
-        "model": model_config.name,
-        "endpoint": model_config.endpoint,
-        "feature": "video_remix",
-        "duration_seconds": round(duration, 2),
-        "file_name": file_name,
-        "title": title_from_filename(file_name),
-        "file_size": size,
-        "width": width,
-        "height": height,
-        "ratio": ratio_label(width, height),
-        "suffix": suffix,
-        "source_key": source_key,
-        "source_url": source_url,
-        "cover_key": cover_key,
-        "cover_url": cover_url,
-    }
    if quote.meta.get("rate"):
-        request_payload["points_per_yuan_snapshot"] = quote.meta["rate"]
+        request_payload = {**request_payload, "points_per_yuan_snapshot": quote.meta["rate"]}

    try:
        with transaction.atomic():
@@ -904,6 +992,53 @@ def submit_team_digest(*, team, user, upload, model_config_id=None) -> dict:
    return serialize_digest_job(task)


+def _payload_from_existing_source(*, team, reuse_task_id, model_config) -> dict:
+    from django.core.exceptions import ValidationError
+
+    from apps.ai.models import AITask
+
+    task_id = str(reuse_task_id or "").strip()
+    if not task_id:
+        raise VideoDigestError("请先上传参考视频")
+    try:
+        previous = AITask.objects.filter(
+            id=task_id,
+            team=team,
+            task_type=AITask.Type.VIDEO_DIGEST,
+            project__isnull=True,
+            is_deleted=False,
+            purged_at__isnull=True,
+        ).first()
+    except (ValueError, ValidationError, TypeError):
+        raise VideoDigestError("请先上传参考视频") from None
+    if previous is None:
+        raise VideoDigestError("请先上传参考视频")
+    req = previous.request_payload or {}
+    source_key = str(req.get("source_key") or "").strip()
+    if not source_key:
+        raise VideoDigestError("参考视频已失效，请重新上传")
+    file_name = str(req.get("file_name") or "") or "参考视频.mp4"
+    return {
+        "model": model_config.name,
+        "endpoint": model_config.endpoint,
+        "feature": "video_remix",
+        "duration_seconds": float(req.get("duration_seconds") or 0),
+        "file_name": file_name,
+        "title": str(req.get("title") or "").strip() or title_from_filename(file_name),
+        "file_size": int(req.get("file_size") or 0),
+        "width": int(req.get("width") or 0),
+        "height": int(req.get("height") or 0),
+        "ratio": str(req.get("ratio") or ""),
+        "suffix": str(req.get("suffix") or ".mp4"),
+        "source_key": source_key,
+        "source_url": str(req.get("source_url") or ""),
+        "cover_key": str(req.get("cover_key") or ""),
+        "cover_url": str(req.get("cover_url") or ""),
+        "video_key": str(req.get("video_key") or ""),
+        "video_url": str(req.get("video_url") or ""),
+    }
+
+
 def run_team_digest_task(*, task_id: str) -> None:
    """Worker:从 TOS 取参考视频,跑抽帧 + Gemini,失败退费。"""
    from django.db import transaction
@@ -1078,6 +1213,11 @@ def _digest_video(*, team, user, upload, project=None, product_hint="", model_co
            raise VideoDigestError(str(exc)) from exc
    reservation = task.credit_reservation
    try:
+        task.refresh_from_db()
+        if task.status == AITask.Status.CANCELLED:
+            if source_path:
+                Path(source_path).unlink(missing_ok=True)
+            return serialize_digest_job(task)
        task.status = AITask.Status.SUBMITTED
        task.submitted_at = timezone.now()
        task.save(update_fields=["status", "submitted_at", "updated_at"])
@@ -1105,6 +1245,9 @@ def _digest_video(*, team, user, upload, project=None, product_hint="", model_co
    except Exception as exc:  # noqa: BLE001
        if source_path:
            Path(source_path).unlink(missing_ok=True)
+        task.refresh_from_db()
+        if task.status == AITask.Status.CANCELLED:
+            return serialize_digest_job(task)
        _fail_digest_task(task, reservation, str(exc))
        raise

@@ -1145,6 +1288,10 @@ def _digest_video(*, team, user, upload, project=None, product_hint="", model_co
    }

    with transaction.atomic():
+        locked = AITask.objects.select_for_update().get(id=task.id)
+        if locked.status == AITask.Status.CANCELLED:
+            return serialize_digest_job(locked)
+        task = locked
        task.status = AITask.Status.SUCCEEDED
        task.request_payload = request_payload
        task.response_payload = {"digest": digest[:32000], "prompt": digest[:32000]}