Spaces:
Running on Zero
Running on Zero
Vansh Chugh commited on
Commit ·
a5e496e
1
Parent(s): 94e6004
move extractor to GPU
Browse files
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 |
|