Update app.py

app.py

@@ -1,113 +1,103 @@
 import gradio as gr
-import numpy as np
-import io
-import tempfile
-from pydub import AudioSegment
-from dataclasses import dataclass, field
-
-@dataclass
-class AppState:
-    stream: np.ndarray | None = None
-    sampling_rate: int = 0
-    pause_detected: bool = False
-    stopped: bool = False
-    started_talking: bool = False
-    conversation: list = field(default_factory=list)
-
-# Process audio input and detect pauses
-def process_audio(audio: tuple, state: AppState):
-    if state.stream is None:
-        state.stream = audio[1]
-        state.sampling_rate = audio[0]
+from huggingface_hub import InferenceClient
+import json
+import uuid
+from PIL import Image
+from bs4 import BeautifulSoup
+import requests
+import random
+from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
+from threading import Thread
+import re
+import time
+import torch
+
+# Initialize model and processor
+model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
+processor = LlavaProcessor.from_pretrained(model_id)
+model = LlavaForConditionalGeneration.from_pretrained(model_id).to("cpu")
+
+# Initialize inference clients for different models
+client_gemma = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+client_mixtral = InferenceClient("NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO")
+client_llama = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")
+client_yi = InferenceClient("01-ai/Yi-1.5-34B-Chat")
+
+def search(query):
+    """Performs a Google search and extracts text from the top results."""
+    session = requests.Session()
+    response = session.get(f"https://www.google.com/search?q={query}", 
+                           headers={"User-Agent": "Mozilla/5.0"})
+    soup = BeautifulSoup(response.text, "html.parser")
+    results = []
+    for result in soup.find_all("div", class_="BNeawe vvjwJb AP7Wnd"):
+        text = result.get_text()
+        link = result.find_parent("a")["href"]
+        results.append(f"{text}: {link}")
+    return "\n".join(results[:3])
+
+def llava(inputs, history):
+    """Processes an image and text input with Llava."""
+    image = Image.open(inputs["files"][0]).convert("RGB")
+    prompt = f"<|im_start|>user <image>\n{inputs['text']}<|im_end|>"
+    processed = processor(prompt, image, return_tensors="pt").to("cpu")
+    return processed
+
+def respond(message, history):
+    """Main response function for the chatbot."""
+    if "files" in message and message["files"]:
+        inputs = llava(message, history)
+        streamer = TextIteratorStreamer(skip_prompt=True, skip_special_tokens=True)
+        thread = Thread(target=model.generate, kwargs=dict(inputs=inputs, max_new_tokens=512, streamer=streamer))
+        thread.start()
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            yield buffer
     else:
-        state.stream = np.concatenate((state.stream, audio[1]))
-
-    pause_detected = len(state.stream) > state.sampling_rate * 1
-    state.pause_detected = pause_detected
-
-    if state.pause_detected:
-        return gr.Audio(recording=False), state  # Stop recording
-    return None, state
-
-# Generate response based on input type (text or audio)
-def response(input_data, state: AppState, input_type: str):
-    if input_type == "text":
-        # Ensure text input is handled correctly
-        user_message = input_data.strip()  # Prevent errors from empty inputs
-        if not user_message:
-            return "Please enter a valid message.", state
-
-        state.conversation.append({"role": "user", "content": user_message})
-        bot_response = f"Echo: {user_message}"  # Simulated bot response
-        state.conversation.append({"role": "assistant", "content": bot_response})
-        return bot_response, state
-
-    if input_type == "audio" and state.pause_detected:
-        # Convert audio to WAV and store in conversation history
-        audio_buffer = io.BytesIO()
-        segment = AudioSegment(
-            state.stream.tobytes(),
-            frame_rate=state.sampling_rate,
-            sample_width=state.stream.dtype.itemsize,
-            channels=1 if len(state.stream.shape) == 1 else state.stream.shape[1]
-        )
-        segment.export(audio_buffer, format="wav")
-
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            f.write(audio_buffer.getvalue())
-        state.conversation.append({"role": "user", "content": {"path": f.name, "mime_type": "audio/wav"}})
-
-        chatbot_response = b"Simulated response audio content"
-        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
-            f.write(chatbot_response)
-        state.conversation.append({"role": "assistant", "content": {"path": f.name, "mime_type": "audio/mp3"}})
-
-        yield None, state
+        prompt = [{"role": "user", "content": msg[0]} for msg in history]
+        prompt.append({"role": "user", "content": message["text"]})
+        response = client_gemma.chat_completion(prompt, max_tokens=200)
+        yield response["choices"][0]["message"]["content"]
 
-    return None, state  # Handle unexpected input cases gracefully
+def generate_image(prompt):
+    """Generates an image using the external model."""
+    client = InferenceClient("KingNish/Image-Gen-Pro")
+    return client.predict("Image Generation", None, prompt, api_name="/image_gen_pro")
 
-# Start recording audio input
-def start_recording_user(state: AppState):
-    if not state.stopped:
-        return gr.Audio(recording=True)
-
-# Gradio app setup
+# Set up Gradio interface
 with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
     with gr.Row():
         with gr.Column():
-            input_audio = gr.Audio(label="Input Audio", type="numpy")
-            text_input = gr.Textbox(label="Text Input", placeholder="Type your message here...")
+            text_input = gr.Textbox(placeholder="Enter your message...")
+            file_input = gr.File(label="Upload an image")
         with gr.Column():
-            chatbot = gr.Chatbot(label="Conversation", type="messages")
-            output_audio = gr.Audio(label="Output Audio", streaming=True, autoplay=True)
-
-    state = gr.State(value=AppState())
-
-    # Handle audio input streaming
-    stream = input_audio.stream(
-        process_audio, [input_audio, state], [input_audio, state], stream_every=0.5, time_limit=30
-    )
-
-    # Handle text input submission
-    text_submit = text_input.submit(
-        lambda txt, s: response(txt, s, "text"), [text_input, state], [chatbot, state]
-    )
-
-    # Handle audio stop recording
-    respond = input_audio.stop_recording(
-        lambda s: response(None, s, "audio"), [state], [output_audio, state]
-    )
-    respond.then(lambda s: s.conversation, [state], [chatbot])
-
-    # Restart recording after audio playback ends
-    restart = output_audio.stop(start_recording_user, [state], [input_audio])
-
-    # Stop conversation button
-    cancel = gr.Button("Stop Conversation", variant="stop")
-    cancel.click(
-        lambda: (AppState(stopped=True), gr.Audio(recording=False)),
-        None, [state, input_audio], cancels=[respond, restart]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+            output = gr.Image(label="Generated Image")
+    with gr.Row():
+        search_button = gr.Button("Search Google")
+        image_button = gr.Button("Generate Image")
+    examples = [
+        {"text": "Who are you?"},
+        {"text": "Generate an image of the Eiffel Tower at night."},
+        {"text": "Search for the latest trends on YouTube."},
+    ]
+
+    def handle_text(text, state):
+        response = respond({"text": text}, state)
+        return response, state
+
+    def handle_file_upload(files, state):
+        response = respond({"files": files, "text": "Describe this image."}, state)
+        return response, state
+
+    # Connect components to callbacks
+    text_input.submit(handle_text, [text_input], [chatbot])
+    file_input.change(handle_file_upload, [file_input], [chatbot])
+
+    # Search button functionality
+    search_button.click(lambda query: search(query), [text_input], [chatbot])
+    image_button.click(lambda text: generate_image(text), [text_input], [output])
+
+# Launch the Gradio interface
+demo.launch()