Update app.py

app.py

@@ -1,98 +1,94 @@
 import streamlit as st
-import os
+import torchaudio
 import torch
-import soundfile as sf
-from groq import Groq
-from diffusers import AutoPipelineForText2Image
-from streamlit_webrtc import webrtc_streamer, AudioRecorder
-
-# Load API keys
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-HF_API_KEY = os.getenv("HF_API_KEY")
-
-# Initialize Groq client
-client = Groq(api_key=GROQ_API_KEY)
-
-# 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 audio
-def transcribe(audio_path):
-    with open(audio_path, "rb") as file:
-        transcription = client.audio.transcriptions.create(
-            file=(audio_path, file.read()),
-            model="whisper-large-v3",
-            language="ta",
-            response_format="verbose_json"
-        )
-    return transcription["text"]
-
-# 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].message.content
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
+from diffusers import StableDiffusionPipeline
+from io import BytesIO
+import tempfile
+import os
+
+st.set_page_config(page_title="Tamil Voice to Story & Image Generator", layout="wide")
+st.title("🎤 Tamil Voice to Story & Image Generator")
 
-# 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}"}]
+# Load models only once
+@st.cache_resource
+def load_models():
+    # 1. Whisper small for speech recognition
+    whisper_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-small", device=0 if torch.cuda.is_available() else -1)
+
+    # 2. NLLB for Tamil to English translation
+    tokenizer_trans = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+    model_trans = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+
+    # 3. Tiny Story Generator
+    story_gen = pipeline("text-generation", model="sshleifer/tiny-gpt2", device=0 if torch.cuda.is_available() else -1)
+
+    # 4. Image Generator
+    image_pipe = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
     )
-    return response.choices[0].message.content
+    if torch.cuda.is_available():
+        image_pipe.to("cuda")
+
+    return whisper_pipe, tokenizer_trans, model_trans, story_gen, image_pipe
 
-# Function to generate an image
+whisper_pipe, tokenizer_trans, model_trans, story_gen, image_pipe = load_models()
+
+# Function: Translate Tamil to English
+def translate_ta_to_en(text):
+    inputs = tokenizer_trans(text, return_tensors="pt", padding=True)
+    translated = model_trans.generate(**inputs, forced_bos_token_id=tokenizer_trans.lang_code_to_id["eng_Latn"])
+    return tokenizer_trans.batch_decode(translated, skip_special_tokens=True)[0]
+
+# Function: Generate story
+def generate_story(prompt):
+    story = story_gen(prompt, max_length=100, num_return_sequences=1)
+    return story[0]['generated_text']
+
+# Function: Generate image
 def generate_image(prompt):
-    img = image_gen(prompt=prompt).images[0]
-    return img
-
-# Streamlit UI
-st.title("Tamil Speech to Image & Story Generator")
-
-# Choose input method
-input_method = st.radio("Choose Input Method:", ("Record Audio", "Upload Audio"))
-
-audio_path = None
-
-if input_method == "Record Audio":
-    st.subheader("Record your Tamil speech")
-    recorder = webrtc_streamer(key="record_audio", audio=True)
-    
-    if recorder.audio_receiver:
-        audio_data = recorder.audio_receiver.get_frames()  # Get recorded audio
-        audio_path = "recorded_audio.wav"
-        sf.write(audio_path, audio_data, 16000)  # Save recorded audio
-elif input_method == "Upload Audio":
-    uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3"])
-    if uploaded_file:
-        audio_path = "uploaded_audio.wav"
-        with open(audio_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-
-if st.button("Generate"):
-    if not audio_path:
-        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")
+    image = image_pipe(prompt).images[0]
+    return image
+
+# Upload or Record
+input_method = st.radio("Select Input Method", ["Upload Audio", "Record Live"])
+
+if input_method == "Upload Audio":
+    audio_file = st.file_uploader("Upload Tamil Audio", type=["wav", "mp3", "m4a"])
+else:
+    audio_bytes = st.audio("Record or Upload Audio Below", format='audio/wav')
+    audio_file = None
+    if audio_bytes:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmpfile:
+            tmpfile.write(audio_bytes.read())
+            audio_file = tmpfile.name
+
+# Process Button
+if st.button("Generate from Audio") and audio_file:
+    with st.spinner("🔄 Transcribing Tamil audio..."):
+        result = whisper_pipe(audio_file)
+        tamil_text = result['text']
+
+    st.success("✅ Tamil Transcription")
     st.write(tamil_text)
-    
-    st.subheader("English Translation")
+
+    with st.spinner("🌐 Translating to English..."):
+        english_text = translate_ta_to_en(tamil_text)
+
+    st.success("✅ English Translation")
     st.write(english_text)
 
-    st.subheader("Generated Story")
+    with st.spinner("✍️ Generating Story..."):
+        story = generate_story(english_text)
+
+    st.success("✅ Story Generated")
     st.write(story)
 
-    st.subheader("Generated Image")
+    with st.spinner("🎨 Generating Image..."):
+        image = generate_image(english_text)
+
     st.image(image, caption="Generated Image")
 
+elif st.button("Generate from Audio") and not audio_file:
+    st.warning("Please upload or record an audio file.")