Update app.py

app.py

@@ -16,11 +16,12 @@ from langchain_community.llms import OpenAI
 from PIL import Image
 from decord import VideoReader, cpu
 import requests
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 MODEL = 'llama3-groq-70b-8192-tool-use-preview'
 
-# Load MiniCPM-V-2_6 with 4-bit quantization
 text_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)
@@ -28,9 +29,15 @@ 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_model = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet")
-image_pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", unet=image_model, torch_dtype=torch.float16)
-image_pipe.scheduler = EulerDiscreteScheduler.from_pretrained(image_pipe.scheduler.config, timestep_spacing="trailing")
+# Corrected image model and pipeline setup
+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").to("cuda", torch.float16)
+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
 def play_voice_output(response):
@@ -177,7 +184,7 @@ def initialize_tools():
 def main_interface(user_prompt, image=None, video=None, audio=None, doc=None, voice_only=False):
     text_model.to(device='cuda', dtype=torch.bfloat16)
     tts_model.to("cuda")
-    image_model.to("cuda", torch.float16)
+    unet.to("cuda", torch.float16)
     image_pipe.to("cuda")
     response = handle_input(user_prompt, image=image, video=video, audio=audio, doc=doc)
     if voice_only: