audio_specialist.py

# audio_specialist.py
# Especialista ADUC para geração de áudio, com gerenciamento de memória GPU.
# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos

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

# Importa as classes e funções necessárias do MMAudio
try:
    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
except ImportError:
    raise ImportError("MMAudio não foi encontrado. Por favor, instale-o a partir do GitHub: git+https://github.com/hkchengrex/MMAudio.git")

logger = logging.getLogger(__name__)

class AudioSpecialist:
    """
    Especialista responsável por gerar áudio para fragmentos de vídeo.
    Gerencia o carregamento e descarregamento de modelos de áudio da VRAM.
    """
    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.model_config: ModelConfig = 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 _load_models_to_cpu(self):
        """Carrega os modelos MMAudio para a memória da CPU na inicialização."""
        try:
            logger.info("Verificando e baixando modelos MMAudio, se necessário...")
            self.model_config.download_if_needed()
            
            self.seq_cfg = self.model_config.seq_cfg
            
            logger.info(f"Carregando modelo MMAudio: {self.model_config.model_name} para a 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("Carregando utilitários de features do MMAudio para a 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("Especialista de áudio pronto na CPU.")
        except Exception as e:
            logger.error(f"Falha ao carregar modelos de áudio: {e}", exc_info=True)
            self.net = None

    def to_gpu(self):
        """Move os modelos e utilitários para a GPU antes da inferência."""
        if self.device == 'cpu': return
        logger.info(f"Movendo especialista de áudio para a GPU ({self.device})...")
        self.net.to(self.device, self.dtype)
        self.feature_utils.to(self.device, self.dtype)

    def to_cpu(self):
        """Move os modelos de volta para a CPU e limpa a VRAM após a inferência."""
        if self.device == 'cpu': return
        logger.info("Descarregando especialista de áudio da 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) -> str:
        """
        Gera áudio para um arquivo de vídeo, aplicando um prompt negativo para evitar fala.

        Args:
            video_path (str): Caminho para o vídeo silencioso.
            prompt (str): Descrição da cena para guiar a geração de SFX.
            duration_seconds (float): Duração do áudio a ser gerado.

        Returns:
            str: Caminho para o novo arquivo de vídeo com áudio.
        """
        if self.net is None:
            raise gr.Error("Modelo MMAudio não está carregado. Não é possível gerar áudio.")

        logger.info("------------------------------------------------------")
        logger.info("--- Gerando Áudio para Fragmento de Vídeo ---")
        logger.info(f"--- Vídeo Fragmento: {os.path.basename(video_path)}")
        logger.info(f"--- Duração: {duration_seconds:.2f}s")
        logger.info(f"--- Prompt (Descrição da Cena): '{prompt}'")
        
        negative_prompt = "speech, human voice, talking, vocals, music, singing, dialogue"
        logger.info(f"--- Negative Prompt: '{negative_prompt}'")
        
        if duration_seconds < 1:
            logger.warning("Fragmento muito curto (<1s). Retornando vídeo silencioso.")
            logger.info("------------------------------------------------------")
            return video_path

        if self.device == 'cpu':
            logger.warning("Gerando áudio na CPU. Isso pode ser muito lento.")

        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]
                
                fragment_name = Path(video_path).stem
                output_video_path = os.path.join(self.workspace_dir, f"{fragment_name}_com_audio.mp4")
                
                make_video(video_info, Path(output_video_path), audio_waveform, sampling_rate=self.seq_cfg.sampling_rate)
                logger.info(f"--- Fragmento com áudio salvo em: {os.path.basename(output_video_path)}")
                logger.info("------------------------------------------------------")
                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"Não foi possível inicializar o AudioSpecialist: {e}", exc_info=True)
    audio_specialist_singleton = None