Add Diffusers model and VAE variant loading option

This commit is contained in:
Disty0
2023-07-31 14:39:32 +03:00
parent 20d94af063
commit 8ffaea76ba
5 changed files with 23 additions and 6 deletions
+7 -2
View File
@@ -548,8 +548,13 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
"load_connected_pipeline": True # always load end-to-end / connected pipelines
# "use_safetensors": True, # TODO(PVP) - we can't enable this for all checkpoints just yet
}
if devices.dtype == torch.float16:
diffusers_load_config['variant'] = 'fp16'
if shared.opts.diffusers_model_load_variant == 'default':
if devices.dtype == torch.float16:
diffusers_load_config['variant'] = 'fp16'
elif shared.opts.diffusers_model_load_variant == 'fp32':
pass
else:
diffusers_load_config['variant'] = shared.opts.diffusers_model_load_variant
if shared.opts.data.get('sd_model_checkpoint', '') == 'model.ckpt' or shared.opts.data.get('sd_model_checkpoint', '') == '':
shared.opts.data['sd_model_checkpoint'] = "runwayml/stable-diffusion-v1-5"
+7 -2
View File
@@ -181,8 +181,13 @@ def load_vae_diffusers(_model, vae_file=None, vae_source="from unknown source"):
"torch_dtype": devices.dtype_vae,
"use_safetensors": True,
}
if devices.dtype_vae == torch.float16:
diffusers_load_config['variant'] = 'fp16'
if shared.opts.diffusers_vae_load_variant == 'default':
if devices.dtype_vae == torch.float16:
diffusers_load_config['variant'] = 'fp16'
elif shared.opts.diffusers_vae_load_variant == 'fp32':
pass
else:
diffusers_load_config['variant'] = shared.opts.diffusers_vae_load_variant
if shared.opts.diffusers_vae_upcast != 'default':
diffusers_load_config['force_upcast'] = True if shared.opts.diffusers_vae_upcast == 'true' else False
+2
View File
@@ -407,6 +407,8 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
"diffusers_vae_slicing": OptionInfo(True, "Enable VAE slicing"),
"diffusers_vae_tiling": OptionInfo(False, "Enable VAE tiling"),
"diffusers_attention_slicing": OptionInfo(False, "Enable attention slicing"),
"diffusers_model_load_variant": OptionInfo("default", "Diffusers model loading variant", gr.Radio, lambda: {"choices": ['default', 'fp32', 'fp16']}),
"diffusers_vae_load_variant": OptionInfo("default", "Diffusers VAE loading variant", gr.Radio, lambda: {"choices": ['default', 'fp32', 'fp16']}),
# "diffusers_force_zeros": OptionInfo(False, "Force zeros for prompts when empty"),
# "diffusers_aesthetics_score": OptionInfo(6.0, "Require aesthetic score", gr.Slider, {"minimum": 0, "maximum": 10, "step": 0.1}),
}))