"""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)