Unverified Commit 863e7416 authored by Symbiomatrix's avatar Symbiomatrix Committed by GitHub
Browse files

Bugfix for SDXL model loading in low ram system. (#4628)



Update convert_from_ckpt.py
Co-authored-by: default avatarPatrick von Platen <patrick.v.platen@gmail.com>
parent 24c5e770
...@@ -1387,8 +1387,9 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -1387,8 +1387,9 @@ def download_from_original_stable_diffusion_ckpt(
unet = UNet2DConditionModel(**unet_config) unet = UNet2DConditionModel(**unet_config)
if is_accelerate_available(): if is_accelerate_available():
for param_name, param in converted_unet_checkpoint.items(): if model_type not in ["SDXL", "SDXL-Refiner"]: # SBM Delay this.
set_module_tensor_to_device(unet, param_name, "cpu", value=param) for param_name, param in converted_unet_checkpoint.items():
set_module_tensor_to_device(unet, param_name, "cpu", value=param)
else: else:
unet.load_state_dict(converted_unet_checkpoint) unet.load_state_dict(converted_unet_checkpoint)
...@@ -1588,6 +1589,11 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -1588,6 +1589,11 @@ def download_from_original_stable_diffusion_ckpt(
checkpoint, config_name, prefix="conditioner.embedders.1.model.", has_projection=True, **config_kwargs checkpoint, config_name, prefix="conditioner.embedders.1.model.", has_projection=True, **config_kwargs
) )
if is_accelerate_available(): # SBM Now move model to cpu.
if model_type in ["SDXL", "SDXL-Refiner"]:
for param_name, param in converted_unet_checkpoint.items():
set_module_tensor_to_device(unet, param_name, "cpu", value=param)
pipe = pipeline_class( pipe = pipeline_class(
vae=vae, vae=vae,
text_encoder=text_encoder, text_encoder=text_encoder,
...@@ -1611,6 +1617,11 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -1611,6 +1617,11 @@ def download_from_original_stable_diffusion_ckpt(
checkpoint, config_name, prefix="conditioner.embedders.0.model.", has_projection=True, **config_kwargs checkpoint, config_name, prefix="conditioner.embedders.0.model.", has_projection=True, **config_kwargs
) )
if is_accelerate_available(): # SBM Now move model to cpu.
if model_type in ["SDXL", "SDXL-Refiner"]:
for param_name, param in converted_unet_checkpoint.items():
set_module_tensor_to_device(unet, param_name, "cpu", value=param)
pipe = StableDiffusionXLImg2ImgPipeline( pipe = StableDiffusionXLImg2ImgPipeline(
vae=vae, vae=vae,
text_encoder=text_encoder, text_encoder=text_encoder,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment