wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
# Initially pulled from https://github.com/black-forest-labs/flux
|
|
|
|
from torch import Tensor, nn
|
|
from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
|
|
|
|
from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
|
|
|
|
|
|
class HFEncoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
encoder: PreTrainedModel,
|
|
tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
|
|
is_clip: bool,
|
|
max_length: int,
|
|
):
|
|
super().__init__()
|
|
self.max_length = max_length
|
|
self.is_clip = is_clip
|
|
self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
|
|
self.tokenizer = tokenizer
|
|
self.hf_module = encoder
|
|
self.hf_module = self.hf_module.eval().requires_grad_(False)
|
|
|
|
def forward(self, text: list[str]) -> Tensor:
|
|
batch_encoding = self.tokenizer(
|
|
text,
|
|
truncation=True,
|
|
max_length=self.max_length,
|
|
return_length=False,
|
|
return_overflowing_tokens=False,
|
|
padding="max_length",
|
|
return_tensors="pt",
|
|
)
|
|
|
|
# Move inputs to the same device as the model to support cpu_only models
|
|
model_device = get_effective_device(self.hf_module)
|
|
|
|
outputs = self.hf_module(
|
|
input_ids=batch_encoding["input_ids"].to(model_device),
|
|
attention_mask=None,
|
|
output_hidden_states=False,
|
|
)
|
|
return outputs[self.output_key]
|