Update readme

README.md

@@ -25,7 +25,7 @@ Please always use the correct checkpoint for the corresponding inference steps.
 
 ```python
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from huggingface_hub import hf_hub_download
 
 base = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -33,8 +33,9 @@ repo = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_4step_unet.pth" # Use the correct ckpt for your step setting!
 
 # Load model.
-pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
-pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
 # Ensure sampler uses "trailing" timesteps.
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
@@ -72,7 +73,7 @@ The 1-step model uses "sample" prediction instead of "epsilon" prediction! The s
 
 ```python
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from huggingface_hub import hf_hub_download
 
 base = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -80,8 +81,9 @@ repo = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_1step_unet_x0.pth" # Use the correct ckpt for your step setting!
 
 # Load model.
-pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
-pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
 # Ensure sampler uses "trailing" timesteps and "sample" prediction type.
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample")
@@ -112,4 +114,4 @@ Please use Euler sampler with sgm_uniform scheduler.
 
 ### 1-Step UNet
 
-ComfyUI does not support changing model formulation to x0-prediction, so it is not usable in ComfyUI yet. Hopefully ComfyUI gets updated soon.
+ComfyUI does not support changing the model formulation to x0-prediction, so it is not usable in ComfyUI yet. Hopefully, ComfyUI will get updated soon.