Update aduc_framework/managers/mmaudio_manager.py

aduc_framework/managers/mmaudio_manager.py

@@ -1,23 +1,12 @@
 # managers/mmaudio_manager.py
-# AducSdr: Uma implementação aberta e funcional da arquitetura ADUC-SDR
-# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
 #
-# Contato:
-# Carlos Rodrigues dos Santos
-# carlex22@gmail.com
-# Rua Eduardo Carlos Pereira, 4125, B1 Ap32, Curitiba, PR, Brazil, CEP 8102025
+# Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Repositórios e Projetos Relacionados:
-# GitHub: https://github.com/carlex22/Aduc-sdr
+# Version: 3.0.0 (GPU Pool Manager)
 #
-# PENDING PATENT NOTICE: Please see NOTICE.md.
-#
-# Version: 2.3.0
-#
-# This file defines the MMAudioManager 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.
+# Esta versão refatora o MMAudioManager para um modelo de Pool com Workers,
+# permitindo o uso de múltiplas GPUs dedicadas para a geração de áudio
+# com um sistema de rodízio para gerenciamento eficiente de VRAM.
 
 import torch
 import logging
@@ -26,183 +15,212 @@ import os
 import time
 import yaml
 import gc
+import threading
 from pathlib import Path
 import gradio as gr
 import sys
 
+# Imports relativos para o hardware_manager
+from ..tools.hardware_manager import hardware_manager
+
 logger = logging.getLogger(__name__)
 
-# --- Dependency Management ---
+# --- Gerenciamento de Dependências ---
 DEPS_DIR = Path("./deps")
 MMAUDIO_REPO_DIR = DEPS_DIR / "MMAudio"
 MMAUDIO_REPO_URL = "https://github.com/hkchengrex/MMAudio.git"
 
-def setup_mmaudio_dependencies():
-    """
-    Ensures the MMAudio repository is cloned and available in the sys.path.
-    This function is run once when the module is first imported.
-    """
-    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", "--depth", "1", 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.")
-
-setup_mmaudio_dependencies()
-
-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
-
-
-class MMAudioManager:
-    """
-    Manages the MMAudio model for audio generation tasks.
-    """
-    def __init__(self, workspace_dir):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+# Lazy-loaded imports
+ModelConfig, all_model_cfg, mmaudio_generate, load_video, make_video = None, None, None, None, None
+MMAudio, get_my_mmaudio = None, None
+FeaturesUtils = None
+SequenceConfig = None
+FlowMatching = None
+
+class MMAudioWorker:
+    """Representa uma única instância do pipeline MMAudio em um dispositivo."""
+    def __init__(self, device_id: str):
+        self.device = torch.device(device_id)
         self.cpu_device = torch.device("cpu")
-        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
-        self.workspace_dir = workspace_dir
+        self.dtype = torch.bfloat16 if 'cuda' in self.device.type else torch.float32
         
-        self.all_model_cfg = all_model_cfg
-        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.model_config: 'ModelConfig' = None
         
-        self._load_models_to_cpu()
-
-    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]
-            # The paths in the original config are relative, so we join them with our repo path
-            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("MMAudioManager ready on CPU.")
-        except Exception as e:
-            logger.error(f"Failed to load audio models: {e}", exc_info=True)
-            self.net = None
+        self._check_and_run_global_setup()
+        self._lazy_load_mmaudio_modules()
+        logger.info(f"MMAudio Worker inicializado para o dispositivo {self.device}.")
+
+    def _lazy_load_mmaudio_modules(self):
+        """Importa dinamicamente os módulos do MMAudio."""
+        global ModelConfig, all_model_cfg, mmaudio_generate, load_video, make_video, MMAudio, get_my_mmaudio, FeaturesUtils, SequenceConfig, FlowMatching
+        if MMAudio is not None: return
+
+        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
+        logger.info("Módulos do MMAudio foram carregados dinamicamente.")
+
+    @staticmethod
+    def _check_and_run_global_setup():
+        """Executa o setup de clonagem do repositório e download de modelos uma única vez."""
+        setup_flag = DEPS_DIR / "mmaudio.setup.complete"
+        if setup_flag.exists():
+            return True
+
+        logger.info("--- Iniciando Setup Global do MMAudio (primeira execução) ---")
+        if not MMAUDIO_REPO_DIR.exists():
+            DEPS_DIR.mkdir(exist_ok=True)
+            subprocess.run(["git", "clone", "--depth", "1", MMAUDIO_REPO_URL, str(MMAUDIO_REPO_DIR)], check=True)
+        
+        if str(MMAUDIO_REPO_DIR.resolve()) not in sys.path:
+            sys.path.insert(0, str(MMAUDIO_REPO_DIR.resolve()))
 
