lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-04-29 02:31:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

support all flux models

Browse Source
This commit is contained in:
lllyasviel
2024-08-13 05:42:17 -07:00
committed by GitHub
parent 3589b57ec1
commit 61f83dd610
8 changed files with 177 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
modules/sd_models.py
View File

@@ -490,27 +490,15 @@ def forge_model_reload():
timer.record("unload existing model")
checkpoint_info = model_data.forge_loading_parameters['checkpoint_info']
state_dict = load_torch_file(checkpoint_info.filename)
timer.record("load state dict")
state_dict_vae = model_data.forge_loading_parameters.get('vae_filename', None)
if state_dict_vae is not None:
state_dict_vae = load_torch_file(state_dict_vae)
timer.record("load vae state dict")
if shared.opts.sd_checkpoint_cache > 0:
# cache newly loaded model
checkpoints_loaded[checkpoint_info] = state_dict.copy()
state_dict = checkpoint_info.filename
additional_state_dicts = model_data.forge_loading_parameters.get('additional_modules', [])
timer.record("cache state dict")
dynamic_args['forge_unet_storage_dtype'] = model_data.forge_loading_parameters.get('unet_storage_dtype', None)
dynamic_args['embedding_dir'] = cmd_opts.embeddings_dir
dynamic_args['emphasis_name'] = opts.emphasis
sd_model = forge_loader(state_dict, sd_vae=state_dict_vae)
del state_dict
sd_model = forge_loader(state_dict, additional_state_dicts=additional_state_dicts)
timer.record("forge model load")
sd_model.extra_generation_params = {}
@@ -520,10 +508,6 @@ def forge_model_reload():
sd_model.sd_model_hash = checkpoint_info.calculate_shorthash()
timer.record("calculate hash")
# clean up cache if limit is reached
while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:
checkpoints_loaded.popitem(last=False)
shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
model_data.set_sd_model(sd_model)