Update app.py

app.py

@@ -1,25 +1,22 @@
-import os
-import torch
 import streamlit as st
+import torch
+import os
+import soundfile as sf
 from groq import Groq
 from diffusers import AutoPipelineForText2Image
-import tempfile
-import soundfile as sf
 
 # Load API keys
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 HF_API_KEY = os.getenv("HF_API_KEY")
 
-# Initialize Groq client with API key
+# Initialize Groq client
 client = Groq(api_key=GROQ_API_KEY)
 
-# Load lightweight Hugging Face image generation model
-image_gen = AutoPipelineForText2Image.from_pretrained(
-    "stabilityai/sdxl-turbo", use_auth_token=HF_API_KEY
-)
-image_gen.to("cuda" if torch.cuda.is_available() else "cpu")
+# Load image generation model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+image_gen = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo").to(device)
 
-# Function to transcribe Tamil audio using Groq's Whisper
+# Function to transcribe audio
 def transcribe(audio_path):
     with open(audio_path, "rb") as file:
         transcription = client.audio.transcriptions.create(
@@ -30,21 +27,21 @@ def transcribe(audio_path):
         )
     return transcription["text"]
 
-# Function to translate Tamil to English using Groq's Gemma
+# Function to translate Tamil to English
 def translate_text(tamil_text):
     response = client.chat.completions.create(
         model="gemma-7b-it",
         messages=[{"role": "user", "content": f"Translate this Tamil text to English: {tamil_text}"}]
     )
-    return response.choices[0].delta.content
+    return response.choices[0].message.content
 
-# Function to generate text using Groq's DeepSeek R1
+# Function to generate text
 def generate_text(prompt):
     response = client.chat.completions.create(
         model="deepseek-coder-r1-7b",
         messages=[{"role": "user", "content": f"Write a short story about: {prompt}"}]
     )
-    return response.choices[0].delta.content
+    return response.choices[0].message.content
 
 # Function to generate an image
 def generate_image(prompt):
@@ -54,36 +51,37 @@ def generate_image(prompt):
 # Streamlit UI
 st.title("Tamil Speech to Image & Story Generator")
 
-# Audio input - Recording or Uploading
-st.subheader("Upload or Record Audio")
-recorded_audio = st.audio("", format='audio/wav', start_time=0)
-uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "m4a"]) 
+# Choose input method
+input_method = st.radio("Choose Input Method:", ("Record Audio", "Upload Audio"))
 
-audio_path = None
+if input_method == "Record Audio":
+    recorded_audio = st.audio(st.file_uploader("Record your Tamil speech", type=["wav", "mp3"]))
+elif input_method == "Upload Audio":
+    uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3"])
 
-if uploaded_file is not None:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
-        temp_audio.write(uploaded_file.read())
-        audio_path = temp_audio.name
-elif recorded_audio:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
+if st.button("Generate"):
+    if input_method == "Record Audio" and recorded_audio:
         audio_data, samplerate = sf.read(recorded_audio)
-        sf.write(temp_audio.name, audio_data, samplerate)
-        audio_path = temp_audio.name
-
-if st.button("Generate") and audio_path:
-    with st.spinner("Transcribing Tamil speech..."):
-        tamil_text = transcribe(audio_path)
-    with st.spinner("Translating to English..."):
-        english_text = translate_text(tamil_text)
-    with st.spinner("Generating story..."):
-        story = generate_text(english_text)
-    with st.spinner("Generating image..."):
-        image = generate_image(english_text)
-
+        audio_path = "recorded_audio.wav"
+        sf.write(audio_path, audio_data, samplerate)
+    elif input_method == "Upload Audio" and uploaded_file:
+        audio_path = "uploaded_audio.wav"
+        with open(audio_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+    else:
+        st.error("Please provide an audio file.")
+        st.stop()
+
+    # Process audio
+    tamil_text = transcribe(audio_path)
+    english_text = translate_text(tamil_text)
+    story = generate_text(english_text)
+    image = generate_image(english_text)
+
+    # Display results
     st.subheader("Tamil Transcription")
     st.write(tamil_text)
-
+    
     st.subheader("English Translation")
     st.write(english_text)
 
@@ -91,4 +89,5 @@ if st.button("Generate") and audio_path:
     st.write(story)
 
     st.subheader("Generated Image")
-    st.image(image)
+    st.image(image, caption="Generated Image")
+