Upload audio_specialist.py

managers/audio_specialist.py

@@ -0,0 +1,163 @@
+# 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, output_path_override: str = None) -> 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 = "human voice"
+        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 = output_path_override if output_path_override else 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