added ausio generation

audio_chat_client.py

@@ -0,0 +1,26 @@
+import requests
+import streamlit as st
+# from streamlitui import StreamlitUI
+
+# stui = StreamlitUI(api_url="http://localhost:8000")  # FastAPI backend URL
+st.title("FastAPI ChatBot") 
+
+if "messages" not in st.session_state:
+    st.session_state.messages = [] 
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        content = message["content"]
+        if isinstance(content, bytes):
+            st.audio(content)
+        else:
+            st.markdown(content)
+
+
+if prompt := st.chat_input("Write your prompt in this input field"):
+    response = requests.get(
+        f"https://ahmed-eisa-genai-service.hf.space/generate/audio", params={"prompt": prompt}
+    )
+    response.raise_for_status()
+    with st.chat_message("assistant"):
+        st.text("Here is your generated audio")
+        st.audio(response.content)

client.py

@@ -0,0 +1,27 @@
+import requests
+import streamlit as st
+# from streamlitui import StreamlitUI
+
+# stui = StreamlitUI(api_url="http://localhost:8000")  # FastAPI backend URL
+st.title("FastAPI ChatBot") 
+
+if "messages" not in st.session_state:
+    st.session_state.messages = [] 
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"]) 
+
+if prompt := st.chat_input("Write your prompt in this input field"): 
+    st.session_state.messages.append({"role": "user", "content": prompt}) 
+
+    with st.chat_message("user"):
+        st.text(prompt) 
+
+    response = requests.get(
+        f"https://ahmed-eisa-genai-service.hf.space//generate/text", params={"prompt": prompt}
+    ) 
+    response.raise_for_status() 
+
+    with st.chat_message("assistant"):
+        st.markdown(response.text)

main.py

@@ -1,7 +1,10 @@
 # main.py
-from fastapi import FastAPI
-from models import load_text_model,generate_text
+from fastapi import FastAPI,status
+from fastapi.responses import StreamingResponse
 
+from models import load_text_model,generate_text,load_audio_model,generate_audio
+from schemas import VoicePresets
+from utils import audio_array_to_buffer
 app = FastAPI() 
 
 @app.get("/")
@@ -12,4 +15,20 @@ def root_controller():
 def serve_language_model_controller(prompt: str) -> str: 
     pipe = load_text_model() 
     output = generate_text(pipe, prompt) 
-    return output 
+    return output 
+
+
+@app.get(
+    "/generate/audio",
+    responses={status.HTTP_200_OK: {"content": {"audio/wav": {}}}},
+    response_class=StreamingResponse,
+) 
+def serve_text_to_audio_model_controller(
+    prompt: str,
+    preset: VoicePresets = "v2/en_speaker_1",
+):
+    processor, model = load_audio_model()
+    output, sample_rate = generate_audio(processor, model, prompt, preset)
+    return StreamingResponse(
+        audio_array_to_buffer(output, sample_rate), media_type="audio/wav"
+    )

models.py

@@ -1,7 +1,9 @@
 # models.py
 
 import torch
-from transformers import Pipeline, pipeline
+from transformers import Pipeline, pipeline,AutoProcessor, AutoModel, BarkProcessor, BarkModel
+from schemas import VoicePresets
+import numpy as np
 
 prompt = "How to set up a FastAPI project?"
 system_prompt = """
@@ -12,6 +14,24 @@ Always respond in markdown.
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
 
+def load_audio_model() -> tuple[BarkProcessor, BarkModel]:
+    processor = AutoProcessor.from_pretrained("suno/bark-small", device=device)
+    model = AutoModel.from_pretrained("suno/bark-small", device=device)
+    return processor, model
+
+def generate_audio(
+    processor: BarkProcessor,
+    model: BarkModel,
+    prompt: str,
+    preset: VoicePresets,
+) -> tuple[np.array, int]:
+    inputs = processor(text=[prompt], return_tensors="pt", voice_preset=preset)
+    output = model.generate(**inputs, do_sample=True).cpu().numpy().squeeze()
+    sample_rate = model.generation_config.sample_rate
+    return output, sample_rate
+
+
+
 def load_text_model():
     pipe = pipeline(
         "text-generation",

requirements.txt

@@ -3,4 +3,5 @@ uvicorn
 transformers
 torch
 pydantic
-bitsandbytes
+bitsandbytes
+soundfile

schemas.py

@@ -0,0 +1,3 @@
+from typing import Literal
+
+VoicePresets = Literal["v2/en_speaker_1", "v2/en_speaker_9"]

utils.py

@@ -0,0 +1,9 @@
+from io import BytesIO
+import soundfile
+import numpy as np
+
+def audio_array_to_buffer(audio_array: np.array, sample_rate: int) -> BytesIO:
+    buffer = BytesIO()
+    soundfile.write(buffer, audio_array, sample_rate, format="wav") 
+    buffer.seek(0)
+    return buffer