Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import os
 import numpy as np
 from groq import Groq
-import spaces
 from transformers import AutoModel, AutoTokenizer
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from parler_tts import ParlerTTSForConditionalGeneration
@@ -12,13 +11,14 @@ from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma 
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import RetrievalQA
-from langchain_community.llms import OpenAI 
 from PIL import Image
 from decord import VideoReader, cpu
+from tavily import TavilyClient
 import requests
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
+# Initialize models
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 MODEL = 'llama3-groq-70b-8192-tool-use-preview'
 
@@ -39,7 +39,10 @@ unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
 image_pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
 image_pipe.scheduler = EulerDiscreteScheduler.from_config(image_pipe.scheduler.config, timestep_spacing="trailing")
 
-# Initialize voice-only mode
+# Tavily Client
+tavily_client = TavilyClient(api_key="tvly-YOUR_API_KEY")
+
+# Voice output function
 def play_voice_output(response):
     description = "Jon's voice is monotone yet slightly fast in delivery, with a very close recording that almost has no background noise."
     input_ids = tts_tokenizer(description, return_tensors="pt").input_ids.to('cuda')
@@ -49,18 +52,6 @@ def play_voice_output(response):
     sf.write("output.wav", audio_arr, tts_model.config.sampling_rate)
     return "output.wav"
 
-# Web search function
-def web_search(query):
-    api_key = os.environ.get("BING_API_KEY")
-    search_url = "https://api.bing.microsoft.com/v7.0/search"
-    headers = {"Ocp-Apim-Subscription-Key": api_key}
-    params = {"q": query, "textDecorations": True, "textFormat": "HTML"}
-    response = requests.get(search_url, headers=headers, params=params)
-    response.raise_for_status()
-    search_results = response.json()
-    snippets = [result['snippet'] for result in search_results.get('webPages', {}).get('value', [])]
-    return "\n".join(snippets)
-
 # NumPy Calculation function
 def numpy_calculate(code: str) -> str:
     try:
@@ -71,37 +62,6 @@ def numpy_calculate(code: str) -> str:
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-# Function to handle different input types
-def handle_input(user_prompt, image=None, video=None, audio=None, doc=None):
-    messages = [{"role": "user", "content": user_prompt}]
-
-    if audio:
-        transcription = client.audio.transcriptions.create(
-            file=(audio.name, audio.read()),
-            model="whisper-large-v3"
-        )
-        user_prompt = transcription.text
-
-    if doc:
-        # RAG with Langchain
-        response = use_langchain_rag(doc.name, doc.read(), user_prompt)
-    elif image and not video:
-        image = Image.open(image).convert('RGB')
-        messages[0]['content'] = [image, user_prompt]
-        response = text_model.chat(image=None, msgs=messages, tokenizer=tokenizer)
-    elif video:
-        frames = encode_video(video.name)
-        messages[0]['content'] = frames + [user_prompt]
-        response = text_model.chat(image=None, msgs=messages, tokenizer=tokenizer)
-    else:
-        response = client.chat.completions.create(
-            model=MODEL,
-            messages=messages,
-            tools=initialize_tools()
-        ).choices[0].message.content
-
-    return response
-
 # Function to use Langchain for RAG
 def use_langchain_rag(file_name, file_content, query):
     # Split the document into chunks
@@ -130,64 +90,58 @@ def encode_video(video_path):
     frames = [Image.fromarray(v.astype('uint8')) for v in frames]
     return frames
 
-# Initialize tools with web search and NumPy calculation
-def initialize_tools():
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "calculate",
-                "description": "Evaluate a mathematical expression",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "expression": {"type": "string", "description": "The mathematical expression to evaluate"}
-                    },
-                    "required": ["expression"]
-                },
-            }
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "web_search",
-                "description": "Perform a web search",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "query": {"type": "string", "description": "The search query"}
-                    },
-                    "required": ["query"]
-                },
-                "implementation": web_search 
-            }
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "numpy_calculate",
-                "description": "Execute NumPy-based Python code for calculations",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "code": {"type": "string", "description": "The Python code with NumPy operations"}
-                    },
-                    "required": ["code"]
-                },
-                "implementation": numpy_calculate
-            }
-        }
-    ]
-    return tools
+# Web search function
+def web_search(query):
+    answer = tavily_client.qna_search(query=query)
+    return answer
+
+# Function to handle different input types
+def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, websearch=False):
+    # Voice input handling
+    if audio:
+        transcription = client.audio.transcriptions.create(
+            file=(audio.name, audio.read()),
+            model="whisper-large-v3"
+        )
+        user_prompt = transcription.text
+
+    # If user uploaded an image and text, use MiniCPM model
+    if image:
+        image = Image.open(image).convert('RGB')
+        messages = [{"role": "user", "content": [image, user_prompt]}]
+        response = text_model.chat(image=None, msgs=messages, tokenizer=tokenizer)
+        return response
+
+    # Determine which tool to use
+    if doc:
+        file_content = doc.read().decode('utf-8')
+        response = use_langchain_rag(doc.name, file_content, user_prompt)
+    elif "calculate" in user_prompt.lower():
+        response = numpy_calculate(user_prompt)
+    elif "generate" in user_prompt.lower() and ("image" in user_prompt.lower() or "picture" in user_prompt.lower()):
+        response = image_pipe(prompt=user_prompt, num_inference_steps=20, guidance_scale=7.5)
+    elif websearch:
+        response = web_search(user_prompt)
+    else:
+        chat_completion = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": user_prompt}
+            ],
+            model=MODEL,
+        )
+        response = chat_completion.choices[0].message.content
+
+    return response
 
 @spaces.GPU()
