fix

app.py

@@ -11,6 +11,9 @@ from dataclasses import dataclass
 from generator import Segment, load_csm_1b
 from huggingface_hub import login
 
+# Tắt tính năng compile của torch để tránh lỗi triton
+torch._dynamo.config.suppress_errors = True
+
 # 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}")
@@ -28,17 +31,29 @@ def login_huggingface():
 # Đăng nhập khi khởi động
 login_huggingface()
 
-# Tải mô hình CSM-1B
+# Biến toàn cục để theo dõi trạng thái mô hình
 generator = None
+model_loaded = False
 
-def load_model():
-    global generator
-    if generator is None:
-        print("Đang tải mô hình CSM-1B...")
-        generator = load_csm_1b(device=device)
+# Hàm tải mô hình được gọi trong ZeroGPU
+@spaces.GPU
+def initialize_model():
+    global generator, model_loaded
+    if not model_loaded:
+        print("Đang tải mô hình CSM-1B trong GPU...")
+        generator = load_csm_1b(device="cuda")
+        model_loaded = True
         print("Đã tải xong mô hình!")
     return generator
 
+# Hàm lấy mô hình đã tải
+@spaces.GPU
+def get_model():
+    global generator, model_loaded
+    if not model_loaded:
+        return initialize_model()
+    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)
@@ -68,8 +83,8 @@ def generate_speech(
     top_k: int = 50,
     progress=gr.Progress()
 ) -> str:
-    # Tải mô hình nếu chưa tải
-    generator = load_model()
+    # Lấy mô hình đã tải
+    generator = get_model()
     
     # Chuẩn bị ngữ cảnh (context)
     context = []
@@ -119,8 +134,8 @@ def generate_speech_simple(
     top_k: int = 50,
     progress=gr.Progress()
 ) -> str:
-    # Tải mô hình nếu chưa tải
-    generator = load_model()
+    # Lấy mô hình đã tải
+    generator = get_model()
     
     progress(0.3, "Đang tạo âm thanh...")
     # Tạo âm thanh từ văn bản
@@ -282,8 +297,11 @@ def create_demo():
             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.
+            
+            Lưu ý: Trong môi trường ZeroGPU, CUDA không được khởi tạo trong quá trình chính, mà chỉ trong các hàm có decorator @spaces.GPU.
             """)
             
+            @spaces.GPU
             def check_gpu():
                 if torch.cuda.is_available():
                     gpu_name = torch.cuda.get_device_name(0)
@@ -296,6 +314,21 @@ def create_demo():
             gpu_info = gr.Textbox(label="Thông tin GPU", interactive=False)
             
             check_gpu_btn.click(fn=check_gpu, inputs=None, outputs=gpu_info)
+            
+            # Thêm nút tải mô hình
+            load_model_btn = gr.Button("Tải mô hình")
+            model_status = gr.Textbox(label="Trạng thái mô hình", interactive=False)
+            
+            @spaces.GPU
+            def load_model_and_report():
+                global model_loaded
+                if model_loaded:
+                    return "Mô hình đã được tải trước đó!"
+                else:
+                    initialize_model()
+                    return "Mô hình đã được tải thành công!"
+            
+            load_model_btn.click(fn=load_model_and_report, inputs=None, outputs=model_status)
         
         # Kết nối các thành phần
         generate_btn.click(
@@ -327,9 +360,6 @@ def create_demo():
             ],
             outputs=output_audio_context
         )
-        
-        # Tải mô hình khi khởi động
-        demo.load(fn=load_model)
     
     return demo
 

generator.py

@@ -10,6 +10,8 @@ from tokenizers.processors import TemplateProcessing
 from transformers import AutoTokenizer
 from watermarking import CSM_1B_GH_WATERMARK, load_watermarker, watermark
 
+# Tắt tính năng compile của torch để tránh lỗi triton
+torch._dynamo.config.suppress_errors = True
 
 @dataclass
 class Segment:
@@ -174,8 +176,14 @@ def load_csm_1b(device: str = "cuda") -> Generator:
         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)
+        # Trong ZeroGPU, không nên khởi tạo CUDA trong quá trình chính
+        # Chỉ chuyển mô hình sang GPU khi được gọi trong hàm có decorator @spaces.GPU
+        if 'cuda' in device and not torch.cuda.is_initialized():
+            # Sử dụng CPU cho quá trình chính
+            model = Model.from_pretrained("sesame/csm-1b")
+        else:
+            model = Model.from_pretrained("sesame/csm-1b")
+            model.to(device=device, dtype=torch.bfloat16)
         
         generator = Generator(model)
         return generator