fix

.gitignore

@@ -0,0 +1,54 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# PyCharm
+.idea/
+
+# VS Code
+.vscode/
+
+# Temporary files
+*.tmp
+*.wav
+*.mp3
+*.ogg
+temp/
+
+# Logs
+logs/
+*.log
+
+# HuggingFace
+.cache/
+huggingface/ 

README.md

@@ -1,12 +1,102 @@
 ---
-title: Csm 1b Gradio V2
-emoji: 📊
+title: CSM-1B Gradio Demo
+emoji: 🎙️
 colorFrom: indigo
 colorTo: green
 sdk: gradio
-sdk_version: 5.21.0
+sdk_version: 4.26.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# CSM-1B Text-to-Speech Demo
+
+Ứng dụng này sử dụng mô hình CSM-1B (Collaborative Speech Model) để chuyển đổi văn bản thành giọng nói với chất lượng cao.
+
+## Tính năng
+
+- **Tạo âm thanh đơn giản**: Chuyển đổi văn bản thành giọng nói với các tùy chọn về ID người nói, thời lượng, temperature và top-k.
+- **Tạo âm thanh với ngữ cảnh**: 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.
+- **Tối ưu GPU**: Sử dụng ZeroGPU của Hugging Face Spaces để tối ưu việc sử dụng GPU.
+
+## Cài đặt và Cấu hình
+
+### Yêu cầu truy cập
+
+Để sử dụng mô hình CSM-1B, bạn cần có quyền truy cập vào các mô hình sau trên Hugging Face:
+
+- [meta-llama/Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B)
+- [sesame/csm-1b](https://huggingface.co/sesame/csm-1b)
+
+### Cấu hình Hugging Face Token
+
+1. Tạo tài khoản Hugging Face nếu bạn chưa có.
+2. Truy cập vào [Hugging Face Settings](https://huggingface.co/settings/tokens) để tạo token.
+3. Yêu cầu quyền truy cập vào các mô hình nếu cần.
+4. Đặt biến môi trường `HF_TOKEN` với giá trị là token của bạn:
+   ```bash
+   export HF_TOKEN=your_token_here
+   ```
+5. Hoặc bạn có thể nhập token trực tiếp trong tab "Cấu hình" của ứng dụng.
+
+### Cài đặt
+
+```bash
+git clone https://github.com/yourusername/csm-1b-gradio.git
+cd csm-1b-gradio
+pip install -r requirements.txt
+```
+
+## Cách sử dụng
+
+1. Khởi động ứng dụng:
+   ```bash
+   python app.py
+   ```
+2. Mở trình duyệt web và truy cập địa chỉ được hiển thị (thường là http://127.0.0.1:7860).
+3. Nhập văn bản bạn muốn chuyển thành giọng nói.
+4. Chọn ID người nói (từ 0-10).
+5. Điều chỉnh các tham số như thời lượng tối đa, temperature và top-k.
+6. Nhấn nút "Tạo âm thanh" để tạo giọng nói.
+
+## Thông tin về mô hình
+
+CSM-1B là một mô hình text-to-speech tiên tiến được phát triển bởi Sesame AI Labs. Mô hình này có khả năng tạo giọng nói tự nhiên từ văn bản với nhiều giọng nói khác nhau.
+
+## ZeroGPU
+
+Ứ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.
+
+```python
+import spaces
+
+@spaces.GPU
+def my_gpu_function():
+    # Hàm này sẽ chỉ sử dụng GPU khi được gọi
+    # và giải phóng GPU sau khi hoàn thành
+    pass
+```
+
+Khi triển khai trên Hugging Face Spaces, ZeroGPU sẽ tự động quản lý việc sử dụng GPU, giúp ứng dụng hoạt động hiệu quả hơn.
+
+## Lưu ý
+
+- Mô hình này sử dụng watermarking để đánh dấu âm thanh được tạo ra bởi AI.
+- Thời gian tạo âm thanh phụ thuộc vào độ dài văn bản và cấu hình phần cứng.
+- Bạn cần có quyền truy cập vào mô hình CSM-1B trên Hugging Face để sử dụng ứng dụng này.
+
+## Triển khai trên Hugging Face Spaces
+
+Để triển khai ứng dụng này trên Hugging Face Spaces:
+
+1. Tạo một Space mới trên Hugging Face với SDK là Gradio.
+2. Tải lên tất cả các file của dự án.
+3. Trong phần cài đặt của Space, thêm biến môi trường `HF_TOKEN` với giá trị là token của bạn.
+4. Chọn cấu hình phần cứng phù hợp (khuyến nghị sử dụng GPU).
+
+## Tài nguyên
+
+- [GitHub Repository](https://github.com/SesameAILabs/csm-1b)
+- [Hugging Face Model](https://huggingface.co/sesame/csm-1b)
+- [Hugging Face Space Demo](https://huggingface.co/spaces/sesame/csm-1b)
+- [Hugging Face Spaces ZeroGPU](https://huggingface.co/docs/hub/spaces-sdks-docker-zero-gpu)

app.py

@@ -0,0 +1,294 @@
+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_files: List[Tuple[str, str, 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()
+    
+    # Chuẩn bị ngữ cảnh (context)
+    context = []
+    progress(0.1, "Đang xử lý ngữ cảnh...")
+    
+    for audio_file, text_content, speaker in context_audio_files:
+        if audio_file and text_content:
+            waveform, sample_rate = audio_to_tensor(audio_file)
+            # 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=speaker, text=text_content, 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
+
+# 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,
+            inputs=[
+                text_input,
+                speaker_id,
+                gr.State([]),  # Không có ngữ cảnh
+                max_duration,
+                temperature,
+                top_k
+            ],
+            outputs=output_audio
+        )
+        
+        generate_context_btn.click(
+            fn=generate_speech,
+            inputs=[
+                text_input_context,
+                speaker_id_context,
+                gr.State([
+                    (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() 

generator.py

@@ -0,0 +1,186 @@
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download, login
+from models import Model
+from moshi.models import loaders
+from tokenizers.processors import TemplateProcessing
+from transformers import AutoTokenizer
+from watermarking import CSM_1B_GH_WATERMARK, load_watermarker, watermark
+
+
+@dataclass
+class Segment:
+    speaker: int
+    text: str
+    # (num_samples,), sample_rate = 24_000
+    audio: torch.Tensor
+
+
+def load_llama3_tokenizer():
+    """
+    https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992
+    """
+    tokenizer_name = "meta-llama/Llama-3.2-1B"
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    bos = tokenizer.bos_token
+    eos = tokenizer.eos_token
+    tokenizer._tokenizer.post_processor = TemplateProcessing(
+        single=f"{bos}:0 $A:0 {eos}:0",
+        pair=f"{bos}:0 $A:0 {eos}:0 {bos}:1 $B:1 {eos}:1",
+        special_tokens=[(f"{bos}", tokenizer.bos_token_id), (f"{eos}", tokenizer.eos_token_id)],
+    )
+
+    return tokenizer
+
+
+class Generator:
+    def __init__(
+        self,
+        model: Model,
+    ):
+        self._model = model
+        self._model.setup_caches(1)
+
+        self._text_tokenizer = load_llama3_tokenizer()
+
+        device = next(model.parameters()).device
+        mimi_weight = hf_hub_download(loaders.DEFAULT_REPO, loaders.MIMI_NAME)
+        mimi = loaders.get_mimi(mimi_weight, device=device)
+        mimi.set_num_codebooks(32)
+        self._audio_tokenizer = mimi
+
+        self._watermarker = load_watermarker(device=device)
+
+        self.sample_rate = mimi.sample_rate
+        self.device = device
+
+    def _tokenize_text_segment(self, text: str, speaker: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        frame_tokens = []
+        frame_masks = []
+
+        text_tokens = self._text_tokenizer.encode(f"[{speaker}]{text}")
+        text_frame = torch.zeros(len(text_tokens), 33).long()
+        text_frame_mask = torch.zeros(len(text_tokens), 33).bool()
+        text_frame[:, -1] = torch.tensor(text_tokens)
+        text_frame_mask[:, -1] = True
+
+        frame_tokens.append(text_frame.to(self.device))
+        frame_masks.append(text_frame_mask.to(self.device))
+
+        return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0)
+
+    def _tokenize_audio(self, audio: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        frame_tokens = []
+        frame_masks = []
+
+        # (K, T)
+        audio = audio.to(self.device)
+        audio_tokens = self._audio_tokenizer.encode(audio.unsqueeze(0).unsqueeze(0))[0]
+        # add EOS frame
+        eos_frame = torch.zeros(audio_tokens.size(0), 1).to(self.device)
+        audio_tokens = torch.cat([audio_tokens, eos_frame], dim=1)
+
+        audio_frame = torch.zeros(audio_tokens.size(1), 33).long().to(self.device)
+        audio_frame_mask = torch.zeros(audio_tokens.size(1), 33).bool().to(self.device)
+        audio_frame[:, :-1] = audio_tokens.transpose(0, 1)
+        audio_frame_mask[:, :-1] = True
+
+        frame_tokens.append(audio_frame)
+        frame_masks.append(audio_frame_mask)
+
+        return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0)
+
+    def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Returns:
+            (seq_len, 33), (seq_len, 33)
+        """
+        text_tokens, text_masks = self._tokenize_text_segment(segment.text, segment.speaker)
+        audio_tokens, audio_masks = self._tokenize_audio(segment.audio)
+
+        return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+    ) -> torch.Tensor:
+        self._model.reset_caches()
+
+        max_audio_frames = int(max_audio_length_ms / 80)
+        tokens, tokens_mask = [], []
+        for segment in context:
+            segment_tokens, segment_tokens_mask = self._tokenize_segment(segment)
+            tokens.append(segment_tokens)
+            tokens_mask.append(segment_tokens_mask)
+
+        gen_segment_tokens, gen_segment_tokens_mask = self._tokenize_text_segment(text, speaker)
+        tokens.append(gen_segment_tokens)
+        tokens_mask.append(gen_segment_tokens_mask)
+
+        prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
+        prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
+
+        samples = []
+        curr_tokens = prompt_tokens.unsqueeze(0)
+        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+
+        max_seq_len = 2048 - max_audio_frames
+        if curr_tokens.size(1) >= max_seq_len:
+            raise ValueError(f"Inputs too long, must be below max_seq_len - max_audio_frames: {max_seq_len}")
+
+        for _ in range(max_audio_frames):
+            sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
+            if torch.all(sample == 0):
+                break  # eos
+
+            samples.append(sample)
+
+            curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
+            curr_tokens_mask = torch.cat(
+                [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
+            ).unsqueeze(1)
+            curr_pos = curr_pos[:, -1:] + 1
+
+        audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+
+        # This applies an imperceptible watermark to identify audio as AI-generated.
+        # Watermarking ensures transparency, dissuades misuse, and enables traceability.
+        # Please be a responsible AI citizen and keep the watermarking in place.
+        # If using CSM 1B in another application, use your own private key and keep it secret.
+        audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, CSM_1B_GH_WATERMARK)
+        audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate)
+
+        return audio
+
+
+def load_csm_1b(device: str = "cuda") -> Generator:
+    """
+    Tải mô hình CSM-1B từ Hugging Face Hub.
+    
+    Args:
+        device: Thiết bị để chạy mô hình (cuda hoặc cpu)
+        
+    Returns:
+        Generator: Đối tượng Generator để tạo âm thanh từ văn bản
+    """
+    try:
+        model = Model.from_pretrained("sesame/csm-1b")
+        model.to(device=device, dtype=torch.bfloat16)
+        
+        generator = Generator(model)
+        return generator
+    except Exception as e:
+        print(f"Lỗi khi tải mô hình: {e}")
+        print("Vui lòng kiểm tra xem bạn đã đăng nhập vào Hugging Face Hub chưa.")
+        print("Bạn có thể cần phải yêu cầu quyền truy cập vào mô hình tại: https://huggingface.co/sesame/csm-1b")
+        raise e

hf_requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.0.0
+torchaudio>=2.0.0
+tokenizers>=0.13.0
+transformers>=4.30.0
+huggingface_hub>=0.16.0
+moshi>=0.2.2
+torchtune>=0.4.0
+torchao>=0.9.0
+silentcipher @ git+https://github.com/SesameAILabs/silentcipher@master
+gradio>=4.13.0
+huggingface-hub-spaces>=0.19.0 

models.py

@@ -0,0 +1,203 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torchtune
+from huggingface_hub import PyTorchModelHubMixin
+from torchtune.models import llama3_2
+
+
+def llama3_2_1B() -> torchtune.modules.transformer.TransformerDecoder:
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=16,
+        num_heads=32,
+        num_kv_heads=8,
+        embed_dim=2048,
+        max_seq_len=2048,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+def llama3_2_100M() -> torchtune.modules.transformer.TransformerDecoder:
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=4,
+        num_heads=8,
+        num_kv_heads=2,
+        embed_dim=1024,
+        max_seq_len=2048,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+FLAVORS = {
+    "llama-1B": llama3_2_1B,
+    "llama-100M": llama3_2_100M,
+}
+
+
+def _prepare_transformer(model):
+    embed_dim = model.tok_embeddings.embedding_dim
+    model.tok_embeddings = nn.Identity()
+    model.output = nn.Identity()
+    return model, embed_dim
+
+
+def _create_causal_mask(seq_len: int, device: torch.device):
+    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
+
+
+def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor):
+    """
+    Args:
+        mask: (max_seq_len, max_seq_len)
+        input_pos: (batch_size, seq_len)
+
+    Returns:
+        (batch_size, seq_len, max_seq_len)
+    """
+    r = mask[input_pos, :]
+    return r
+
+
+def _multinomial_sample_one_no_sync(probs):  # Does multinomial sampling without a cuda synchronization
+    q = torch.empty_like(probs).exponential_(1)
+    return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def sample_topk(logits: torch.Tensor, topk: int, temperature: float):
+    logits = logits / temperature
+
+    filter_value: float = -float("Inf")
+    indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None]
+    scores_processed = logits.masked_fill(indices_to_remove, filter_value)
+    scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1)
+    probs = torch.nn.functional.softmax(scores_processed, dim=-1)
+
+    sample_token = _multinomial_sample_one_no_sync(probs)
+    return sample_token
+
+
+@dataclass
+class ModelArgs:
+    backbone_flavor: str
+    decoder_flavor: str
+    text_vocab_size: int
+    audio_vocab_size: int
+    audio_num_codebooks: int
+
+
+class Model(
+    nn.Module,
+    PyTorchModelHubMixin,
+    repo_url="https://github.com/SesameAILabs/csm",
+    pipeline_tag="text-to-speech",
+    license="apache-2.0",
+):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+
+        self.backbone, backbone_dim = _prepare_transformer(FLAVORS[config.backbone_flavor]())
+        self.decoder, decoder_dim = _prepare_transformer(FLAVORS[config.decoder_flavor]())
+
+        self.text_embeddings = nn.Embedding(config.text_vocab_size, backbone_dim)
+        self.audio_embeddings = nn.Embedding(config.audio_vocab_size * config.audio_num_codebooks, backbone_dim)
+
+        self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False)
+        self.codebook0_head = nn.Linear(backbone_dim, config.audio_vocab_size, bias=False)
+        self.audio_head = nn.Parameter(torch.empty(config.audio_num_codebooks - 1, decoder_dim, config.audio_vocab_size))
+
+    def setup_caches(self, max_batch_size: int) -> torch.Tensor:
+        """Setup KV caches and return a causal mask."""
+        dtype = next(self.parameters()).dtype
+        device = next(self.parameters()).device
+
+        with device:
+            self.backbone.setup_caches(max_batch_size, dtype)
+            self.decoder.setup_caches(max_batch_size, dtype, decoder_max_seq_len=self.config.audio_num_codebooks)
+
+        self.register_buffer("backbone_causal_mask", _create_causal_mask(self.backbone.max_seq_len, device))
+        self.register_buffer("decoder_causal_mask", _create_causal_mask(self.config.audio_num_codebooks, device))
+
+    def generate_frame(
+        self,
+        tokens: torch.Tensor,
+        tokens_mask: torch.Tensor,
+        input_pos: torch.Tensor,
+        temperature: float,
+        topk: int,
+    ) -> torch.Tensor:
+        """
+        Args:
+            tokens: (batch_size, seq_len, audio_num_codebooks+1)
+            tokens_mask: (batch_size, seq_len, audio_num_codebooks+1)
+            input_pos: (batch_size, seq_len) positions for each token
+            mask: (batch_size, seq_len, max_seq_len
+
+        Returns:
+            (batch_size, audio_num_codebooks) sampled tokens
+        """
+        dtype = next(self.parameters()).dtype
+        b, s, _ = tokens.size()
+
+        assert self.backbone.caches_are_enabled(), "backbone caches are not enabled"
+        curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos)
+        embeds = self._embed_tokens(tokens)
+        masked_embeds = embeds * tokens_mask.unsqueeze(-1)
+        h = masked_embeds.sum(dim=2)
+        h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype)
+
+        last_h = h[:, -1, :]
+        c0_logits = self.codebook0_head(last_h)
+        c0_sample = sample_topk(c0_logits, topk, temperature)
+        c0_embed = self._embed_audio(0, c0_sample)
+
+        curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1)
+        curr_sample = c0_sample.clone()
+        curr_pos = torch.arange(0, curr_h.size(1), device=curr_h.device).unsqueeze(0).repeat(curr_h.size(0), 1)
+
+        # Decoder caches must be reset every frame.
+        self.decoder.reset_caches()
+        for i in range(1, self.config.audio_num_codebooks):
+            curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos)
+            decoder_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask).to(
+                dtype=dtype
+            )
+            ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1])
+            ci_sample = sample_topk(ci_logits, topk, temperature)
+            ci_embed = self._embed_audio(i, ci_sample)
+
+            curr_h = ci_embed
+            curr_sample = torch.cat([curr_sample, ci_sample], dim=1)
+            curr_pos = curr_pos[:, -1:] + 1
+
+        return curr_sample
+
+    def reset_caches(self):
+        self.backbone.reset_caches()
+        self.decoder.reset_caches()
+
+    def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor:
+        return self.audio_embeddings(tokens + codebook * self.config.audio_vocab_size)
+
+    def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        text_embeds = self.text_embeddings(tokens[:, :, -1]).unsqueeze(-2)
+
+        audio_tokens = tokens[:, :, :-1] + (
+            self.config.audio_vocab_size * torch.arange(self.config.audio_num_codebooks, device=tokens.device)
+        )
+        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
+            tokens.size(0), tokens.size(1), self.config.audio_num_codebooks, -1
+        )
+
+        return torch.cat([audio_embeds, text_embeds], dim=-2)

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.4.0
+torchaudio==2.4.0
+tokenizers==0.21.0
+transformers==4.49.0
+huggingface_hub==0.28.1
+moshi==0.2.2
+torchtune==0.4.0
+torchao==0.9.0
+silentcipher @ git+https://github.com/SesameAILabs/silentcipher@master
+gradio==4.26.0
+huggingface-hub-spaces==0.22.0

test_model.py

@@ -0,0 +1,73 @@
+import os
+import torch
+import torchaudio
+import spaces
+from generator import Segment, load_csm_1b
+from huggingface_hub import login
+
+def login_huggingface():
+    """Đăng nhập vào Hugging Face Hub sử dụng token từ biến môi trường hoặc nhập từ người dùng"""
+    hf_token = os.environ.get("HF_TOKEN")
+    
+    if not hf_token:
+        print("Không tìm thấy HF_TOKEN trong biến môi trường.")
+        hf_token = input("Vui lòng nhập Hugging Face token của bạn: ")
+    
+    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!")
+        return True
+    else:
+        print("Không có token. Một số mô hình có thể không truy cập được.")
+        return False
+
+@spaces.GPU
+def generate_test_audio(text, speaker_id, device):
+    """Tạo âm thanh kiểm tra sử dụng ZeroGPU"""
+    generator = load_csm_1b(device=device)
+    print("Đã tải xong mô hình!")
+    
+    print(f"Đang tạo âm thanh cho văn bản: '{text}'")
+    audio = generator.generate(
+        text=text,
+        speaker=speaker_id,
+        context=[],
+        max_audio_length_ms=10000,
+        temperature=0.9,
+        topk=50
+    )
+    
+    return audio, generator.sample_rate
+
+def test_model():
+    print("Kiểm tra mô hình CSM-1B...")
+    
+    # Đăng nhập vào Hugging Face Hub
+    login_huggingface()
+    
+    # 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}")
+    
+    # Tải mô hình CSM-1B và tạo âm thanh
+    print("Đang tải mô hình CSM-1B...")
+    try:
+        # Sử dụng ZeroGPU để tạo âm thanh
+        text = "Xin chào, đây là bài kiểm tra mô hình CSM-1B."
+        speaker_id = 0
+        
+        audio, sample_rate = generate_test_audio(text, speaker_id, device)
+        
+        # Lưu âm thanh thành file
+        output_path = "test_output.wav"
+        torchaudio.save(output_path, audio.unsqueeze(0), sample_rate)
+        print(f"Đã lưu âm thanh vào file: {output_path}")
+        
+        print("Kiểm tra hoàn tất!")
+    except Exception as e:
+        print(f"Lỗi khi kiểm tra mô hình: {e}")
+        print("Vui lòng kiểm tra lại token và quyền truy cập của bạn.")
+
+if __name__ == "__main__":
+    test_model() 

watermarking.py

@@ -0,0 +1,79 @@
+import argparse
+
+import silentcipher
+import torch
+import torchaudio
+
+# This watermark key is public, it is not secure.
+# If using CSM 1B in another application, use a new private key and keep it secret.
+CSM_1B_GH_WATERMARK = [212, 211, 146, 56, 201]
+
+
+def cli_check_audio() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--audio_path", type=str, required=True)
+    args = parser.parse_args()
+
+    check_audio_from_file(args.audio_path)
+
+
+def load_watermarker(device: str = "cuda") -> silentcipher.server.Model:
+    model = silentcipher.get_model(
+        model_type="44.1k",
+        device=device,
+    )
+    return model
+
+
+@torch.inference_mode()
+def watermark(
+    watermarker: silentcipher.server.Model,
+    audio_array: torch.Tensor,
+    sample_rate: int,
+    watermark_key: list[int],
+) -> tuple[torch.Tensor, int]:
+    audio_array_44khz = torchaudio.functional.resample(audio_array, orig_freq=sample_rate, new_freq=44100)
+    encoded, _ = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=36)
+
+    output_sample_rate = min(44100, sample_rate)
+    encoded = torchaudio.functional.resample(encoded, orig_freq=44100, new_freq=output_sample_rate)
+    return encoded, output_sample_rate
+
+
+@torch.inference_mode()
+def verify(
+    watermarker: silentcipher.server.Model,
+    watermarked_audio: torch.Tensor,
+    sample_rate: int,
+    watermark_key: list[int],
+) -> bool:
+    watermarked_audio_44khz = torchaudio.functional.resample(watermarked_audio, orig_freq=sample_rate, new_freq=44100)
+    result = watermarker.decode_wav(watermarked_audio_44khz, 44100, phase_shift_decoding=True)
+
+    is_watermarked = result["status"]
+    if is_watermarked:
+        is_csm_watermarked = result["messages"][0] == watermark_key
+    else:
+        is_csm_watermarked = False
+
+    return is_watermarked and is_csm_watermarked
+
+
+def check_audio_from_file(audio_path: str) -> None:
+    watermarker = load_watermarker(device="cuda")
+
+    audio_array, sample_rate = load_audio(audio_path)
+    is_watermarked = verify(watermarker, audio_array, sample_rate, CSM_1B_GH_WATERMARK)
+
+    outcome = "Watermarked" if is_watermarked else "Not watermarked"
+    print(f"{outcome}: {audio_path}")
+
+
+def load_audio(audio_path: str) -> tuple[torch.Tensor, int]:
+    audio_array, sample_rate = torchaudio.load(audio_path)
+    audio_array = audio_array.mean(dim=0)
+    return audio_array, int(sample_rate)
+
+
+if __name__ == "__main__":
+    cli_check_audio()