Loading weight time is too long

[Un1Lee]

Hello, I found the time is too long for load the weight for reference_unet,denoising_unet and so on. For example, it takes more than 10 mins to load reference_unet. But the loading for vae is quick. What's the problem it maybe? I have download the models needed and put them in 'pretrained_model' dir.

reference_unet = UNet2DConditionModel.from_pretrained( config.pretrained_base_model_path, subfolder="unet", ).to(dtype=weight_dtype, device="cuda") reference_unet.load_state_dict( torch.load(config.reference_unet_path, map_location="cpu"), ) Both of them waste time. In fact, UNet2DConditionModel init is quick. after the unet_2d_condition.py 655, it falls slow. I think it is going on from_pretrained(). It's normal?

I found it's the pretrained. if device_map is None: param_device = "cpu" state_dict = load_state_dict(model_file, variant=variant) where model_file is './pretrained_model/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin' the load_state_dict is really very very slow.

[Zejun-Yang]

Thank you for your interest. I'm sorry about that! Perhaps you can wait for our upcoming hugging face demo page.