update readme

README.md

@@ -40,23 +40,27 @@ Usage example:
 from transformers import VitsModel, AutoTokenizer, set_seed
 import torch
 import scipy
-from ruaccent import RUAccent
+
+device = 'cuda' #  'cpu' or 'cuda'
+
+speaker = 0 # 0-woman, 1-man  
+
+set_seed(555)  # make deterministic
 
 # load model
 model_name = "utrobinmv/tts_ru_free_hf_vits_low_multispeaker"
 
-model = VitsModel.from_pretrained(model_name)
+model = VitsModel.from_pretrained(model_name).to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model.eval()
 
+from ruaccent import RUAccent
+
 # load accentizer
 accentizer = RUAccent()
-accentizer.load(omograph_model_size='turbo', use_dictionary=True)
-
-set_seed(555)  # make deterministic
-
-speaker = 0 # 0-woman, 1-man  
+accentizer.load(omograph_model_size='turbo', use_dictionary=True, device=device)
 
+# text
 text = """Ночью двадцать третьего июня начал извергаться самый высокий 
 действующий вулкан в Евразии - Кл+ючевской. Об этом сообщила руководитель 
 Камчатской группы реагирования на вулканические извержения, ведущий 
@@ -78,13 +82,12 @@ print(text)
 # ни насел+ению, ни ави+ации » поясн+ила тасс госпож+а г+ирина.
 
 inputs = tokenizer(text, return_tensors="pt")
-inputs['speaker_id'] = speaker
 
 with torch.no_grad():
-    output = model(**inputs).waveform
+    output = model(**inputs.to(device), speaker_id=speaker).waveform.detach().cpu().numpy()
     
 scipy.io.wavfile.write("tts_audio.wav", rate=model.config.sampling_rate,
-                       data=output[0].cpu().numpy())
+                       data=output[0])
 ```