Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
-from diffusers import DiffusionPipeline
 import torch
+from diffusers import KolorsPipeline
 
-# Load the Kwai-Kolors diffusion model
-pipeline = DiffusionPipeline.from_pretrained("Kwai-Kolors/Kolors").to("cuda" if torch.cuda.is_available() else "cpu")
+pipeline= KolorsPipeline.from_pretrained(
+    "Kwai-Kolors/Kolors-diffusers", 
+    torch_dtype=torch.float16, 
+    variant="fp16"
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
 # Function to generate image from prompt
 def generate_image(prompt):