Update app.py

app.py

@@ -22,7 +22,7 @@ import requests
 import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
-from diffusers import AutoencoderKL, DPMSolverMultistepScheduler
+from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, UNet2DConditionModel
 import cv2
 import torch
 import numpy as np
@@ -129,7 +129,16 @@ et = time.time()
 elapsed_time = et - st
 print('Loading VAE took: ', elapsed_time, 'seconds')
 st = time.time()
-pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("rubbrband/albedobaseXL_v21",
+
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_4step_unet.safetensors" 
+
+# Load model.
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(base,
+                                                                 unet=unet,
                                                                  vae=vae,
                                                                  controlnet=[identitynet, zoedepthnet],
                                                                  torch_dtype=torch.float16)