Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from transformers import BarkModel
-#from optimum.bettertransformer import BetterTransformer
+from optimum.bettertransformer import BetterTransformer
 
 model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16)
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -11,15 +11,15 @@ from transformers import AutoProcessor
 processor = AutoProcessor.from_pretrained("suno/bark-small")
 
 # Use bettertransform for flash attention
-#model = BetterTransformer.transform(model, keep_original_model=False)
+model = BetterTransformer.transform(model, keep_original_model=False)
 
 # Enable CPU offload
-#model.enable_cpu_offload()
+model.enable_cpu_offload()
 
 import numpy as np
-#from scipy.io import wavfile
+from scipy.io.wavfile import write as write_wav
 #from pydub import AudioSegment
-import soundfile as sf
+#import soundfile as sf
 
 def infer(text_prompt):
     text_prompt = text_prompt
@@ -28,14 +28,14 @@ def infer(text_prompt):
     with torch.inference_mode():
         speech_output = model.generate(**inputs, do_sample = True, fine_temperature = 0.4, coarse_temperature = 0.8)
         #audio_out = speech_output[0].cpu().numpy()
-        audio_out = speech_output[0]
+        audio_array = speech_output[0].cpu().numpy().squeeze()
         
         # Assuming audio_out contains audio data and the sampling rate
         sampling_rate = model.generation_config.sample_rate
+        write_wav("output.wav", sample_rate, audio_array)
         #wavfile.write("output.wav", sampling_rate, audio_out)
         
-        # Save the audio data as an audio file using soundfile library
-        sf.write("output.wav", audio_out, sampling_rate)
+