API Improvements: Modules Change AND Restore override_settings (#2027)

* Improve API modules change
* Restore override_settings and make it work
* Simplify some memory management
This commit is contained in:
altoiddealer
2024-10-13 07:29:02 -04:00
committed by GitHub
parent ae8187bf2d
commit 862c7a589e
3 changed files with 89 additions and 22 deletions

View File

@@ -32,6 +32,7 @@ from einops import repeat, rearrange
from blendmodes.blend import blendLayers, BlendType
from modules.sd_models import apply_token_merging, forge_model_reload
from modules_forge.utils import apply_circular_forge
from modules_forge import main_entry
from backend import memory_management
@@ -810,11 +811,52 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
if p.scripts is not None:
p.scripts.before_process(p)
# backwards compatibility, fix sampler and scheduler if invalid
sd_samplers.fix_p_invalid_sampler_and_scheduler(p)
stored_opts = {k: opts.data[k] if k in opts.data else opts.get_default(k) for k in p.override_settings.keys() if k in opts.data}
with profiling.Profiler():
res = process_images_inner(p)
try:
# if no checkpoint override or the override checkpoint can't be found, remove override entry and load opts checkpoint
# and if after running refiner, the refiner model is not unloaded - webui swaps back to main model here, if model over is present it will be reloaded afterwards
if sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None:
p.override_settings.pop('sd_model_checkpoint', None)
temp_memory_changes = {}
memory_keys = ['forge_inference_memory', 'forge_async_loading', 'forge_pin_shared_memory']
for k, v in p.override_settings.items():
# options for memory/modules/checkpoints are set in their dedicated functions
if k in memory_keys:
mem_k = k[len('forge_'):] # remove 'forge_' prefix
temp_memory_changes[mem_k] = v
elif k == 'forge_additional_modules':
main_entry.modules_change(v, refresh_params=False)
elif k == 'sd_model_checkpoint':
main_entry.checkpoint_change(v)
# set all other options
else:
opts.set(k, v, is_api=True, run_callbacks=False)
if temp_memory_changes:
main_entry.refresh_memory_management_settings(**temp_memory_changes)
# backwards compatibility, fix sampler and scheduler if invalid
sd_samplers.fix_p_invalid_sampler_and_scheduler(p)
with profiling.Profiler():
res = process_images_inner(p)
finally:
# restore opts to original state
if p.override_settings_restore_afterwards:
for k, v in stored_opts.items():
if k == 'forge_additional_modules':
main_entry.modules_change(v, refresh_params=False)
elif k == 'sd_model_checkpoint':
main_entry.checkpoint_change(v)
else:
setattr(opts, k, v)
if temp_memory_changes:
main_entry.refresh_memory_management_settings() # applies the set options by default
return res