managers/audio_specialist.py

# audio_specialist.py
#
# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
#
# Version: 2.2.0
#
# This file defines the Audio Specialist for the ADUC-SDR framework. It is responsible
# for generating audio synchronized with video clips. This version has been refactored
# to be self-contained by automatically cloning the MMAudio dependency from its
# official repository, making the framework more portable and easier to set up.

import torch
import logging
import subprocess
import os
import time
import yaml
import gc
from pathlib import Path
import gradio as gr
import sys

logger = logging.getLogger(__name__)

# --- Dependency Management ---
DEPS_DIR = Path("./deps")
MMAUDIO_REPO_DIR = DEPS_DIR / "MMAudio"
MMAUDIO_REPO_URL = "https://github.com/hkchengrex/MMAudio.git"

class AudioSpecialist:
    """
    Specialist responsible for generating audio for video fragments.
    Manages the loading and unloading of audio models from VRAM and handles
    its own code dependencies by cloning the MMAudio repository.
    """
    def __init__(self, workspace_dir):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.cpu_device = torch.device("cpu")
        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
        self.workspace_dir = workspace_dir
        
        self._mmaudio_modules_loaded = False
        self._setup_dependencies()
        self._lazy_load_mmaudio_modules()

        self.model_config: 'ModelConfig' = self.all_model_cfg['large_44k_v2']
        self.net: 'MMAudio' = None
        self.feature_utils: 'FeaturesUtils' = None
        self.seq_cfg: 'SequenceConfig' = None
        
        self._load_models_to_cpu()

    def _setup_dependencies(self):
        """
        Checks for the MMAudio repository locally. If not found, clones it.
        Then, it adds the repository to the Python path to make its modules importable.
        """
        if not MMAUDIO_REPO_DIR.exists():
            logger.info(f"MMAudio repository not found at '{MMAUDIO_REPO_DIR}'. Cloning from GitHub...")
            try:
                DEPS_DIR.mkdir(exist_ok=True)
                subprocess.run(
                    ["git", "clone", MMAUDIO_REPO_URL, str(MMAUDIO_REPO_DIR)],
                    check=True, capture_output=True, text=True
                )
                logger.info("MMAudio repository cloned successfully.")
            except subprocess.CalledProcessError as e:
                logger.error(f"Failed to clone MMAudio repository. Git stderr: {e.stderr}")
                raise RuntimeError("Could not clone the required MMAudio dependency from GitHub.")
        else:
            logger.info("Found local MMAudio repository.")
        
        if str(MMAUDIO_REPO_DIR.resolve()) not in sys.path:
            sys.path.insert(0, str(MMAUDIO_REPO_DIR.resolve()))
            logger.info(f"Added '{MMAUDIO_REPO_DIR.resolve()}' to sys.path.")
    
    def _lazy_load_mmaudio_modules(self):
        """Dynamically imports MMAudio modules only when needed."""
        if self._mmaudio_modules_loaded:
            return
        
        # These globals are now populated by the lazy loader
        global ModelConfig, all_model_cfg, mmaudio_generate, load_video, make_video
        global FlowMatching, MMAudio, get_my_mmaudio, FeaturesUtils, SequenceConfig
        
        from mmaudio.eval_utils import ModelConfig, all_model_cfg, generate as mmaudio_generate, load_video, make_video
        from mmaudio.model.flow_matching import FlowMatching
        from mmaudio.model.networks import MMAudio, get_my_mmaudio
        from mmaudio.model.utils.features_utils import FeaturesUtils
        from mmaudio.model.sequence_config import SequenceConfig

        self.all_model_cfg = all_model_cfg
        self._mmaudio_modules_loaded = True
        logger.info("MMAudio modules have been dynamically loaded.")

    def _adjust_paths_for_repo(self):
        """Adjusts the checkpoint paths in the model config to point inside the cloned repo."""
        for cfg_key in self.all_model_cfg:
            cfg = self.all_model_cfg[cfg_key]
            cfg.model_path = MMAUDIO_REPO_DIR / cfg.model_path
            cfg.vae_path = MMAUDIO_REPO_DIR / cfg.vae_path
            if cfg.bigvgan_16k_path is not None:
                cfg.bigvgan_16k_path = MMAUDIO_REPO_DIR / cfg.bigvgan_16k_path
            cfg.synchformer_ckpt = MMAUDIO_REPO_DIR / cfg.synchformer_ckpt

    def _load_models_to_cpu(self):
        """Loads the MMAudio models to CPU memory on initialization."""
        try:
            self._adjust_paths_for_repo()
            logger.info("Verifying and downloading MMAudio models, if necessary...")
            self.model_config.download_if_needed()
            
            self.seq_cfg = self.model_config.seq_cfg
            
            logger.info(f"Loading MMAudio model: {self.model_config.model_name} to CPU...")
            self.net = get_my_mmaudio(self.model_config.model_name).eval()
            self.net.load_weights(torch.load(self.model_config.model_path, map_location=self.cpu_device, weights_only=True))
            
            logger.info("Loading MMAudio feature utils to CPU...")
            self.feature_utils = FeaturesUtils(
                tod_vae_ckpt=self.model_config.vae_path,
                synchformer_ckpt=self.model_config.synchformer_ckpt,
                enable_conditions=True,
                mode=self.model_config.mode,
                bigvgan_vocoder_ckpt=self.model_config.bigvgan_16k_path,
                need_vae_encoder=False
            )
            self.feature_utils = self.feature_utils.eval()
            self.net.to(self.cpu_device)
            self.feature_utils.to(self.cpu_device)
            logger.info("Audio Specialist ready on CPU.")
        except Exception as e:
            logger.error(f"Failed to load audio models: {e}", exc_info=True)
            self.net = None

    def to_gpu(self):
        """Moves the models and utilities to the GPU before inference."""
        if self.device == 'cpu': return
        logger.info(f"Moving Audio Specialist to GPU ({self.device})...")
        self.net.to(self.device, self.dtype)
        self.feature_utils.to(self.device, self.dtype)

    def to_cpu(self):
        """Moves the models back to CPU and clears VRAM after inference."""
        if self.device == 'cpu': return
        logger.info("Unloading Audio Specialist from GPU...")
        self.net.to(self.cpu_device)
        self.feature_utils.to(self.cpu_device)
        gc.collect()
        if torch.cuda.is_available(): torch.cuda.empty_cache()

    def generate_audio_for_video(self, video_path: str, prompt: str, duration_seconds: float, output_path_override: str = None) -> str:
        """
        Generates audio for a video file, applying a negative prompt to avoid speech.
        """
        if self.net is None:
            raise gr.Error("MMAudio model is not loaded. Cannot generate audio.")

        logger.info("--- Generating Audio for Video Fragment ---")
        logger.info(f"--- Video: {os.path.basename(video_path)}")
        logger.info(f"--- Duration: {duration_seconds:.2f}s")
        
        negative_prompt = "human voice, speech, talking, singing, narration"
        logger.info(f"--- Prompt: '{prompt}' | Negative Prompt: '{negative_prompt}'")
        
        if duration_seconds < 1:
            logger.warning("Fragment too short (<1s). Returning original video.")
            return video_path
        
        if self.device == 'cpu':
            logger.warning("Generating audio on CPU. This may be very slow.")

        try:
            self.to_gpu()
            with torch.no_grad():
                rng = torch.Generator(device=self.device).manual_seed(int(time.time()))
                fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=25)
                
                video_info = load_video(Path(video_path), duration_seconds)
                self.seq_cfg.duration = video_info.duration_sec
                self.net.update_seq_lengths(self.seq_cfg.latent_seq_len, self.seq_cfg.clip_seq_len, self.seq_cfg.sync_seq_len)
                
                audios = mmaudio_generate(
                    clip_video=video_info.clip_frames.unsqueeze(0),
                    sync_video=video_info.sync_frames.unsqueeze(0),
                    text=[prompt],
                    negative_text=[negative_prompt],
                    feature_utils=self.feature_utils,
                    net=self.net,
                    fm=fm,
                    rng=rng,
                    cfg_strength=4.5
                )
                audio_waveform = audios.float().cpu()[0]
                
                output_video_path = output_path_override if output_path_override else os.path.join(self.workspace_dir, f"{Path(video_path).stem}_with_audio.mp4")
                
                make_video(video_info, Path(output_video_path), audio_waveform, sampling_rate=self.seq_cfg.sampling_rate)
                logger.info(f"--- Fragment with audio saved to: {os.path.basename(output_video_path)}")
                return output_video_path
        finally:
            self.to_cpu()

# Singleton instantiation
try:
    with open("config.yaml", 'r') as f:
        config = yaml.safe_load(f)
    WORKSPACE_DIR = config['application']['workspace_dir']
    audio_specialist_singleton = AudioSpecialist(workspace_dir=WORKSPACE_DIR)
except Exception as e:
    logger.error(f"Could not initialize AudioSpecialist: {e}", exc_info=True)
    audio_specialist_singleton = None