xiaoxia-saas/apps/worker/worker_app/tasks/generation.py

import os
import subprocess
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlparse

import oss2

from packages.adapters.sqlalchemy_impl import (
    SQLAlchemyAssetRepository,
    SQLAlchemyGeneratedVideoRepository,
    SQLAlchemyGenerationTaskRepository,
)
from packages.domain import GeneratedVideo, GenerationTaskStatus
from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

OUTPUT_WIDTH = 1280
OUTPUT_HEIGHT = 720
OUTPUT_FPS = 25.0
OUTPUT_DURATION_SECONDS = 5.0


def _run_ffmpeg(command: list[str]) -> None:
    subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)


def _oss_bucket() -> oss2.Bucket:
    access_key_id = os.environ["OSS_ACCESS_KEY_ID"]
    access_key_secret = os.environ["OSS_ACCESS_KEY_SECRET"]
    endpoint = os.environ["OSS_ENDPOINT"]
    bucket_name = os.environ["OSS_BUCKET_NAME"]
    return oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)


def _public_oss_url(storage_key: str) -> str:
    endpoint = os.environ["OSS_ENDPOINT"].removeprefix("https://").removeprefix("http://")
    bucket_name = os.environ["OSS_BUCKET_NAME"]
    return f"https://{bucket_name}.{endpoint}/{storage_key}"


def _normalize_storage_key(storage_key_or_url: str) -> str:
    if storage_key_or_url.startswith(("http://", "https://")):
        return urlparse(storage_key_or_url).path.lstrip("/")
    return storage_key_or_url.lstrip("/")


def _download_asset(asset_storage_key: str, local_path: Path) -> bool:
    try:
        _oss_bucket().get_object_to_file(_normalize_storage_key(asset_storage_key), str(local_path))
        return local_path.exists() and local_path.stat().st_size > 0
    except Exception:
        return False


def _upload_generated_video(local_path: Path, storage_key: str) -> str:
    _oss_bucket().put_object_from_file(
        storage_key,
        str(local_path),
        headers={"Content-Type": "video/mp4"},
    )
    return _public_oss_url(storage_key)


def _probe_duration(local_path: Path) -> float:
    try:
        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(local_path),
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        return round(float(result.stdout.strip()), 3)
    except Exception:
        return OUTPUT_DURATION_SECONDS


def _create_fallback_clip(output_path: Path, title: str) -> None:
    safe_title = title.replace(":", "\\:").replace("'", "\\'")[:80]
    _run_ffmpeg(
        [
            "ffmpeg",
            "-y",
            "-f",
            "lavfi",
            "-i",
            f"color=c=#111827:s={OUTPUT_WIDTH}x{OUTPUT_HEIGHT}:d={OUTPUT_DURATION_SECONDS}:r={int(OUTPUT_FPS)}",
            "-vf",
            f"drawtext=text='{safe_title}':fontcolor=white:fontsize=48:x=(w-text_w)/2:y=(h-text_h)/2",
            "-c:v",
            "libx264",
            "-pix_fmt",
            "yuv420p",
            "-movflags",
            "+faststart",
            str(output_path),
        ]
    )


def _compose_from_asset(input_path: Path, output_path: Path) -> None:
    _run_ffmpeg(
        [
            "ffmpeg",
            "-y",
            "-i",
            str(input_path),
            "-t",
            str(OUTPUT_DURATION_SECONDS),
            "-vf",
            f"scale={OUTPUT_WIDTH}:{OUTPUT_HEIGHT}:force_original_aspect_ratio=decrease,pad={OUTPUT_WIDTH}:{OUTPUT_HEIGHT}:(ow-iw)/2:(oh-ih)/2,setsar=1",
            "-r",
            str(int(OUTPUT_FPS)),
            "-an",
            "-c:v",
            "libx264",
            "-pix_fmt",
            "yuv420p",
            "-movflags",
            "+faststart",
            str(output_path),
        ]
    )


@celery_app.task(name="worker.generate_video")
def generate_video(task_id: str) -> dict:
    """Generate and persist a real MP4 video for a generation task."""
    db = SessionLocal()
    task_repo = SQLAlchemyGenerationTaskRepository(db)
    asset_repo = SQLAlchemyAssetRepository(db)
    video_repo = SQLAlchemyGeneratedVideoRepository(db)

    task = task_repo.get(task_id)
    if task is None:
        db.close()
        return {"status": "failed", "error": "generation task not found", "task_id": task_id}

    try:
        task.status = GenerationTaskStatus.RUNNING
        task.progress = 10.0
        task.started_at = task.started_at or datetime.now(timezone.utc)
        task_repo.update(task)

        assets = [asset for asset in asset_repo.list_by_library(task.asset_library_id) if asset.mime_type.startswith("video")]

        output_name = f"generated-{task.id}.mp4"
        storage_key = f"generated/workspaces/{task.workspace_id}/projects/{task.project_id}/tasks/{task.id}/{output_name}"

        with tempfile.TemporaryDirectory(prefix="xiaoxia-generation-") as temp_dir:
            temp_path = Path(temp_dir)
            output_path = temp_path / output_name
            input_path = temp_path / "source.mp4"

            task.progress = 35.0
            task_repo.update(task)

            source_downloaded = False
            if assets:
                source_downloaded = _download_asset(assets[0].storage_key, input_path)

            if source_downloaded:
                _compose_from_asset(input_path, output_path)
            else:
                _create_fallback_clip(output_path, f"Xiaoxia Generated Video {task.id[:8]}")

            task.progress = 70.0
            task_repo.update(task)

            file_url = _upload_generated_video(output_path, storage_key)
            file_size = output_path.stat().st_size
            duration = _probe_duration(output_path)

        video = GeneratedVideo.create(
            workspace_id=task.workspace_id,
            project_id=task.project_id,
            generation_task_id=task.id,
            name=output_name,
            file_url=file_url,
            file_size=file_size,
            duration=duration,
            width=OUTPUT_WIDTH,
            height=OUTPUT_HEIGHT,
            fps=OUTPUT_FPS,
            thumbnail_url=None,
        )
        video_repo.create(video)

        task.status = GenerationTaskStatus.COMPLETED
        task.progress = 100.0
        task.result_count = 1
        task.error_message = ""
        task.completed_at = datetime.now(timezone.utc)
        task_repo.update(task)

        return {"status": "completed", "task_id": task.id, "video_id": video.id, "file_url": file_url}
    except Exception as error:
        task.status = GenerationTaskStatus.FAILED
        task.error_message = str(error)
        task.completed_at = datetime.now(timezone.utc)
        task_repo.update(task)
        return {"status": "failed", "task_id": task.id, "error": str(error)}
    finally:
        db.close()