Spaces:
Running
Running
| # Imports | |
| import gradio as gr | |
| import spaces | |
| import os | |
| import torch | |
| import torchaudio | |
| import time | |
| from zonos.model import Zonos | |
| from zonos.conditioning import make_cond_dict, supported_language_codes | |
| # Variables | |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") | |
| device = "cuda" | |
| REPO = "Zyphra/Zonos-v0.1-transformer" | |
| model = Zonos.from_pretrained(REPO, device=device) | |
| # Functions | |
| def patch_cuda(): | |
| if torch.cuda.is_available(): | |
| for i in range(torch.cuda.device_count()): | |
| p = torch.cuda.get_device_properties(i) | |
| if not hasattr(p, "regs_per_multiprocessor"): | |
| setattr(p, "regs_per_multiprocessor", 65536) | |
| if not hasattr(p, "max_threads_per_multi_processor"): | |
| setattr(p, "max_threads_per_multi_processor", 2048) | |
| def generate(input, language, speaker_audio, emotion_happy, emotion_sad, emotion_disgust, emotion_fear, emotion_surprise, emotion_anger, emotion_other, emotion_neutral, clarity, fmax, pitch_std, speaking_rate, dnsmos_ovrl, cfg_scale, min_p, steps, seed, randomize_seed): | |
| if randomize_seed: seed = int(time.time()) | |
| torch.manual_seed(seed) | |
| speaker_embedding = None | |
| if speaker_audio is not None: | |
| wav, sr = torchaudio.load(speaker_audio) | |
| speaker_embedding = (model.make_speaker_embedding(wav, sr).to(device, dtype=torch.bfloat16)) | |
| emotion_tensor = torch.tensor([emotion_happy, emotion_sad, emotion_disgust, emotion_fear, emotion_surprise, emotion_anger, emotion_other, emotion_neutral], device=device, dtype=torch.bfloat16) | |
| vq_tensor = torch.tensor([clarity] * 8, device=device, dtype=torch.bfloat16).unsqueeze(0) | |
| cond_dict = make_cond_dict( | |
| text=input, | |
| language=language, | |
| speaker=speaker_embedding, | |
| emotion=emotion_tensor, | |
| vqscore_8=vq_tensor, | |
| fmax=float(fmax), | |
| pitch_std=float(pitch_std), | |
| speaking_rate=float(speaking_rate), | |
| dnsmos_ovrl=float(dnsmos_ovrl), | |
| device=device, | |
| ) | |
| conditioning = model.prepare_conditioning(cond_dict) | |
| codes = model.generate( | |
| prefix_conditioning=conditioning, | |
| max_new_tokens=int(steps), | |
| cfg_scale=float(cfg_scale), | |
| batch_size=1, | |
| sampling_params=dict(min_p=float(min_p)), | |
| ) | |
| wav_out = model.autoencoder.decode(codes).cpu().detach() | |
| sr_out = model.autoencoder.sampling_rate | |
| if wav_out.dim() == 2 and wav_out.size(0) > 1: wav_out = wav_out[0:1, :] | |
| return (sr_out, wav_out.squeeze().numpy()) | |
| # Initialize | |
| patch_cuda() | |
| with gr.Blocks() as main: | |
| text = gr.Textbox(label="text", value="hello, world!") | |
| language = gr.Dropdown(choices=supported_language_codes, value="en-us", label="language") | |
| speaker_audio = gr.Audio(label="voice reference", type="filepath") | |
| clarity_slider = gr.Slider(0.5, 0.8, 0.8, 0.01, label="clarity") | |
| steps_slider = gr.Slider(1, 3000, 300, 1, label="steps") | |
| dnsmos_slider = gr.Slider(1.0, 5.0, 5.0, 0.1, label="quality") | |
| fmax_slider = gr.Slider(0, 24000, 24000, 1, label="fmax") | |
| pitch_std_slider = gr.Slider(0.0, 300.0, 30.0, 1, label="pitch std") | |
| speaking_rate_slider = gr.Slider(5.0, 30.0, 15.0, 0.1, label="rate") | |
| cfg_scale_slider = gr.Slider(1.0, 5.0, 2.5, 0.1, label="guidance") | |
| min_p_slider = gr.Slider(0.0, 1.0, 0.05, 0.01, label="min p") | |
| with gr.Row(): | |
| e1 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="happy") | |
| e2 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="sad") | |
| e3 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="disgust") | |
| e4 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="fear") | |
| e5 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="surprise") | |
| e6 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="anger") | |
| e7 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="other") | |
| e8 = gr.Slider(0.0, 1.0, 1.0, 0.01, label="neutral") | |
| seed_number = gr.Number(label="seed", value=42, precision=0) | |
| randomize_seed_toggle = gr.Checkbox(label="randomize seed", value=True) | |
| generate_button = gr.Button("generate") | |
| output_audio = gr.Audio(label="output", type="numpy", autoplay=True) | |
| generate_button.click(fn=generate, inputs=[text, language, speaker_audio, e1, e2, e3, e4, e5, e6, e7, e8, clarity_slider, fmax_slider, pitch_std_slider, speaking_rate_slider, dnsmos_slider, cfg_scale_slider, min_p_slider, steps_slider, seed_number, randomize_seed_toggle], outputs=output_audio) | |
| main.launch() |