Load custom UNet2DConditionModel takes too long

I’m loading a UNet2DConditionModel safe-tensors which are results from this text_to_image training. The code takes way too long to load.

unet = UNet2DConditionModel.from_pretrained('checkpoint-3000', subfolder='unet_ema', use_safetensors=True, torch_dtype=torch.float16)
pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", unet=unet, torch_dtype=torch.float16, use_safetensors=True)

Is there something wrong with my installation or my code?