Vansh Chugh commited on
Commit
a5e496e
·
1 Parent(s): 94e6004

move extractor to GPU

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -124,6 +124,7 @@ def _activate_variant(entry):
124
  is where it's actually safe to touch CUDA."""
125
  model, extractor, audio_vae, vae_scale_factor, sample_rate, _, scheduler = entry
126
  extractor.text_backbone = extractor.text_backbone.to(DEVICE)
 
127
  model = model.to(device=DEVICE, dtype=torch.bfloat16).eval()
128
  audio_vae = audio_vae.to(DEVICE).eval()
129
 
 
124
  is where it's actually safe to touch CUDA."""
125
  model, extractor, audio_vae, vae_scale_factor, sample_rate, _, scheduler = entry
126
  extractor.text_backbone = extractor.text_backbone.to(DEVICE)
127
+ extractor.device = DEVICE
128
  model = model.to(device=DEVICE, dtype=torch.bfloat16).eval()
129
  audio_vae = audio_vae.to(DEVICE).eval()
130