fix

app.py

@@ -8,7 +8,7 @@ import torchaudio
 import gradio as gr
 import numpy as np
 
-from generator import load_csm_1b, Segment
+from generator import Segment, Model, Generator
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -25,7 +25,9 @@ def initialize_model():
     logger.info(f"Using device: {device}")
     
     try:
-        generator = load_csm_1b(device=device)
+        model = Model.from_pretrained("sesame/csm-1b")
+        model = model.to(device=device)
+        generator = Generator(model)
         logger.info(f"Model loaded successfully on device: {device}")
         return True
     except Exception as e:

app_huggingface.py

@@ -25,7 +25,8 @@ class MockGenerator:
 try:
     import torch
     import torchaudio
-    from generator import load_csm_1b, Segment
+    # Chỉ import các thành phần cần thiết
+    from generator import Segment
     TORCH_AVAILABLE = True
 except ImportError:
     TORCH_AVAILABLE = False
@@ -57,9 +58,21 @@ def initialize_model():
         logger.info(f"Using device: {device}")
         
         try:
-            # Try to use the actual model
-            generator = load_csm_1b(device=device)
-            logger.info(f"Model loaded successfully on device: {device}")
+            # Cố gắng tải model theo cách khác, không sử dụng load_csm_1b
+            from generator import Model, Generator
+            from huggingface_hub import hf_hub_download
+            
+            try:
+                # Trực tiếp khởi tạo mô hình từ pretrained
+                model = Model.from_pretrained("sesame/csm-1b")
+                model = model.to(device=device)
+                generator = Generator(model)
+                logger.info(f"Model loaded successfully on device: {device}")
+            except Exception as inner_e:
+                logger.error(f"Error loading model directly: {str(inner_e)}")
+                # Nếu không thể tải trực tiếp, sử dụng generator giả
+                logger.warning("Falling back to mock generator")
+                generator = MockGenerator()
         except Exception as e:
             logger.error(f"Error loading actual model: {str(e)}")
             # Fall back to mock generator

generator.py

@@ -164,8 +164,26 @@ class Generator:
 
 
 def load_csm_1b(device: str = "cuda") -> Generator:
-    model = Model.from_pretrained("sesame/csm-1b")
-    model.to(device=device, dtype=torch.bfloat16)
-
-    generator = Generator(model)
-    return generator
+    try:
+        # Try the simple approach first
+        model = Model.from_pretrained("sesame/csm-1b")
+        model = model.to(device=device)
+        
+        generator = Generator(model)
+        return generator
+    except Exception as e:
+        # Log the error for debugging
+        import logging
+        logging.error(f"Error in standard model loading: {str(e)}")
+        
+        # Try alternative loading if config is available
+        try:
+            from silentcipher import Config
+            model_path = "sesame/csm-1b"
+            config = Config.from_pretrained(model_path)
+            model = Model.from_pretrained(model_path, config=config)
+            model = model.to(device=device)
+            return Generator(model)
+        except ImportError:
+            # Config not available, try direct initialization
+            raise RuntimeError("Could not load model with either method. Original error: " + str(e))