Update app.py

app.py

@@ -1,8 +1,9 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoProcessor, AutoModel
 import torch
-from transformers import pipeline
 from diffusers import FluxPipeline
+from IPython.display import Audio
+
 
 # llama
 model_id = "meta-llama/Llama-3.2-3B-Instruct"
@@ -20,6 +21,10 @@ messages = [
 flux = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
 flux.enable_model_cpu_offload()
 
+#suno
+processor = AutoProcessor.from_pretrained("suno/bark-small")
+vc = AutoModel.from_pretrained("suno/bark-small")
+
 def poet(text):
     messages.append({"role": "user", "content": text})
     outputs = pipe(
@@ -41,10 +46,28 @@ def poet_image(poetry):
     ).images[0]
     return image
 
+def poet_bard(poetry):
+    inputs = processor(
+        text=[poetry],
+        return_tensors="pt",
+    )
+    speech_values = vc.generate(**inputs, do_sample=True)
+    sampling_rate = vc.generation_config.sample_rate
+    return Audio(speech_values.cpu().numpy().squeeze(), rate=sampling_rate)
+
 st.title("Shakespeare Ai")
 st.write("A space made to allow people to create shakespeare like text with images!")
+
+# get prompt
 prompt = st.text_input("Enter your prompt: ")
 st.button("Generate Shakespeare")
+
+# analyze prompt
 shakespeare = poet(prompt)
+bard = poet_bard(shakespeare)
+img = poet_image(shakespeare)
+
+# write content
 st.write(shakespeare)
-st.image(poet_image(shakespeare))
+st.audio(bard)
+st.image(img)