github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Reworked freeing memory manager for removing text encoder completly when not needed.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-09 07:08:24 -06:00
parent 3a28c4b1b7
commit 682b27c6ee
2 changed files with 92 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
toolkit/unloader.py
View File

@@ -1,6 +1,4 @@
import gc
import torch
from toolkit.basic import flush
from toolkit.memory_management import MemoryManager
from typing import TYPE_CHECKING
@@ -35,12 +33,6 @@ class FakeTextEncoder(torch.nn.Module):
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
@@ -52,8 +44,8 @@ def unload_text_encoder(model: "BaseModel"):
pipe = model.pipeline
# 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)
if getattr(pipe, "text_encoder", None) is not None:
MemoryManager.free(pipe.text_encoder)
te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype)
text_encoder_list.append(te)
pipe.text_encoder = te
@@ -61,20 +53,19 @@ def unload_text_encoder(model: "BaseModel"):
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)
if real_te is not None:
MemoryManager.free(real_te)
te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype)
text_encoder_list.append(te)
setattr(pipe, f"text_encoder_{i}", te)
i += 1
model.text_encoder = text_encoder_list
else:
# only has a single text encoder
_detach_and_cpu(model.text_encoder)
MemoryManager.free(model.text_encoder)
model.text_encoder = FakeTextEncoder(
device=model.device_torch,
dtype=model.torch_dtype
)
torch.cuda.empty_cache()
gc.collect()
flush()
MemoryManager.release_cached_memory()