Update app.py

app.py

@@ -1,93 +1,48 @@
-import os
-import json
-import tempfile
-import torch
-import soundfile as sf
 import gradio as gr
 import requests
-import io
-
-from dotenv import load_dotenv
+import os
+import tempfile
 from groq import Groq
-from PIL import Image
 
+import torch
+import soundfile as sf
 from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
+from datasets import load_dataset
 
-load_dotenv()
 
-# =============================
+# ==============================
 # API KEYS
-# =============================
+# ==============================
 
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 HF_TOKEN = os.getenv("HF_TOKEN")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
-client = Groq(api_key=GROQ_API_KEY)
-
-CHAT_FILE = "chat_history.json"
-PREF_FILE = "preferences.json"
-
-# =============================
-# JSON HELPERS
-# =============================
-
-def load_json(file, default):
-    if os.path.exists(file):
-        try:
-            with open(file, "r") as f:
-                return json.load(f)
-        except:
-            return default
-    return default
-
-
-def save_json(file, data):
-    with open(file, "w") as f:
-        json.dump(data, f, indent=4)
-
+groq_client = Groq(api_key=GROQ_API_KEY)
 
-conversation_history = load_json(CHAT_FILE, [])
-user_preferences = load_json(PREF_FILE, {"style": "Default"})
 
-# =============================
-# LOAD TTS MODEL
-# =============================
+# ==============================
+# LOAD TTS MODELS
+# ==============================
 
 processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-tts_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
-
-speaker_embeddings = torch.randn(1, 512)
-
-# =============================
-# HUGGING FACE IMAGE API
-# =============================
-
-HF_API_URL = "https://api-inference.huggingface.co/models/stabilityai/sdxl-turbo"
 
-headers = {
-    "Authorization": f"Bearer {HF_TOKEN}"
-}
-
-# =============================
-# SPEECH TO TEXT
-# =============================
-
-def transcribe_audio(audio_path):
+tts_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
 
-    with open(audio_path, "rb") as audio_file:
+vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
 
-        transcription = client.audio.transcriptions.create(
-            file=audio_file,
-            model="whisper-large-v3"
-        )
+embeddings_dataset = load_dataset(
+    "Matthijs/cmu-arctic-xvectors",
+    split="validation"
+)
 
-    return transcription.text
+speaker_embeddings = torch.tensor(
+    embeddings_dataset[7306]["xvector"]
+).unsqueeze(0)
 
 
-# =============================
-# TEXT TO SPEECH
-# =============================
+# ==============================
+# TEXT → SPEECH
+# ==============================
 
 def text_to_speech(text):
 
@@ -111,41 +66,23 @@ def text_to_speech(text):
     return temp_audio.name
 
 
-# =============================
-# IMAGE TRIGGER
-# =============================
-
-def should_generate_image(user_prompt):
-
-    keywords = [
-        "draw",
-        "diagram",
-        "visualize",
-        "show me",
-        "illustration",
-        "picture",
-        "image",
-        "architecture"
-    ]
-
-    for word in keywords:
-        if word in user_prompt.lower():
-            return True
+# ==============================
+# IMAGE GENERATION
+# ==============================
 
-    return False
+def generate_image(prompt):
 
+    API_URL = "https://router.huggingface.co/hf-inference/models/stabilityai/stable-diffusion-xl-base-1.0"
 
-# =============================
-# IMAGE GENERATION
-# =============================
+    headers = {
+        "Authorization": f"Bearer {HF_TOKEN}"
+    }
 
-def generate_image(prompt):
+    payload = {
+        "inputs": prompt
+    }
 
-    response = requests.post(
-        HF_API_URL,
-        headers=headers,
-        json={"inputs": prompt}
-    )
+    response = requests.post(API_URL, headers=headers, json=payload)
 
     print("HF STATUS:", response.status_code)
 
@@ -153,157 +90,83 @@ def generate_image(prompt):
         print(response.text)
         return None
 
-    image = Image.open(io.BytesIO(response.content))
+    image_bytes = response.content
 
-    return image
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
 
+    temp_file.write(image_bytes)
 
-# =============================
-# CHAT FUNCTION
-# =============================
+    temp_file.close()
 
-def chat_with_memory(user_message, preference_text):
+    return temp_file.name
 
-    global conversation_history, user_preferences
 
-    if preference_text and preference_text.strip():
-        user_preferences["style"] = preference_text
-        save_json(PREF_FILE, user_preferences)
+# ==============================
+# GROQ CHATBOT
+# ==============================
 
