Spaces:
Running on Zero
Running on Zero
| from create_env import setup_dependencies | |
| setup_dependencies() | |
| import spaces | |
| import gradio as gr | |
| from util import InitModels, load_config, Examples, SpeakerManager | |
| import numpy as np | |
| import torch | |
| config = load_config("./model_config.yaml") | |
| models_configs = config.models | |
| examples_cfg = load_config("./examples.yaml") | |
| examples_maker = Examples(examples_cfg) | |
| examples = examples_maker() | |
| init_models = InitModels(models_configs) | |
| models = init_models() | |
| # Initialize speaker manager | |
| speaker_manager = SpeakerManager() | |
| def generate_embedding_gpu(audio_data): | |
| """ | |
| Generate speaker embedding from audio on GPU | |
| """ | |
| try: | |
| if audio_data is None: | |
| return "No audio provided" | |
| embedding = speaker_manager.generate_embedding(audio_data) | |
| print("Embedding generated successfully!") | |
| return speaker_manager.get_status() | |
| except Exception as e: | |
| print(f"Error generating embedding: {str(e)}") | |
| return f"Error: {str(e)}" | |
| def generate_speech_gpu(text, model_choice, mode, speaker_choice, t, top_p, rp): | |
| """ | |
| Generate speech from text using the selected model on GPU | |
| """ | |
| if not text.strip(): | |
| return None | |
| if not model_choice: | |
| return None | |
| try: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {device}") | |
| selected_model = models[model_choice] | |
| # Get speaker embedding based on mode | |
| print(f"[generate_speech_gpu] Mode: {mode}, Speaker choice: {speaker_choice}") | |
| speaker_emb = speaker_manager.get_speaker_emb(mode, speaker_choice) | |
| print(f"[generate_speech_gpu] Speaker emb type: {type(speaker_emb)}") | |
| if speaker_emb is not None: | |
| if isinstance(speaker_emb, str): | |
| print(f"[generate_speech_gpu] Speaker emb is path: {speaker_emb}") | |
| elif torch.is_tensor(speaker_emb): | |
| print(f"[generate_speech_gpu] Speaker emb is tensor: shape={speaker_emb.shape}, device={speaker_emb.device}") | |
| print(f"Generating speech with {model_choice}...") | |
| audio, _ = selected_model( | |
| text, | |
| speaker_emb=speaker_emb, | |
| temperature=t, | |
| top_p=top_p, | |
| repetition_penalty=rp | |
| ) | |
| sample_rate = 22050 | |
| print("Speech generation completed!") | |
| return (sample_rate, audio) | |
| except Exception as e: | |
| print(f"Error during generation: {str(e)}") | |
| return None | |
| # Create Gradio interface | |
| with gr.Blocks(title="😻 KaniTTS - Text to Speech", theme=gr.themes.Ocean()) as demo: | |
| gr.Markdown("# 😻 KaniTTS: Fast and Expressive Speech Generation Model") | |
| gr.Markdown("Select a model and enter text to generate emotional speech") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| model_dropdown = gr.Dropdown( | |
| choices=list(models_configs.keys()), | |
| value=list(models_configs.keys())[0], | |
| label="Selected Model" | |
| ) | |
| # Speaker mode selector | |
| speaker_mode = gr.Radio( | |
| choices=["select", "generate"], | |
| value="select", | |
| label="Speaker Mode" | |
| ) | |
| # Speaker selection (visible in "select" mode) | |
| speaker_dropdown = gr.Dropdown( | |
| choices=speaker_manager.get_speaker_names(), | |
| value=speaker_manager.get_speaker_names()[0] if speaker_manager.get_speaker_names() else None, | |
| label="Speaker", | |
| visible=True | |
| ) | |
| # Audio upload and embedding generation (visible in "generate" mode) | |
| with gr.Group(visible=False) as embedding_group: | |
| audio_input = gr.Audio( | |
| label="Upload or Record Audio (will be resampled to 16kHz)", | |
| type="numpy", | |
| sources=["upload", "microphone"], | |
| format="wav", | |
| ) | |
| with gr.Row(): | |
| run_embedding_btn = gr.Button("Run Embedding", variant="secondary") | |
| clean_embedding_btn = gr.Button("Clean", variant="stop") | |
| embedding_status = gr.Textbox( | |
| label="Embedding Status", | |
| value="No embedding generated", | |
| interactive=False | |
| ) | |
| with gr.Column(scale=1): | |
| text_input = gr.Textbox( | |
| label="Text", | |
| placeholder="Enter your text ...", | |
| lines=3, | |
| max_lines=10 | |
| ) | |
| with gr.Accordion("Settings", open=False): | |
| temp = gr.Slider( | |
| minimum=0.1, maximum=1.5, value=1.0, step=0.05, | |
| label="Temp", | |
| ) | |
| top_p = gr.Slider( | |
| minimum=0.1, maximum=1.0, value=0.95, step=0.05, | |
| label="Top P", | |
| ) | |
| rp = gr.Slider( | |
| minimum=1.0, maximum=2.0, value=1.1, step=0.05, | |
| label="Repetition Penalty", | |
| ) | |
| generate_btn = gr.Button("Run", variant="primary", size="lg") | |
| audio_output = gr.Audio( | |
| label="Generated Audio", | |
| type="numpy" | |
| ) | |
| # Toggle visibility based on speaker mode | |
| def toggle_speaker_mode(mode): | |
| if mode == "select": | |
| return gr.update(visible=True), gr.update(visible=False) | |
| else: # generate | |
| return gr.update(visible=False), gr.update(visible=True) | |
| speaker_mode.change( | |
| fn=toggle_speaker_mode, | |
| inputs=[speaker_mode], | |
| outputs=[speaker_dropdown, embedding_group] | |
| ) | |
| # Embedding generation events | |
| run_embedding_btn.click( | |
| fn=generate_embedding_gpu, | |
| inputs=[audio_input], | |
| outputs=[embedding_status] | |
| ) | |
| clean_embedding_btn.click( | |
| fn=speaker_manager.clean, | |
| inputs=[], | |
| outputs=[embedding_status] | |
| ) | |
| # GPU generation event | |
| generate_btn.click( | |
| fn=generate_speech_gpu, | |
| inputs=[text_input, model_dropdown, speaker_mode, speaker_dropdown, temp, top_p, rp], | |
| outputs=[audio_output] | |
| ) | |
| with gr.Row(): | |
| examples = examples | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[text_input, model_dropdown, speaker_mode, speaker_dropdown, temp, top_p, rp], | |
| fn=generate_speech_gpu, | |
| outputs=[audio_output], | |
| cache_examples=True, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| show_error=True | |
| ) |