Fix text encoder offload bug when caching embeddings (#868)

toolkit/unloader.py

@@ -1,5 +1,7 @@
+import gc
 import torch
 from toolkit.basic import flush
+from toolkit.memory_management import MemoryManager
 from typing import TYPE_CHECKING
 
 
@@ -24,15 +26,21 @@ class FakeTextEncoder(torch.nn.Module):
     @property
     def device(self):
         return self._device
-    
+
     @property
     def dtype(self):
         return self._dtype
-    
+
     def to(self, *args, **kwargs):
         return self
 
 
+def _detach_and_cpu(te: torch.nn.Module):
+    MemoryManager.detach(te)
+    # bypass any nopped-out .to() override and force an actual CPU move
+    torch.nn.Module.to(te, 'cpu')
+
+
 def unload_text_encoder(model: "BaseModel"):
     # unload the text encoder in a way that will work with all models and will not throw errors
     # we need to make it appear as a text encoder module without actually having one so all
@@ -45,13 +53,15 @@ def unload_text_encoder(model: "BaseModel"):
 
             # the pipeline stores text encoders like text_encoder, text_encoder_2, text_encoder_3, etc.
             if hasattr(pipe, "text_encoder"):
+                _detach_and_cpu(pipe.text_encoder)
                 te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype)
                 text_encoder_list.append(te)
-                pipe.text_encoder.to('cpu')
                 pipe.text_encoder = te
 
             i = 2
             while hasattr(pipe, f"text_encoder_{i}"):
+                real_te = getattr(pipe, f"text_encoder_{i}")
+                _detach_and_cpu(real_te)
                 te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype)
                 text_encoder_list.append(te)
                 setattr(pipe, f"text_encoder_{i}", te)
@@ -59,6 +69,12 @@ def unload_text_encoder(model: "BaseModel"):
             model.text_encoder = text_encoder_list
         else:
             # only has a single text encoder
-            model.text_encoder = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype)
+            _detach_and_cpu(model.text_encoder)
+            model.text_encoder = FakeTextEncoder(
+                device=model.device_torch,
+                dtype=model.torch_dtype
+            )
 
+    torch.cuda.empty_cache()
+    gc.collect()
     flush()