Spaces:
Running
Running
| import os | |
| import tempfile | |
| import time | |
| from typing import List, Tuple | |
| import gradio as gr | |
| import torch | |
| import torchaudio | |
| # import spaces | |
| from dataclasses import dataclass | |
| from generator import Segment, load_csm_1b | |
| from huggingface_hub import login | |
| # Disable torch compile feature to avoid triton error | |
| torch._dynamo.config.suppress_errors = True | |
| # Check if GPU is available and configure the device | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {device}") | |
| # Login to Hugging Face Hub if token is available | |
| def login_huggingface(): | |
| hf_token = os.environ.get("HF_TOKEN") | |
| if hf_token: | |
| print("Logging in to Hugging Face Hub...") | |
| login(token=hf_token) | |
| print("Login successful!") | |
| else: | |
| print("HF_TOKEN not found in environment variables. Some models may not be accessible.") | |
| # Login at startup | |
| login_huggingface() | |
| # Global variables to track model state | |
| generator = None | |
| model_loaded = False | |
| # Function to load model in ZeroGPU | |
| # @spaces.GPU(duration=30) | |
| def initialize_model(): | |
| global generator, model_loaded | |
| if not model_loaded: | |
| print("Loading CSM-1B model in GPU...") | |
| generator = load_csm_1b(device="cuda") | |
| model_loaded = True | |
| print("Model loaded successfully!") | |
| return generator | |
| # Function to get the loaded model | |
| # @spaces.GPU(duration=30) | |
| def get_model(): | |
| global generator, model_loaded | |
| if not model_loaded: | |
| return initialize_model() | |
| return generator | |
| # Preload model if environment variable is set | |
| def preload_model_if_needed(): | |
| if os.environ.get("PRELOAD_MODEL", "").lower() in ("true", "1", "yes"): | |
| print("PRELOAD_MODEL is set. Attempting to preload model...") | |
| try: | |
| # We can't directly call initialize_model() here because it's decorated with @spaces.GPU | |
| # Instead, we'll set a flag that will be checked when the first request comes in | |
| global model_loaded | |
| model_loaded = False | |
| print("Model will be loaded on first request.") | |
| except Exception as e: | |
| print(f"Error during model preloading setup: {e}") | |
| else: | |
| print("PRELOAD_MODEL is not set. Model will be loaded on demand.") | |
| # Call preload function at startup | |
| preload_model_if_needed() | |
| # Function to convert audio to tensor | |
| def audio_to_tensor(audio_path: str) -> Tuple[torch.Tensor, int]: | |
| waveform, sample_rate = torchaudio.load(audio_path) | |
| waveform = waveform.mean(dim=0) # Convert stereo to mono if needed | |
| return waveform, sample_rate | |
| # Function to save audio tensor to file | |
| def save_audio(audio_tensor: torch.Tensor, sample_rate: int) -> str: | |
| # Lưu file vào thư mục hiện tại hoặc thư mục files mà Gradio mặc định sử dụng | |
| output_path = f"csm1b_output_{int(time.time())}.wav" | |
| torchaudio.save(output_path, audio_tensor.unsqueeze(0), sample_rate) | |
| return output_path | |
| # Function to generate speech from text using ZeroGPU | |
| # @spaces.GPU(duration=30) | |
| def generate_speech( | |
| text: str, | |
| speaker_id: int, | |
| context_audio_path1: str = None, | |
| context_text1: str = None, | |
| context_speaker1: int = 0, | |
| context_audio_path2: str = None, | |
| context_text2: str = None, | |
| context_speaker2: int = 1, | |
| max_duration_ms: float = 30000, | |
| temperature: float = 0.9, | |
| top_k: int = 50, | |
| progress=gr.Progress() | |
| ) -> str: | |
| try: | |
| # Get the loaded model | |
| generator = get_model() | |
| # Prepare context | |
| context = [] | |
| progress(0.1, "Processing context...") | |
| # Process context 1 | |
| if context_audio_path1 and context_text1: | |
| waveform, sample_rate = audio_to_tensor(context_audio_path1) | |
| # Resample if needed | |
| if sample_rate != generator.sample_rate: | |
| waveform = torchaudio.functional.resample(waveform, orig_freq=sample_rate, new_freq=generator.sample_rate) | |
| context.append(Segment(speaker=context_speaker1, text=context_text1, audio=waveform)) | |
| # Process context 2 | |
| if context_audio_path2 and context_text2: | |
| waveform, sample_rate = audio_to_tensor(context_audio_path2) | |
| # Resample if needed | |
| if sample_rate != generator.sample_rate: | |
| waveform = torchaudio.functional.resample(waveform, orig_freq=sample_rate, new_freq=generator.sample_rate) | |
| context.append(Segment(speaker=context_speaker2, text=context_text2, audio=waveform)) | |
| progress(0.3, "Generating audio...") | |
| # Generate audio from text | |
| audio = generator.generate( | |
| text=text, | |
| speaker=speaker_id, | |
| context=context, | |
| max_audio_length_ms=max_duration_ms, | |
| # temperature=temperature, | |
| # topk=top_k | |
| ) | |
| progress(0.8, "Saving audio...") | |
| # Save audio to file | |
| # output_path = save_audio(audio, generator.sample_rate) | |
| output_path = f"csm1b_output_{int(time.time())}.wav" | |
| progress(1.0, "Completed!") | |
| return output_path | |
| except Exception as e: | |
| # Handle ZeroGPU quota exceeded error | |
| error_message = str(e) | |
| if "GPU quota exceeded" in error_message: | |
| # Extract wait time from error message | |
| import re | |
| wait_time_match = re.search(r"Try again in (\d+:\d+:\d+)", error_message) | |
| wait_time = wait_time_match.group(1) if wait_time_match else "some time" | |
| return f"GPU quota exceeded. Please try again in {wait_time}." | |
| return f"GPU error: {error_message}" | |
| except Exception as e: | |
| return f"Error generating speech: {str(e)}" | |
| # Function to generate simple speech without context | |
| # @spaces.GPU(duration=30) | |
| def generate_speech_simple( | |
| text: str, | |
| speaker_id: int, | |
| max_duration_ms: float = 30000, | |
| temperature: float = 0.9, | |
| top_k: int = 50, | |
| progress=gr.Progress() | |
| ) -> str: | |
| try: | |
| # Get the loaded model | |
| generator = get_model() | |
| progress(0.3, "Generating audio...") | |
| # Generate audio from text | |
| audio = generator.generate( | |
| text=text, | |
| speaker=speaker_id, | |
| context=[], # No context | |
| max_audio_length_ms=max_duration_ms, | |
| # temperature=temperature, | |
| # topk=top_k | |
| ) | |
| progress(0.8, "Saving audio...") | |
| # Save audio to file | |
| # output_path = save_audio(audio, generator.sample_rate) | |
| output_path = f"csm1b_output_{int(time.time())}.wav" | |
| torchaudio.save(output_path, audio.unsqueeze(0).cpu(), generator.sample_rate) | |
| print(f"Audio saved to {output_path}") | |
| progress(1.0, "Completed!") | |
| return output_path | |
| except Exception as e: | |
| # Handle ZeroGPU quota exceeded error | |
| error_message = str(e) | |
| if "GPU quota exceeded" in error_message: | |
| # Extract wait time from error message | |
| import re | |
| wait_time_match = re.search(r"Try again in (\d+:\d+:\d+)", error_message) | |
| wait_time = wait_time_match.group(1) if wait_time_match else "some time" | |
| return f"GPU quota exceeded. Please try again in {wait_time}." | |
| return f"GPU error: {error_message}" | |
| except Exception as e: | |
| return f"Error generating speech: {str(e)}" | |
| # Create Gradio interface | |
| def create_demo(): | |
| with gr.Blocks(title="CSM-1B Text-to-Speech") as demo: | |
| gr.Markdown("# CSM-1B Text-to-Speech Demo") | |
| gr.Markdown("CSM-1B (Collaborative Speech Model) is an advanced text-to-speech model capable of generating natural-sounding speech from text.") | |
| with gr.Tab("Simple Audio Generation"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| text_input = gr.Textbox( | |
| label="Text to convert to speech", | |
| placeholder="Enter the text you want to convert to speech...", | |
| lines=5 | |
| ) | |
| speaker_id = gr.Number( | |
| label="Speaker ID", | |
| value=0, | |
| precision=0, | |
| minimum=0, | |
| maximum=10 | |
| ) | |
| with gr.Row(): | |
| max_duration = gr.Slider( | |
| label="Maximum Duration (ms)", | |
| minimum=1000, | |
| maximum=90000, | |
| value=30000, | |
| step=1000 | |
| ) | |
| # temperature = gr.Slider( | |
| # label="Temperature", | |
| # minimum=0.1, | |
| # maximum=1.5, | |
| # value=0.9, | |
| # step=0.1 | |
| # ) | |
| # top_k = gr.Slider( | |
| # label="Top-K", | |
| # minimum=1, | |
| # maximum=100, | |
| # value=50, | |
| # step=1 | |
| # ) | |
| generate_btn = gr.Button("Generate Audio") | |
| with gr.Column(): | |
| output_audio = gr.Audio(label="Output Audio", type="filepath", autoplay=True) | |
| with gr.Tab("Audio Generation with Context"): | |
| gr.Markdown("This feature allows you to provide audio clips and text as context to help the model generate more appropriate speech.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| context_text1 = gr.Textbox(label="Context Text 1", lines=2) | |
| context_audio1 = gr.Audio(label="Context Audio 1", type="filepath") | |
| context_speaker1 = gr.Number(label="Speaker ID 1", value=0, precision=0) | |
| context_text2 = gr.Textbox(label="Context Text 2", lines=2) | |
| context_audio2 = gr.Audio(label="Context Audio 2", type="filepath") | |
| context_speaker2 = gr.Number(label="Speaker ID 2", value=1, precision=0) | |
| text_input_context = gr.Textbox( | |
| label="Text to convert to speech", | |
| placeholder="Enter the text you want to convert to speech...", | |
| lines=3 | |
| ) | |
| speaker_id_context = gr.Number( | |
| label="Speaker ID", | |
| value=0, | |
| precision=0 | |
| ) | |
| with gr.Row(): | |
| max_duration_context = gr.Slider( | |
| label="Maximum Duration (ms)", | |
| minimum=1000, | |
| maximum=90000, | |
| value=30000, | |
| step=1000 | |
| ) | |
| # temperature_context = gr.Slider( | |
| # label="Temperature", | |
| # minimum=0.1, | |
| # maximum=1.5, | |
| # value=0.9, | |
| # step=0.1 | |
| # ) | |
| # top_k_context = gr.Slider( | |
| # label="Top-K", | |
| # minimum=1, | |
| # maximum=100, | |
| # value=50, | |
| # step=1 | |
| # ) | |
| generate_context_btn = gr.Button("Generate Audio with Context") | |
| with gr.Column(): | |
| output_audio_context = gr.Audio(label="Output Audio", type="filepath", autoplay=True) | |
| # Add Hugging Face configuration tab | |
| with gr.Tab("Configuration"): | |
| gr.Markdown("### Hugging Face Token Configuration") | |
| gr.Markdown(""" | |
| To use the CSM-1B model, you need access to the model on Hugging Face. | |
| You can configure your token by: | |
| 1. Create a token at [Hugging Face Settings](https://huggingface.co/settings/tokens) | |
| 2. Set the `HF_TOKEN` environment variable with your token value | |
| Note: In Hugging Face Spaces, you can set environment variables in the Space Settings. | |
| """) | |
| hf_token_input = gr.Textbox( | |
| label="Hugging Face Token (Only for this session)", | |
| placeholder="Enter your token...", | |
| type="password" | |
| ) | |
| def set_token(token): | |
| if token: | |
| os.environ["HF_TOKEN"] = token | |
| login(token=token) | |
| return "Token set successfully! You can now load the model." | |
| return "Invalid token. Please enter a valid token." | |
| set_token_btn = gr.Button("Set Token") | |
| token_status = gr.Textbox(label="Status", interactive=False) | |
| set_token_btn.click(fn=set_token, inputs=hf_token_input, outputs=token_status) | |
| # Add GPU information tab | |
| with gr.Tab("GPU Information"): | |
| gr.Markdown("### About ZeroGPU") | |
| gr.Markdown(""" | |
| This application uses Hugging Face Spaces' ZeroGPU to optimize GPU usage. | |
| ZeroGPU helps free up GPU memory when not in use, saving resources and improving performance. | |
| When you generate audio, the GPU will be used automatically and released after completion. | |
| Note: In the ZeroGPU environment, CUDA is not initialized in the main process, but only in functions with the @spaces.GPU decorator. | |
| """) | |
| gr.Markdown("### GPU Quota Information") | |
| gr.Markdown(""" | |
| Hugging Face Spaces has GPU quota limitations: | |
| - Each GPU operation has a default duration of 60 seconds | |
| - We've reduced this to 30 seconds for audio generation and 10 seconds for GPU checks | |
| - If you exceed your quota, you'll need to wait for it to reset (usually a few hours) | |
| - For better performance, try generating shorter audio clips | |
| If you encounter a "GPU quota exceeded" error, please wait for the specified time and try again. | |
| """) | |
| # @spaces.GPU(duration=10) | |
| def check_gpu(): | |
| if torch.cuda.is_available(): | |
| gpu_name = torch.cuda.get_device_name(0) | |
| gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3) | |
| return f"GPU: {gpu_name}\nMemory: {gpu_memory:.2f} GB" | |
| else: | |
| return "No GPU found. The application will run on CPU." | |
| check_gpu_btn = gr.Button("Check GPU") | |
| gpu_info = gr.Textbox(label="GPU Information", interactive=False) | |
| check_gpu_btn.click(fn=check_gpu, inputs=None, outputs=gpu_info) | |
| # Add model loading button | |
| load_model_btn = gr.Button("Load Model") | |
| model_status = gr.Textbox(label="Model Status", interactive=False) | |
| # @spaces.GPU(duration=10) | |
| def load_model_and_report(): | |
| global model_loaded | |
| if model_loaded: | |
| return "Model has already been loaded!" | |
| else: | |
| initialize_model() | |
| return "Model loaded successfully!" | |
| load_model_btn.click(fn=load_model_and_report, inputs=None, outputs=model_status) | |
| # Connect components | |
| generate_btn.click( | |
| fn=generate_speech_simple, | |
| inputs=[ | |
| text_input, | |
| speaker_id, | |
| max_duration, | |
| # temperature, | |
| # top_k | |
| ], | |
| outputs=output_audio | |
| ) | |
| generate_context_btn.click( | |
| fn=generate_speech, | |
| inputs=[ | |
| text_input_context, | |
| speaker_id_context, | |
| context_audio1, | |
| context_text1, | |
| context_speaker1, | |
| context_audio2, | |
| context_text2, | |
| context_speaker2, | |
| max_duration_context, | |
| # temperature_context, | |
| # top_k_context | |
| ], | |
| outputs=output_audio_context | |
| ) | |
| return demo | |
| # Launch the application | |
| if __name__ == "__main__": | |
| demo = create_demo() | |
| demo.queue().launch() |