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

"""Voice extraction tasks - extract voice tracks and background music from videos."""

import json
import logging
import os
import subprocess
import tempfile
from typing import Optional

from celery import Task
from sqlalchemy.orm import Session
from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

from packages.adapters.sqlalchemy_impl.asset_repository import SQLAlchemyAssetRepository
from packages.shared.storage import get_storage_service

logger = logging.getLogger(__name__)


class VoiceExtractor:
    """Extract voice tracks and background music from videos using FFmpeg."""

    @staticmethod
    def _run_ffmpeg(cmd: list[str]) -> subprocess.CompletedProcess:
        """Run FFmpeg command and return result."""
        logger.info(f"Running FFmpeg: {chr(39).join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            logger.error(f"FFmpeg error: {result.stderr}")
            raise RuntimeError(f"FFmpeg failed: {result.stderr}")
        return result

    def extract_voice(
        self,
        input_path: str,
        output_path: str,
        highpass: int = 200,
        bandpass_freq: int = 300,
        bandpass_width: int = 3000,
        noise_reduction: int = 20,
    ) -> str:
        """Extract voice track from video using FFmpeg."""
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        cmd = [
            "ffmpeg",
            "-y",
            "-i",
            input_path,
            "-af",
            f"highpass=f={highpass},afftdn=bn={noise_reduction},bandpass=f={bandpass_freq}:width_type=h:width={bandpass_width},loudnorm",
            "-vn",
            "-acodec",
            "libmp3lame",
            "-q:a",
            "2",
            output_path,
        ]
        self._run_ffmpeg(cmd)
        logger.info(f"Voice extracted to: {output_path}")
        return output_path

    def extract_background(self, input_path: str, output_path: str, lowpass: int = 200) -> str:
        """Extract background music from video."""
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        cmd = [
            "ffmpeg",
            "-y",
            "-i",
            input_path,
            "-af",
            f"lowpass=f={lowpass},loudnorm",
            "-vn",
            "-acodec",
            "libmp3lame",
            "-q:a",
            "2",
            output_path,
        ]
        self._run_ffmpeg(cmd)
        logger.info(f"Background extracted to: {output_path}")
        return output_path


@celery_app.task(bind=True, max_retries=3, name="worker.extract_voice")
def extract_voice_task(self: Task, asset_id: str) -> dict:
    session = SessionLocal()
    temp_dir = tempfile.mkdtemp()
    try:
        asset_repo = SQLAlchemyAssetRepository(session)
        storage_service = get_storage_service()
        extractor = VoiceExtractor()
        asset = asset_repo.get(asset_id)
        if asset is None:
            raise ValueError(f"Asset {asset_id} not found")
        local_video_path = os.path.join(temp_dir, f"{asset_id}.mp4")
        storage_service.download_file(asset.storage_key, local_video_path)
        voice_output_path = os.path.join(temp_dir, f"{asset_id}_voice.mp3")
        extractor.extract_voice(local_video_path, voice_output_path)
        voice_storage_key = f"assets/{asset_id}/voice.mp3"
        storage_service.upload_file(voice_output_path, voice_storage_key)
        voice_url = storage_service.get_url(voice_storage_key)
        if asset.metadata is None:
            asset.metadata = {}
        asset.metadata["voice_url"] = voice_url
        asset_repo.update(asset)
        session.commit()
        logger.info(f"Voice extraction completed for asset {asset_id}: {voice_url}")
        return {"ok": True, "asset_id": asset_id, "voice_url": voice_url}
    except Exception as e:
        logger.error(f"Voice extraction failed for {asset_id}: {str(e)}")
        session.rollback()
        raise self.retry(exc=e, countdown=60)
    finally:
        session.close()
        import shutil

        shutil.rmtree(temp_dir, ignore_errors=True)


@celery_app.task(bind=True, max_retries=3, name="worker.extract_background")
def extract_background_task(self: Task, asset_id: str) -> dict:
    session = SessionLocal()
    temp_dir = tempfile.mkdtemp()
    try:
        asset_repo = SQLAlchemyAssetRepository(session)
        storage_service = get_storage_service()
        extractor = VoiceExtractor()
        asset = asset_repo.get(asset_id)
        if asset is None:
            raise ValueError(f"Asset {asset_id} not found")
        local_video_path = os.path.join(temp_dir, f"{asset_id}.mp4")
        storage_service.download_file(asset.storage_key, local_video_path)
        bg_output_path = os.path.join(temp_dir, f"{asset_id}_background.mp3")
        extractor.extract_background(local_video_path, bg_output_path)
        bg_storage_key = f"assets/{asset_id}/background.mp3"
        storage_service.upload_file(bg_output_path, bg_storage_key)
        bg_url = storage_service.get_url(bg_storage_key)
        if asset.metadata is None:
            asset.metadata = {}
        asset.metadata["background_url"] = bg_url
        asset_repo.update(asset)
        session.commit()
        logger.info(f"Background extraction completed for asset {asset_id}: {bg_url}")
        return {"ok": True, "asset_id": asset_id, "background_url": bg_url}
    except Exception as e:
        logger.error(f"Background extraction failed for {asset_id}: {str(e)}")
        session.rollback()
        raise self.retry(exc=e, countdown=60)
    finally:
        session.close()
        import shutil

        shutil.rmtree(temp_dir, ignore_errors=True)