diff options
| author | unknown <mcgpapu@gmail.com> | 2022-12-25 02:03:55 -0600 |
|---|---|---|
| committer | unknown <mcgpapu@gmail.com> | 2022-12-25 02:03:55 -0600 |
| commit | 876da1259965130603f2a7fea505cfa0fce09e2e (patch) | |
| tree | ccb8b89d64480a4bd224b311702ffeb13b8fe754 /modules/lowvram.py | |
| parent | d6fdfde9d70f1b86b696240fb0a0c8f2a4d024f6 (diff) | |
| parent | c6f347b81f584b6c0d44af7a209983284dbb52d2 (diff) | |
Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui
Diffstat (limited to 'modules/lowvram.py')
| -rw-r--r-- | modules/lowvram.py | 12 |
1 files changed, 7 insertions, 5 deletions
diff --git a/modules/lowvram.py b/modules/lowvram.py index aa464a95..042a0254 100644 --- a/modules/lowvram.py +++ b/modules/lowvram.py @@ -55,18 +55,20 @@ def setup_for_low_vram(sd_model, use_medvram): if hasattr(sd_model.cond_stage_model, 'model'):
sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model
- # remove three big modules, cond, first_stage, and unet from the model and then
+ # remove four big modules, cond, first_stage, depth (if applicable), and unet from the model and then
# send the model to GPU. Then put modules back. the modules will be in CPU.
- stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model
- sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = None, None, None
+ stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, getattr(sd_model, 'depth_model', None), sd_model.model
+ sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.model = None, None, None, None
sd_model.to(devices.device)
- sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = stored
+ sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.model = stored
- # register hooks for those the first two models
+ # register hooks for those the first three models
sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)
sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)
sd_model.first_stage_model.encode = first_stage_model_encode_wrap
sd_model.first_stage_model.decode = first_stage_model_decode_wrap
+ if sd_model.depth_model:
+ sd_model.depth_model.register_forward_pre_hook(send_me_to_gpu)
parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
if hasattr(sd_model.cond_stage_model, 'model'):
|