import concurrent import os import tempfile from typing import Optional, Tuple import numpy as np import spaces from transformers import pipeline import gradio as gr import torch import torchaudio from resemble_enhance.enhancer.inference import denoise, enhance from flore200_codes import flores_codes from tts import BambaraTTS # Check if CUDA is available device = "cuda" if torch.cuda.is_available() else "cpu" # Translation pipeline translation_model = "oza75/nllb-600M-mt-french-bambara" translator = pipeline("translation", model=translation_model, max_length=512) # Text-to-Speech pipeline tts_model = "oza75/bambara-tts" tts = BambaraTTS(tts_model) # Function to translate text to Bambara @spaces.GPU def translate_to_bambara(text, src_lang): translation = translator(text, src_lang=src_lang, tgt_lang="bam_Latn") return str(translation[0]['translation_text']) # Function to convert text to speech @spaces.GPU def text_to_speech(bambara_text, reference_audio: Optional[Tuple] = None): if reference_audio is not None: ref_sr, ref_audio = reference_audio ref_audio = torch.from_numpy(ref_audio) # Add a channel dimension if the audio is 1D if ref_audio.ndim == 1: ref_audio = ref_audio.unsqueeze(0) # Save the reference audio to a temporary file if it's not None with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp: torchaudio.save(tmp.name, ref_audio, ref_sr) tmp_path = tmp.name # Use the temporary file as the speaker reference sr, audio = tts.text_to_speech(bambara_text, speaker_reference_wav_path=tmp_path) # Clean up the temporary file os.unlink(tmp_path) else: # If no reference audio provided, proceed with the default sr, audio = tts.text_to_speech(bambara_text) audio = audio.mean(dim=0) return audio, sr # Function to enhance speech # @spaces.GPU # def enhance_speech(audio_array, sampling_rate, solver, nfe, tau, denoise_before_enhancement): # solver = solver.lower() # nfe = int(nfe) # lambd = 0.9 if denoise_before_enhancement else 0.1 # # @spaces.GPU(duration=360) # def denoise_audio(): # try: # return denoise(audio_array, sampling_rate, device) # except Exception as e: # print("> Error while denoising : ", str(e)) # return audio_array, sampling_rate # # @spaces.GPU(duration=360) # def enhance_audio(): # try: # return enhance(audio_array, sampling_rate, device, nfe=nfe, solver=solver, lambd=lambd, tau=tau) # except Exception as e: # print("> Error while enhancement : ", str(e)) # return audio_array, sampling_rate # # with concurrent.futures.ThreadPoolExecutor() as executor: # future_denoise = executor.submit(denoise_audio) # future_enhance = executor.submit(enhance_audio) # # denoised_audio, new_sr1 = future_denoise.result() # enhanced_audio, new_sr2 = future_enhance.result() # # # Convert to numpy and return # return (new_sr1, denoised_audio.cpu().numpy()), (new_sr2, enhanced_audio.cpu().numpy()) @spaces.GPU def enhance_speech(audio_array, sampling_rate, solver, nfe, tau, denoise_before_enhancement): solver = solver.lower() nfe = int(nfe) lambd = 0.9 if denoise_before_enhancement else 0.1 denoised_audio, new_sr1 = denoise(audio_array, sampling_rate, device) enhanced_audio, new_sr2 = enhance(audio_array, sampling_rate, device, nfe=nfe, solver=solver, lambd=lambd, tau=tau) # Convert to numpy and return return (new_sr1, denoised_audio.cpu().numpy()), (new_sr2, enhanced_audio.cpu().numpy()) def convert_to_int16(audio_array): if audio_array.dtype == torch.float32: # Assuming audio_array values are in the range [-1.0, 1.0] # Scale to int16 range and convert the datatype audio_array = (audio_array * 32767).to(torch.int16) return audio_array # Define the Gradio interface def _fn( src_lang, text, reference_audio=None, solver="Midpoint", nfe=64, prior_temp=0.5, denoise_before_enhancement=False ): source_lang = flores_codes[src_lang] # Step 1: Translate the text to Bambara bambara_text = translate_to_bambara(text, source_lang) # Step 2: Convert the translated text to speech with reference audio if reference_audio is not None: audio_array, sampling_rate = text_to_speech(bambara_text, reference_audio) else: audio_array, sampling_rate = text_to_speech(bambara_text) # # Step 3: Enhance the audio # denoised_audio, enhanced_audio = enhance_speech( # audio_array, # sampling_rate, # solver, # nfe, # prior_temp, # denoise_before_enhancement # ) # Return all outputs return ( bambara_text, # (sampling_rate, audio_array.numpy()), # (denoised_audio[0], convert_to_int16(denoised_audio[1])), # (enhanced_audio[0], convert_to_int16(enhanced_audio[1])) ) def main(): lang_codes = list(flores_codes.keys()) # Build Gradio app app = gr.Interface( fn=_fn, inputs=[ gr.Dropdown(label="Source Language", choices=lang_codes, value='French'), gr.Textbox(label="Text to Translate", lines=3), gr.Audio(label="Clone your voice (optional)", type="numpy", format="wav"), gr.Dropdown( choices=["Midpoint", "RK4", "Euler"], value="Midpoint", label="ODE Solver (Midpoint is recommended)" ), gr.Slider(minimum=1, maximum=128, value=64, step=1, label="Number of Function Evaluations"), gr.Slider(minimum=0.1, maximum=1, value=0.5, step=0.01, label="Prior Temperature"), gr.Checkbox(value=False, label="Denoise Before Enhancement") ], outputs=[ gr.Textbox(label="Translated Text"), # gr.Audio(label="Original TTS Audio"), # gr.Audio(label="Denoised Audio"), # gr.Audio(label="Enhanced Audio") ], title="Bambara Translation and Text to Speech with Audio Enhancement", description="Translate text to Bambara and convert it to speech with options to enhance audio quality." ) app.launch(share=False) if __name__ == "__main__": main()