Update app.py

app.py

@@ -6,9 +6,9 @@ from safetensors.torch import load_file
 
 # Define a function to generate the image
 def generate_image(prompt, num_inference_steps):
-    base = "stable-diffusion"
-    repo = "Geek7/testing"
-    ckpt = "dreamshaper_8.safetensors" # Use the correct ckpt for your step setting!
+    base = "stabilityai/stable-diffusion-xl-base-1.0"
+    repo = "ByteDance/SDXL-Lightning"
+    ckpt = "sdxl_lightning_2step_unet.safetensors" # Use the correct ckpt for your step setting!
 
     # Load model.
     unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)