app.py

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

# Kiểm tra xem có GPU không và cấu hình thiết bị phù hợp
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Sử dụng thiết bị: {device}")

# Đăng nhập vào Hugging Face Hub nếu có token
def login_huggingface():
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        print("Đang đăng nhập vào Hugging Face Hub...")
        login(token=hf_token)
        print("Đã đăng nhập thành công!")
    else:
        print("Không tìm thấy HF_TOKEN trong biến môi trường. Một số mô hình có thể không truy cập được.")

# Đăng nhập khi khởi động
login_huggingface()

# Tải mô hình CSM-1B
generator = None

def load_model():
    global generator
    if generator is None:
        print("Đang tải mô hình CSM-1B...")
        generator = load_csm_1b(device=device)
        print("Đã tải xong mô hình!")
    return generator

# Hàm chuyển đổi âm thanh thành tensor
def audio_to_tensor(audio_path: str) -> Tuple[torch.Tensor, int]:
    waveform, sample_rate = torchaudio.load(audio_path)
    waveform = waveform.mean(dim=0)  # Chuyển stereo thành mono nếu cần
    return waveform, sample_rate

# Hàm lưu tensor âm thanh thành file
def save_audio(audio_tensor: torch.Tensor, sample_rate: int) -> str:
    temp_dir = tempfile.gettempdir()
    output_path = os.path.join(temp_dir, f"csm1b_output_{int(time.time())}.wav")
    torchaudio.save(output_path, audio_tensor.unsqueeze(0), sample_rate)
    return output_path

# Hàm tạo âm thanh từ văn bản sử dụng ZeroGPU
@spaces.GPU
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:
    # Tải mô hình nếu chưa tải
    generator = load_model()
    
    # Chuẩn bị ngữ cảnh (context)
    context = []
    progress(0.1, "Đang xử lý ngữ cảnh...")
    
    # Xử lý ngữ cảnh 1
    if context_audio_path1 and context_text1:
        waveform, sample_rate = audio_to_tensor(context_audio_path1)
        # Resample nếu cần
        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))
    
    # Xử lý ngữ cảnh 2
    if context_audio_path2 and context_text2:
        waveform, sample_rate = audio_to_tensor(context_audio_path2)
        # Resample nếu cần
        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, "Đang tạo âm thanh...")
    # Tạo âm thanh từ văn bản
    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, "Đang lưu âm thanh...")
    # Lưu âm thanh thành file
    output_path = save_audio(audio, generator.sample_rate)
    
    progress(1.0, "Hoàn thành!")
    return output_path

# Hàm tạo âm thanh đơn giản không có ngữ cảnh
@spaces.GPU
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:
    # Tải mô hình nếu chưa tải
    generator = load_model()
    
    progress(0.3, "Đang tạo âm thanh...")
    # Tạo âm thanh từ văn bản
    audio = generator.generate(
        text=text,
        speaker=speaker_id,
        context=[],  # Không có ngữ cảnh
        max_audio_length_ms=max_duration_ms,
        temperature=temperature,
        topk=top_k
    )
    
    progress(0.8, "Đang lưu âm thanh...")
    # Lưu âm thanh thành file
    output_path = save_audio(audio, generator.sample_rate)
    
    progress(1.0, "Hoàn thành!")
    return output_path

