jacob-c commited on
Commit
d1ab29e
·
1 Parent(s): fbe6556

diffuserfix

Browse files
Files changed (1) hide show
  1. audioldm/bigvgan/__init__.py +6 -2
audioldm/bigvgan/__init__.py CHANGED
@@ -3,11 +3,15 @@ import bigvgan
3
  from huggingface_hub import hf_hub_download
4
 
5
  class BigVGANVocoder:
6
- def __init__(self, device='cuda', use_cuda_kernel=False):
 
 
 
 
7
  # Load the pretrained model
8
  self.model = bigvgan.BigVGAN.from_pretrained(
9
  'nvidia/bigvgan_v2_44khz_128band_512x',
10
- use_cuda_kernel=use_cuda_kernel
11
  )
12
  self.model.remove_weight_norm()
13
  self.model.eval().to(device)
 
3
  from huggingface_hub import hf_hub_download
4
 
5
  class BigVGANVocoder:
6
+ def __init__(self, device=None):
7
+ # Set default device if none provided
8
+ if device is None:
9
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
10
+
11
  # Load the pretrained model
12
  self.model = bigvgan.BigVGAN.from_pretrained(
13
  'nvidia/bigvgan_v2_44khz_128band_512x',
14
+ use_cuda_kernel=(device == 'cuda')
15
  )
16
  self.model.remove_weight_norm()
17
  self.model.eval().to(device)