-def main_interface(user_prompt, image=None, video=None, audio=None, doc=None, voice_only=False):
+def main_interface(user_prompt, image=None, video=None, audio=None, doc=None, voice_only=False, websearch=False):
     text_model.to(device='cuda', dtype=torch.bfloat16)
     tts_model.to("cuda")
     unet.to("cuda", torch.float16)
     image_pipe.to("cuda")
     
-    response = handle_input(user_prompt, image=image, video=video, audio=audio, doc=doc)
+    response = handle_input(user_prompt, image=image, video=video, audio=audio, doc=doc, websearch=websearch)
     
     if voice_only:
         audio_file = play_voice_output(response)
@@ -195,22 +149,46 @@ def main_interface(user_prompt, image=None, video=None, audio=None, doc=None, vo
     else:
         return response, None  # Return only the text output, no audio
 
-# Gradio App Setup
-with gr.Blocks() as demo:
-    user_prompt = gr.Textbox(placeholder="Type your message here...", lines=1)
-    image_input = gr.Image(type="filepath", label="Upload an image")
-    video_input = gr.Video(label="Upload a video")
-    audio_input = gr.Audio(type="filepath", label="Upload audio")
-    doc_input = gr.File(type="filepath", label="Upload a document")
-    voice_only_mode = gr.Checkbox(label="Enable Voice Only Mode")
-    output_label = gr.Label(label="Output")
-    audio_output = gr.Audio(label="Audio Output", visible=False)
-    
-    submit = gr.Button("Submit")
-    submit.click(
-        fn=main_interface,
-        inputs=[user_prompt, image_input, video_input, audio_input, doc_input, voice_only_mode],
-        outputs=[output_label, audio_output]  # Expecting a string and audio file
-    )
+# Gradio UI Setup
+def create_ui():
+    with gr.Blocks() as demo:
+        gr.Markdown("# AI Assistant")
+        with gr.Row():
+            with gr.Column(scale=2):
+                user_prompt = gr.Textbox(placeholder="Type your message here...", lines=1)
+            with gr.Column(scale=1):
+                image_input = gr.Image(type="filepath", label="Upload an image", elem_id="image-icon")
+                video_input = gr.Video(label="Upload a video", elem_id="video-icon")
+                audio_input = gr.Audio(type="filepath", label="Upload audio", elem_id="mic-icon")
+                doc_input = gr.File(type="filepath", label="Upload a document", elem_id="document-icon")
+                voice_only_mode = gr.Checkbox(label="Enable Voice Only Mode", elem_id="voice-only-mode")
+                websearch_mode = gr.Checkbox(label="Enable Web Search", elem_id="websearch-mode")
+            with gr.Column(scale=1):
+                submit = gr.Button("Submit")
+        
+        output_label = gr.Label(label="Output")
+        audio_output = gr.Audio(label="Audio Output", visible=False)
+
+        submit.click(
+            fn=main_interface,
+            inputs=[user_prompt, image_input, video_input, audio_input, doc_input, voice_only_mode, websearch_mode],
+            outputs=[output_label, audio_output]  # Expecting a string and audio file
+        )
+
+        # Voice-only mode UI
+        voice_only_mode.change(
+            lambda x: gr.update(visible=not x),
+            inputs=voice_only_mode,
+            outputs=[user_prompt, image_input, video_input, doc_input, websearch_mode, submit]
+        )
+        voice_only_mode.change(
+            lambda x: gr.update(visible=x),
+            inputs=voice_only_mode,
+            outputs=[audio_input]
+        )
+
+    return demo
 
+# Launch the app
+demo = create_ui()
 demo.launch(inline=False)