-    def to_gpu(self):
-        """Moves the models and utilities to the GPU before inference."""
-        if self.device == 'cpu': return
-        logger.info(f"Moving MMAudioManager to GPU ({self.device})...")
+        # Importar após adicionar ao path
+        from mmaudio.eval_utils import all_model_cfg as cfg
+        
+        # Ajustar caminhos e baixar modelos
+        for cfg_key in cfg:
+            config = cfg[cfg_key]
+            config.model_path = MMAUDIO_REPO_DIR / config.model_path
+            config.vae_path = MMAUDIO_REPO_DIR / config.vae_path
+            if config.bigvgan_16k_path:
+                config.bigvgan_16k_path = MMAUDIO_REPO_DIR / config.bigvgan_16k_path
+            config.synchformer_ckpt = MMAUDIO_REPO_DIR / config.synchformer_ckpt
+            config.download_if_needed()
+
+        setup_flag.touch()
+        logger.info("--- Setup Global do MMAudio Concluído ---")
+        return True
+
+    def initialize_models(self):
+        """Carrega os modelos do worker para a CPU e depois para a GPU designada."""
+        if self.net is not None: return
+
+        self.model_config = all_model_cfg['large_44k_v2']
+        self.seq_cfg = self.model_config.seq_cfg
+        
+        logger.info(f"Worker {self.device}: Carregando modelo MMAudio 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))
+        
+        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
+        ).eval()
+        
         self.net.to(self.device, self.dtype)
         self.feature_utils.to(self.device, self.dtype)
+        logger.info(f"Worker {self.device}: Modelos MMAudio prontos na VRAM.")
 
-    def to_cpu(self):
-        """Moves the models back to CPU and clears VRAM after inference."""
-        if self.device == 'cpu': return
-        logger.info("Unloading MMAudioManager from GPU...")
+    def unload_models(self):
+        """Descarrega os modelos da VRAM, movendo-os para a CPU."""
+        if self.net is None: return
+        logger.info(f"Worker {self.device}: Descarregando modelos MMAudio da VRAM...")
         self.net.to(self.cpu_device)
         self.feature_utils.to(self.cpu_device)
+        del self.net, self.feature_utils, self.seq_cfg, self.model_config
+        self.net, self.feature_utils, self.seq_cfg, self.model_config = None, None, None, None
         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")
-        
+    def generate_audio_internal(self, video_path: str, prompt: str, duration_seconds: float, output_path: str) -> str:
+        """Lógica de geração de áudio que roda na GPU do worker."""
         negative_prompt = "human voice, speech, talking, singing, narration"
-        logger.info(f"--- Prompt: '{prompt}' | Negative Prompt: '{negative_prompt}'")
+        rng = torch.Generator(device=self.device).manual_seed(int(time.time()))
+        fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=25)
         
-        if duration_seconds < 1:
-            logger.warning("Fragment too short (<1s). Returning original video.")
-            return video_path
+        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)
         
-        if self.device == 'cpu':
-            logger.warning("Generating audio on CPU. This may be very slow.")
+        with torch.no_grad():
+            audios = mmaudio_generate(
+                clip_video=video_info.clip_frames.unsqueeze(0).to(self.device, self.dtype),
+                sync_video=video_info.sync_frames.unsqueeze(0).to(self.device, self.dtype),
+                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]
+        
+        make_video(video_info, Path(output_path), audio_waveform, sampling_rate=self.seq_cfg.sampling_rate)
+        return output_path
 
