Update app.py

app.py

@@ -5,7 +5,6 @@ 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
 import soundfile as sf
 from llama_index.core.agent import ReActAgent
@@ -16,12 +15,12 @@ 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)
 tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_code=True)
@@ -29,15 +28,9 @@ tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_co
 tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1")
 tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-large-v1")
 
-# Image generation model
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "ByteDance/SDXL-Lightning"
-ckpt = "sdxl_lightning_4step_unet.safetensors"
-
-unet = UNet2DConditionModel.from_config(base, subfolder="unet")
-unet.load_state_dict(load_file(hf_hub_download(repo, ckpt)))
-image_pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16")
-image_pipe.scheduler = EulerDiscreteScheduler.from_config(image_pipe.scheduler.config, timestep_spacing="trailing")
+# Updated Image Generation Model
+pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
 
 # Tavily Client for web search
 tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API"))
@@ -79,7 +72,12 @@ def web_search(query):
 
 # Image Generation Tool
 def image_generation(query):
-    image = image_pipe(prompt=query, num_inference_steps=20, guidance_scale=7.5).images[0]
+    image = pipe(
+        query,
+        negative_prompt="",
+        num_inference_steps=28,
+        guidance_scale=7.0,
+    ).images[0]
     image.save("output.jpg")
     return "output.jpg"
 
@@ -97,7 +95,7 @@ def handle_input(user_prompt, image=None, audio=None, websearch=False):
         user_prompt = transcription.text
 
     tools = [
-        FunctionTool.from_defaults(fn=numpy_code_calculator, name="Numpy Code Calculator"),
+        FunctionTool.from_defaults(fn=numpy_code_calculator, name="Numpy"),
         FunctionTool.from_defaults(fn=image_generation, name="Image"),
     ]
     
@@ -166,8 +164,7 @@ def main_interface(user_prompt, image=None, audio=None, voice_only=False, websea
     print("Starting main_interface function")
     vqa_model.to(device='cuda', dtype=torch.bfloat16)
     tts_model.to("cuda")
-    unet.to("cuda")
-    image_pipe.to("cuda")
+    pipe.to("cuda")
 
     print(f"user_prompt: {user_prompt}, image: {image}, audio: {audio}, voice_only: {voice_only}, websearch: {websearch}")
     
@@ -192,4 +189,4 @@ def main_interface(user_prompt, image=None, audio=None, voice_only=False, websea
 
 # Launch the UI
 demo = create_ui()
-demo.launch()
+demo.launch()