use RealVisXL_V5.0_Lightning

pulid/pipeline.py

@@ -8,6 +8,7 @@ from basicsr.utils import img2tensor, tensor2img
 from diffusers import (
     DPMSolverMultistepScheduler,
     StableDiffusionXLPipeline,
+    AutoPipelineForText2Image,
     UNet2DConditionModel,
 )
 from facexlib.parsing import init_parsing_model
@@ -34,22 +35,28 @@ class PuLIDPipeline:
     def __init__(self, *args, **kwargs):
         super().__init__()
         self.device = 'cuda'
-        sdxl_base_repo = 'stabilityai/stable-diffusion-xl-base-1.0'
-        sdxl_lightning_repo = 'ByteDance/SDXL-Lightning'
-        self.sdxl_base_repo = sdxl_base_repo
+        #sdxl_base_repo = 'stabilityai/stable-diffusion-xl-base-1.0'
+        #sdxl_lightning_repo = 'ByteDance/SDXL-Lightning'
+        #self.sdxl_base_repo = sdxl_base_repo
 
         # load base model
-        unet = UNet2DConditionModel.from_config(sdxl_base_repo, subfolder='unet').to(self.device, torch.float16)
-        unet.load_state_dict(
-            load_file(
-                hf_hub_download(sdxl_lightning_repo, 'sdxl_lightning_4step_unet.safetensors'), device=self.device
-            )
-        )
-        unet.half()
-        self.hack_unet_attn_layers(unet)
-        self.pipe = StableDiffusionXLPipeline.from_pretrained(
-            sdxl_base_repo, unet=unet, torch_dtype=torch.float16, variant="fp16"
-        ).to(self.device)
+        #unet = UNet2DConditionModel.from_config(sdxl_base_repo, subfolder='unet').to(self.device, torch.float16)
+        #unet.load_state_dict(
+        #    load_file(
+        #        hf_hub_download(sdxl_lightning_repo, 'sdxl_lightning_4step_unet.safetensors'), device=self.device
+        #    )
+        #)
+        #unet.half()
+        #self.hack_unet_attn_layers(unet)
+        #self.pipe = StableDiffusionXLPipeline.from_pretrained(
+        #    sdxl_base_repo, unet=unet, torch_dtype=torch.float16, variant="fp16"
+        #).to(self.device)
+        
+        # SG161222/RealVisXL_V5.0_Lightning, lykon/dreamshaper-xl-lightning
+        self.pipe = AutoPipelineForText2Image.from_pretrained('SG161222/RealVisXL_V5.0_Lightning', torch_dtype=torch.float16, variant="fp16")
+        self.hack_unet_attn_layers(self.pipe.unet)
+        pipe = pipe.to(self.device)
+
         self.pipe.watermark = None
 
         # scheduler