Update app.py

app.py

@@ -1,11 +1,19 @@
 import gradio as gr
 import torch
-from transformers import AutoProcessor, AutoModel
+from transformers import AutoProcessor, BarkModel
 import scipy
 
 
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16).to(device)
+# model.enable_cpu_offload()
+
+device = "cpu"
+
+
 processor = AutoProcessor.from_pretrained("suno/bark-small")
-model = AutoModel.from_pretrained("suno/bark-small")
+model = BarkModel.from_pretrained("suno/bark-small").to(device)
+
 
 
 def run_bark(text, lang="en", n=1):