Update README.md

README.md

@@ -59,13 +59,12 @@ pip install diffusers
 
 ```python
 import torch
-from transformers import VibeVoiceFeatureExtractor, VibeVoiceAcousticTokenizerModel
+from transformers import AutoFeatureExtractor, VibeVoiceAcousticTokenizerModel
 from transformers.audio_utils import load_audio_librosa
 from scipy.io import wavfile
 
 
-model_path = "bezzam/VibeVoice-AcousticTokenizer"
-fe_path = "bezzam/VibeVoice-1.5B"
+model_id = "bezzam/VibeVoice-AcousticTokenizer"
 sampling_rate = 24000
 
 # load audio
@@ -75,11 +74,11 @@ audio = load_audio_librosa(
 )
 
 # load model
-torch_device = "cuda" if torch.cuda.is_available() else "cpu"
-feature_extractor = VibeVoiceFeatureExtractor.from_pretrained(fe_path)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
 model = VibeVoiceAcousticTokenizerModel.from_pretrained(
-    model_path, device_map=torch_device,
-).to(torch_device).eval()
+    model_id, device_map=device,
+).eval()
 
 # preprocess audio
 inputs = feature_extractor(
@@ -89,7 +88,7 @@ inputs = feature_extractor(
     pad_to_multiple_of=3200,
     return_attention_mask=False,
     return_tensors="pt"
-).to(torch_device)
+).to(device)
 print("Input audio shape:", inputs.input_features.shape)
 # Input audio shape: torch.Size([1, 1, 224000])