diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2022-12-24 09:53:44 +0300 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-12-24 09:53:44 +0300 |
| commit | 94450b88775ee14392d05ef6a8fc915b93946b9e (patch) | |
| tree | 800b9b28660402d53f1baca8adb30f8d0f1ab6e5 /modules/lowvram.py | |
| parent | 9441c28c947588d756e279a8cd5db6c0b4a8d2e4 (diff) | |
| parent | ec0a48826fb41c1b1baab45a9030f7eb55568fd0 (diff) | |
Merge pull request #5589 from MrCheeze/better-special-model-support
Better support for 2.0-inpainting and 2.0-depth special models
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'):
|