Update README.md

README.md

@@ -8,14 +8,25 @@ SNAC in safetensors format.
 import json
 from pathlib import Path
 
+import safetensors.torch
 import torch
-from safetensors.torch import load_model
+import torchaudio
 from snac import SNAC
 
-device = "cuda"
-dtype = torch.float32
 config = json.loads(Path("config.json").read_text(encoding="utf-8"))
 model = SNAC(**config)
-load_model(model, "model.safetensors", device=device)
-model.to(device, dtype).eval()
+state_dict = safetensors.torch.load_file("model.safetensors")
+model.load_state_dict(state_dict)
+model.cuda().eval()
+
+input, sr = torchaudio.load("input.wav")
+input = torchaudio.functional.resample(input, sr, model.sampling_rate)
+input = input.cuda().unsqueeze(0)
+
+with torch.inference_mode():
+    codes = model.encode(input)
+    output = model.decode(codes)
+
+output = output.cpu().squeeze(0)
+torchaudio.save("output.wav", output, model.sampling_rate)
 ```