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

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

import logging
import os
import subprocess
import tempfile

from celery import Task
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)