Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 from groq import Groq
 import spaces
 from transformers import AutoModel, AutoTokenizer
+from diffusers import StableDiffusion3Pipeline
 from parler_tts import ParlerTTSForConditionalGeneration
 import soundfile as sf
 from llama_index.core.agent import ReActAgent
@@ -15,20 +16,19 @@ from tavily import TavilyClient
 import requests
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from diffusers import StableDiffusion3Pipeline
 
 # Initialize models and clients
 MODEL = 'llama3-groq-70b-8192-tool-use-preview'
 client = Groq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY"))
 
 vqa_model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_code=True,
-                                       device_map="auto", torch_dtype=torch.bfloat16)
+                                      device_map="auto", torch_dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_code=True)
 
 tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1")
 tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-large-v1")
 
-# Updated Image Generation Model
+# Updated Image generation model
 pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
 
@@ -75,7 +75,7 @@ def image_generation(query):
     image = pipe(
         query,
         negative_prompt="",
-        num_inference_steps=28,
+        num_inference_steps=15,
         guidance_scale=7.0,
     ).images[0]
     image.save("output.jpg")
@@ -111,12 +111,11 @@ def handle_input(user_prompt, image=None, audio=None, websearch=False):
         messages = [{"role": "user", "content": [image, user_prompt]}]
         response = vqa_model.chat(image=None, msgs=messages, tokenizer=tokenizer)
     else:
-        # Modify this part to check if a tool is required or if a direct answer suffices
         response = agent.chat(user_prompt)
     
     # Extract the content from AgentChatResponse to return as a string
     if isinstance(response, AgentChatResponse):
-        response = response.response_text
+        response = response.final_response # Use 'final_response' to access the text response
     
     return response
 
@@ -189,4 +188,4 @@ def main_interface(user_prompt, image=None, audio=None, voice_only=False, websea
 
 # Launch the UI
 demo = create_ui()
-demo.launch()
+demo.launch()