Create app.py

app.py

@@ -0,0 +1,33 @@
+import gradio as gr
+import torch
+from transformers import BarkModel
+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"
+model = model.to(device)
+
+# Use bettertransform for flash attention
+model = BetterTransformer.transform(model, keep_original_model=False)
+
+# Enable CPU offload
+model.enable_cpu_offload()
+
+from transformers import AutoProcessor
+
+def infer(text_prompt):
+    text_prompt = text_prompt
+    inputs = processor(text_prompt).to(device)
+    with torch.inference_mode():
+        speech_output = measure_latency_and_memory_use(model, inputs, nb_loops = 5)
+        audio_out = speech_output[0].cpu().numpy()
+    return audio_out
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        prompt = gr.Textbox(label="prompt")
+        submit.btn = gr.Button("Submit")
+        audio_out = gr.Audio(type="numpy")
+
+demo.launch()
+