jameshuntercarter commited on
Commit
d9382cf
1 Parent(s): 37eb59e

Upload predict.py

Browse files
Files changed (1) hide show
  1. predict.py +4 -2
predict.py CHANGED
@@ -19,8 +19,7 @@ class Predictor(BasePredictor):
19
  ) -> ModelOutput:
20
  """Run a single prediction on the model"""
21
  # SETUP
22
- large_quant_model = False # Use the larger pretrained model
23
- device = torch.device('cuda') # 'cuda', 'cpu', 'cuda:0', 0, -1, torch.device('cuda')
24
 
25
  import numpy as np
26
  import torch
@@ -31,6 +30,9 @@ class Predictor(BasePredictor):
31
  from bark_hubert_quantizer.pre_kmeans_hubert import CustomHubert
32
  from bark_hubert_quantizer.customtokenizer import CustomTokenizer
33
 
 
 
 
34
  model = ('quantifier_V1_hubert_base_ls960_23.pth', 'tokenizer_large.pth') if large_quant_model else (
35
  'quantifier_hubert_base_ls960_14.pth', 'tokenizer.pth')
36
 
 
19
  ) -> ModelOutput:
20
  """Run a single prediction on the model"""
21
  # SETUP
22
+
 
23
 
24
  import numpy as np
25
  import torch
 
30
  from bark_hubert_quantizer.pre_kmeans_hubert import CustomHubert
31
  from bark_hubert_quantizer.customtokenizer import CustomTokenizer
32
 
33
+ large_quant_model = False # Use the larger pretrained model
34
+ device = torch.device('cuda') # 'cuda', 'cpu', 'cuda:0', 0, -1, torch.device('cuda')
35
+
36
  model = ('quantifier_V1_hubert_base_ls960_23.pth', 'tokenizer_large.pth') if large_quant_model else (
37
  'quantifier_hubert_base_ls960_14.pth', 'tokenizer.pth')
38