Merge branch 'master' into master

1 files changed, 32 insertions, 18 deletions

diff --git a/modules/lowvram.py b/modules/lowvram.py
index 7eba1349..042a0254 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -1,9 +1,8 @@
 import torch

-from modules.devices import get_optimal_device

+from modules import devices

 

 module_in_gpu = None

 cpu = torch.device("cpu")

-device = gpu = get_optimal_device()

 

 

 def send_everything_to_cpu():

@@ -33,34 +32,49 @@ def setup_for_low_vram(sd_model, use_medvram):
         if module_in_gpu is not None:

             module_in_gpu.to(cpu)

 

-        module.to(gpu)

+        module.to(devices.device)

         module_in_gpu = module

 

     # see below for register_forward_pre_hook;

     # first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is

     # useless here, and we just replace those methods

-    def first_stage_model_encode_wrap(self, encoder, x):

-        send_me_to_gpu(self, None)

-        return encoder(x)

 

-    def first_stage_model_decode_wrap(self, decoder, z):

-        send_me_to_gpu(self, None)

-        return decoder(z)

+    first_stage_model = sd_model.first_stage_model

+    first_stage_model_encode = sd_model.first_stage_model.encode

+    first_stage_model_decode = sd_model.first_stage_model.decode

 

-    # remove three big modules, cond, first_stage, and unet from the model and then

+    def first_stage_model_encode_wrap(x):

+        send_me_to_gpu(first_stage_model, None)

+        return first_stage_model_encode(x)

+

+    def first_stage_model_decode_wrap(z):

+        send_me_to_gpu(first_stage_model, None)

+        return first_stage_model_decode(z)

+

+    # for SD1, cond_stage_model is CLIP and its NN is in the tranformer frield, but for SD2, it's open clip, and it's in model field

+    if hasattr(sd_model.cond_stage_model, 'model'):

+        sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model

+

+    # 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

-    sd_model.to(device)

-    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = stored

+    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.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 = lambda x, en=sd_model.first_stage_model.encode: first_stage_model_encode_wrap(sd_model.first_stage_model, en, x)

-    sd_model.first_stage_model.decode = lambda z, de=sd_model.first_stage_model.decode: first_stage_model_decode_wrap(sd_model.first_stage_model, de, z)

+    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'):

+        sd_model.cond_stage_model.model = sd_model.cond_stage_model.transformer

+        del sd_model.cond_stage_model.transformer

+

     if use_medvram:

         sd_model.model.register_forward_pre_hook(send_me_to_gpu)

     else:

@@ -70,7 +84,7 @@ def setup_for_low_vram(sd_model, use_medvram):
         # so that only one of them is in GPU at a time

         stored = diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed

         diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = None, None, None, None

-        sd_model.model.to(device)

+        sd_model.model.to(devices.device)

         diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = stored

 

         # install hooks for bits of third model