+class MMAudioPoolManager:
+    def __init__(self, device_ids: list[str], workspace_dir: str):
+        logger.info(f"MMAUDIO POOL MANAGER: Criando workers para os dispositivos: {device_ids}")
+        self.workspace_dir = workspace_dir
+        if not device_ids or 'cpu' in device_ids:
+            raise ValueError("MMAudioPoolManager requer GPUs dedicadas.")
+        self.workers = [MMAudioWorker(device_id) for device_id in device_ids]
+        self.current_worker_index = 0
+        self.lock = threading.Lock()
+        self.last_cleanup_thread = None
+
+    def _cleanup_worker_thread(self, worker: MMAudioWorker):
+        logger.info(f"MMAUDIO CLEANUP THREAD: Iniciando limpeza de {worker.device} em background...")
+        worker.unload_models()
+
+    def generate_audio_for_video(self, video_path: str, prompt: str, duration_seconds: float, output_path_override: str = None) -> str:
+        if duration_seconds < 1:
+            logger.warning(f"Vídeo muito curto ({duration_seconds:.2f}s). Pulando geração de áudio.")
+            return video_path
+            
+        worker_to_use = None
         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)
+            with self.lock:
+                if self.last_cleanup_thread and self.last_cleanup_thread.is_alive():
+                    self.last_cleanup_thread.join()
                 
-                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")
+                worker_to_use = self.workers[self.current_worker_index]
+                previous_worker_index = (self.current_worker_index - 1 + len(self.workers)) % len(self.workers)
+                worker_to_cleanup = self.workers[previous_worker_index]
+
+                cleanup_thread = threading.Thread(target=self._cleanup_worker_thread, args=(worker_to_cleanup,))
+                cleanup_thread.start()
+                self.last_cleanup_thread = cleanup_thread
                 
-                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()
+                worker_to_use.initialize_models()
+                self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
+
+            logger.info(f"MMAUDIO POOL MANAGER: Gerando áudio em {worker_to_use.device}...")
+            
+            output_path = output_path_override or os.path.join(self.workspace_dir, f"{Path(video_path).stem}_with_audio.mp4")
+            
+            return worker_to_use.generate_audio_internal(
+                video_path=video_path, prompt=prompt, duration_seconds=duration_seconds, output_path=output_path
+            )
+        except Exception as e:
+            logger.error(f"MMAUDIO POOL MANAGER: Erro durante a geração de áudio: {e}", exc_info=True)
+            raise gr.Error(f"Falha na geração de áudio: {e}")
+
+# --- Instanciação Singleton ---
+class MMAudioPlaceholder:
+    def generate_audio_for_video(self, video_path, *args, **kwargs):
+        logger.error("MMAudio não foi inicializado pois nenhuma GPU foi alocada. Pulando etapa de áudio.")
+        return video_path
 
-# --- Singleton Instantiation ---
 try:
     with open("config.yaml", 'r') as f:
         config = yaml.safe_load(f)
     WORKSPACE_DIR = config['application']['workspace_dir']
-    mmaudio_manager_singleton = MMAudioManager(workspace_dir=WORKSPACE_DIR)
+    
+    mmaudio_gpus_required = config['specialists'].get('mmaudio', {}).get('gpus_required', 0)
+    mmaudio_device_ids = hardware_manager.allocate_gpus('MMAudio', mmaudio_gpus_required)
+
+    if mmaudio_gpus_required > 0 and 'cpu' not in mmaudio_device_ids:
+        mmaudio_manager_singleton = MMAudioPoolManager(device_ids=mmaudio_device_ids, workspace_dir=WORKSPACE_DIR)
+        logger.info("Especialista de Áudio (MMAudio Pool) pronto.")
+    else:
+        mmaudio_manager_singleton = MMAudioPlaceholder()
+        logger.warning("MMAudio Pool Manager não foi inicializado. Nenhuma GPU foi requisitada na config.yaml.")
 except Exception as e:
-    logger.error(f"Could not initialize MMAudioManager: {e}", exc_info=True)
-    mmaudio_manager_singleton = None
+    logger.critical(f"Falha CRÍTICA ao inicializar o MMAudioManager: {e}", exc_info=True)
+    mmaudio_manager_singleton = MMAudioPlaceholder()