feat(model): load model from safetensors

inference.py

@@ -24,6 +24,7 @@ from utils.snac_utils import get_snac, generate_audio_data
 import whisper
 from tqdm import tqdm
 from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
 
 
 torch.set_printoptions(sci_mode=False)
@@ -351,14 +352,14 @@ def load_model(ckpt_dir, device):
     whispermodel = whisper.load_model("small").to(device)
     text_tokenizer = Tokenizer(ckpt_dir)
     fabric = L.Fabric(devices=1, strategy="auto")
-    config = Config.from_file(ckpt_dir + "/model_config.yaml")
+    config = Config.from_file(ckpt_dir + "/config.json")
     config.post_adapter = False
 
     with fabric.init_module(empty_init=False):
         model = GPT(config)
 
     model = fabric.setup(model)
-    state_dict = lazy_load(ckpt_dir + "/lit_model.pth")
+    state_dict = load_file(ckpt_dir + "/lit_model.safetensors")
     model.load_state_dict(state_dict, strict=True)
     model.to(device).eval()
 
@@ -366,7 +367,7 @@ def load_model(ckpt_dir, device):
 
     
 def download_model(ckpt_dir):
-    repo_id = "gpt-omni/mini-omni"
+    repo_id = "leafspark/mini-omni-safetensors"
     snapshot_download(repo_id, local_dir=ckpt_dir, revision="main")