-    system_prompt = f"""
-You are a helpful AI assistant.
+def ask_llm(question):
 
-User Preferences:
-{user_preferences.get("style", "Default")}
-
-Maintain conversational memory.
-"""
-
-    messages = [{"role": "system", "content": system_prompt}]
-    messages.extend(conversation_history)
-    messages.append({"role": "user", "content": user_message})
-
-    response = client.chat.completions.create(
+    response = groq_client.chat.completions.create(
         model="llama-3.1-8b-instant",
         max_tokens=200,
-        messages=messages
+        messages=[
+            {"role": "user", "content": question}
+        ]
     )
 
-    assistant_reply = response.choices[0].message.content
-
-    conversation_history.append({"role": "user", "content": user_message})
-    conversation_history.append({"role": "assistant", "content": assistant_reply})
-
-    save_json(CHAT_FILE, conversation_history)
-
-    return assistant_reply
-
-
-# =============================
-# PROCESS TEXT
-# =============================
-
-def process_text(user_message, preference_text, chat_display):
-
-    if not user_message.strip():
-        return "", chat_display, None, None
-
-    assistant_reply = chat_with_memory(user_message, preference_text)
-
-    chat_display.append({"role": "user", "content": user_message})
-    chat_display.append({"role": "assistant", "content": assistant_reply})
-
-    audio_output = text_to_speech(assistant_reply)
-
-    image_output = None
-
-    if should_generate_image(user_message):
-        image_output = generate_image(user_message)
-
-    return "", chat_display, audio_output, image_output
-
+    return response.choices[0].message.content
 
-# =============================
-# PROCESS VOICE
-# =============================
 
-def process_voice(audio_file, preference_text, chat_display):
+# ==============================
+# MAIN ASSISTANT FUNCTION
+# ==============================
 
-    user_text = transcribe_audio(audio_file)
+def ai_assistant(user_input):
 
-    assistant_reply = chat_with_memory(user_text, preference_text)
+    reply = ask_llm(user_input)
 
-    chat_display.append({"role": "user", "content": user_text})
-    chat_display.append({"role": "assistant", "content": assistant_reply})
+    image = None
 
-    audio_output = text_to_speech(assistant_reply)
+    if "image" in user_input.lower() or "generate" in user_input.lower():
+        image = generate_image(user_input)
 
-    image_output = None
+    audio = text_to_speech(reply)
 
-    if should_generate_image(user_text):
-        image_output = generate_image(user_text)
+    return reply, audio, image
 
-    return chat_display, audio_output, image_output
 
-
-# =============================
-# CLEAR MEMORY
-# =============================
-
-def clear_memory():
-
-    global conversation_history
-
-    conversation_history = []
-
-    save_json(CHAT_FILE, [])
-
-    return []
-
-
-# =============================
+# ==============================
 # GRADIO UI
-# =============================
+# ==============================
 
 with gr.Blocks() as demo:
 
-    gr.Markdown("# 🤖 Version 5 — Multimodal AI Assistant (Voice + Images)")
+    gr.Markdown("# 🤖 AI Assistant (Chat + Voice + Image)")
 
-    chatbot = gr.Chatbot(label="Conversation", value=conversation_history)
-
-    preference_input = gr.Textbox(label="User Preferences")
-
-    user_message = gr.Textbox(label="Type message")
-
-    audio_input = gr.Audio(
-        sources=["microphone"],
-        type="filepath",
-        label="Voice Input"
+    user_input = gr.Textbox(
+        label="Ask something or request an image"
     )
 
-    audio_output = gr.Audio(label="Voice Response")
-
-    image_output = gr.Image(label="Generated Image")
-
-    send_btn = gr.Button("Send Text")
-    voice_btn = gr.Button("Send Voice")
-    clear_btn = gr.Button("Clear Memory")
+    text_output = gr.Textbox(
+        label="Assistant Response"
+    )
 
-    send_btn.click(
-        process_text,
-        inputs=[user_message, preference_input, chatbot],
-        outputs=[user_message, chatbot, audio_output, image_output]
+    audio_output = gr.Audio(
+        label="Voice Response"
     )
 
-    voice_btn.click(
-        process_voice,
-        inputs=[audio_input, preference_input, chatbot],
-        outputs=[chatbot, audio_output, image_output]
+    image_output = gr.Image(
+        label="Generated Image"
     )
 
-    clear_btn.click(
-        clear_memory,
-        outputs=chatbot
+    submit_btn = gr.Button("Submit")
+
+    submit_btn.click(
+        fn=ai_assistant,
+        inputs=user_input,
+        outputs=[text_output, audio_output, image_output]
     )
 
+
 demo.launch()