File size: 6,432 Bytes
c0e5fc7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
# audio_specialist.py (Versão final para áudio dinâmico por fragmento)
# Especialista ADUC para geração de áudio, com gerenciamento de memória GPU.
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:
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):
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):
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):
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, negative_prompt: str, duration_seconds: float) -> str:
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}'")
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}")
audio_specialist_singleton = None |