86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
import torch
|
|
from toolkit.memory_management import MemoryManager
|
|
from typing import TYPE_CHECKING
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from toolkit.models.base_model import BaseModel
|
|
|
|
|
|
class FakeTextEncoder(torch.nn.Module):
|
|
def __init__(self, device, dtype):
|
|
super().__init__()
|
|
# register a dummy parameter to avoid errors in some cases
|
|
self.dummy_param = torch.nn.Parameter(torch.zeros(1))
|
|
self._device = device
|
|
self._dtype = dtype
|
|
|
|
def forward(self, *args, **kwargs):
|
|
raise NotImplementedError(
|
|
"This is a fake text encoder and should not be used for inference."
|
|
)
|
|
return None
|
|
|
|
@property
|
|
def device(self):
|
|
return self._device
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self._dtype
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
|
|
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
|
|
# to functions and what not will work.
|
|
|
|
if model.text_encoder is not None:
|
|
if isinstance(model.text_encoder, list):
|
|
pipe = model.pipeline
|
|
freed = set()
|
|
|
|
def _free(te):
|
|
if te is None or isinstance(te, FakeTextEncoder) or id(te) in freed:
|
|
return
|
|
MemoryManager.free(te)
|
|
freed.add(id(te))
|
|
|
|
def _new_fake():
|
|
return FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype)
|
|
|
|
# model.text_encoder is the source of truth; some pipelines (self-referencing
|
|
# or built with text_encoder=None) never expose the TE as an attribute
|
|
for real_te in model.text_encoder:
|
|
_free(real_te)
|
|
fakes = [_new_fake() for _ in model.text_encoder]
|
|
model.text_encoder = fakes
|
|
|
|
def _fake(idx):
|
|
return fakes[idx] if idx < len(fakes) else _new_fake()
|
|
|
|
# the pipeline stores text encoders like text_encoder, text_encoder_2, text_encoder_3, etc.
|
|
if getattr(pipe, "text_encoder", None) is not None:
|
|
_free(pipe.text_encoder)
|
|
pipe.text_encoder = _fake(0)
|
|
|
|
i = 2
|
|
while hasattr(pipe, f"text_encoder_{i}"):
|
|
real_te = getattr(pipe, f"text_encoder_{i}")
|
|
if real_te is not None:
|
|
_free(real_te)
|
|
setattr(pipe, f"text_encoder_{i}", _fake(i - 1))
|
|
i += 1
|
|
else:
|
|
# only has a single text encoder
|
|
MemoryManager.free(model.text_encoder)
|
|
model.text_encoder = FakeTextEncoder(
|
|
device=model.device_torch,
|
|
dtype=model.torch_dtype
|
|
)
|
|
|
|
MemoryManager.release_cached_memory()
|