# Tạo giao diện Gradio
def create_demo():
    with gr.Blocks(title="CSM-1B Text-to-Speech") as demo:
        gr.Markdown("# CSM-1B Text-to-Speech Demo")
        gr.Markdown("Mô hình CSM-1B (Collaborative Speech Model) là một mô hình text-to-speech tiên tiến có khả năng tạo giọng nói tự nhiên từ văn bản.")
        
        with gr.Tab("Tạo âm thanh đơn giản"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Văn bản cần chuyển thành giọng nói",
                        placeholder="Nhập văn bản bạn muốn chuyển thành giọng nói...",
                        lines=5
                    )
                    speaker_id = gr.Number(
                        label="ID người nói",
                        value=0,
                        precision=0,
                        minimum=0,
                        maximum=10
                    )
                    
                    with gr.Row():
                        max_duration = gr.Slider(
                            label="Thời lượng tối đa (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("Tạo âm thanh")
                
                with gr.Column():
                    output_audio = gr.Audio(label="Âm thanh đầu ra", type="filepath")
        
        with gr.Tab("Tạo âm thanh với ngữ cảnh"):
            gr.Markdown("Tính năng này cho phép bạn cung cấp các đoạn âm thanh và văn bản làm ngữ cảnh để mô hình tạo ra âm thanh phù hợp hơn.")
            
            with gr.Row():
                with gr.Column():
                    context_text1 = gr.Textbox(label="Văn bản ngữ cảnh 1", lines=2)
                    context_audio1 = gr.Audio(label="Âm thanh ngữ cảnh 1", type="filepath")
                    context_speaker1 = gr.Number(label="ID người nói 1", value=0, precision=0)
                    
                    context_text2 = gr.Textbox(label="Văn bản ngữ cảnh 2", lines=2)
                    context_audio2 = gr.Audio(label="Âm thanh ngữ cảnh 2", type="filepath")
                    context_speaker2 = gr.Number(label="ID người nói 2", value=1, precision=0)
                    
                    text_input_context = gr.Textbox(
                        label="Văn bản cần chuyển thành giọng nói",
                        placeholder="Nhập văn bản bạn muốn chuyển thành giọng nói...",
                        lines=3
                    )
                    speaker_id_context = gr.Number(
                        label="ID người nói",
                        value=0,
                        precision=0
                    )
                    
                    with gr.Row():
                        max_duration_context = gr.Slider(
                            label="Thời lượng tối đa (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("Tạo âm thanh với ngữ cảnh")
                
                with gr.Column():
                    output_audio_context = gr.Audio(label="Âm thanh đầu ra", type="filepath")
        
        # Thêm tab cấu hình Hugging Face
        with gr.Tab("Cấu hình"):
            gr.Markdown("### Cấu hình Hugging Face Token")
            gr.Markdown("""
            Để sử dụng mô hình CSM-1B, bạn cần có quyền truy cập vào mô hình trên Hugging Face.
            
            Bạn có thể cấu hình token của mình bằng cách:
            1. Tạo token tại [Hugging Face Settings](https://huggingface.co/settings/tokens)
            2. Đặt biến môi trường `HF_TOKEN` với giá trị là token của bạn
            
            Lưu ý: Trong Hugging Face Spaces, bạn có thể đặt biến môi trường trong phần Cài đặt của Space.
            """)
            
            hf_token_input = gr.Textbox(
                label="Hugging Face Token (Chỉ sử dụng trong phiên này)",
                placeholder="Nhập token của bạn...",
                type="password"
            )
            
            def set_token(token):
                if token:
                    os.environ["HF_TOKEN"] = token
                    login(token=token)
                    return "Đã đặt token thành công! Bạn có thể tải mô hình bây giờ."
                return "Token không hợp lệ. Vui lòng nhập token hợp lệ."
            
            set_token_btn = gr.Button("Đặt Token")
            token_status = gr.Textbox(label="Trạng thái", interactive=False)
            
            set_token_btn.click(fn=set_token, inputs=hf_token_input, outputs=token_status)
            
        # Thêm tab thông tin về ZeroGPU
        with gr.Tab("Thông tin GPU"):
            gr.Markdown("### Thông tin về ZeroGPU")
            gr.Markdown("""
            Ứng dụng này sử dụng ZeroGPU của Hugging Face Spaces để tối ưu việc sử dụng GPU.
            
            ZeroGPU giúp giải phóng bộ nhớ GPU khi không sử dụng, giúp tiết kiệm tài nguyên và cải thiện hiệu suất.
            
            Khi bạn tạo âm thanh, GPU sẽ được sử dụng tự động và giải phóng sau khi hoàn thành.
            """)
            
            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}\nBộ nhớ: {gpu_memory:.2f} GB"
                else:
                    return "Không tìm thấy GPU. Ứng dụng sẽ chạy trên CPU."
            
            check_gpu_btn = gr.Button("Kiểm tra GPU")
            gpu_info = gr.Textbox(label="Thông tin GPU", interactive=False)
            
            check_gpu_btn.click(fn=check_gpu, inputs=None, outputs=gpu_info)
        
        # Kết nối các thành phần
        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
        )
        
        # Tải mô hình khi khởi động
        demo.load(fn=load_model)
    
    return demo

# Khởi chạy ứng dụng
if __name__ == "__main__":
    demo = create_demo()
    demo.queue().launch()