Update vits_inference.py

vits_inference.py

@@ -1,35 +1,22 @@
-import torch
-import torchaudio
-from vits import SynthesizerTrn
-from mel_processing import spectrogram_torch
 import argparse
+from vits import SynthesizerTrn
+import torchaudio
+import os
 
-def load_model(model_path, config_path):
-    hps = torch.load(config_path)
-    model = SynthesizerTrn(
-        hps.data.filter_length // 2 + 1,
-        hps.train.segment_size // hps.data.hop_length,
-        **hps.model).cuda()
-    
-    model.load_state_dict(torch.load(model_path, map_location="cuda"))
-    model.eval()
-    return model, hps
+# Argumentos
+parser = argparse.ArgumentParser(description="Dublagem com VITS")
+parser.add_argument("--text", type=str, required=True, help="Texto para dublagem")
+parser.add_argument("--input_audio", type=str, required=True, help="Áudio original para clonagem de voz")
+parser.add_argument("--output_audio", type=str, required=True, help="Áudio de saída dublado")
+parser.add_argument("--language", type=str, required=True, help="Idioma da dublagem")
 
-def synthesize_text(text, model, hps):
-    text = torch.LongTensor([ord(c) for c in text]).unsqueeze(0).cuda()
-    with torch.no_grad():
-        audio = model.infer(text, noise_scale=0.667, length_scale=1.0)[0]
-    return audio.cpu().numpy()
+args = parser.parse_args()
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--text", type=str, required=True, help="Texto a ser dublado")
-    parser.add_argument("--output_audio", type=str, required=True, help="Arquivo de saída")
-    args = parser.parse_args()
+# Carregar modelo VITS
+model = SynthesizerTrn(args.language)
 
-    model, hps = load_model("vits_model.pth", "vits_config.pth")
-    audio_data = synthesize_text(args.text, model, hps)
-    torchaudio.save(args.output_audio, torch.tensor(audio_data), sample_rate=22050)
+# Processar dublagem
+waveform, sample_rate = model.synthesize(args.text, args.input_audio)
+torchaudio.save(args.output_audio, waveform, sample_rate)
 
-if __name__ == "__main__":
-    main()
+print(f"✅ Dublagem concluída: {args.output_audio}")