| import logging |
| from collections import OrderedDict |
| from typing import List, Union |
|
|
| import torch |
| from torch import nn |
|
|
| from .clip_wrapper import build_clip_text_encoder, get_clip_embeddings |
| from .clip_wrapper_open import build_openclip_text_encoder, get_openclip_embeddings |
|
|
|
|
| class TextModel(nn.Module): |
| def __init__( |
| self, |
| model_type, |
| model_name, |
| model_path, |
| ): |
| super().__init__() |
|
|
| self.model_type = model_type |
| self.model_name = model_name |
| self.model_path = model_path |
|
|
| if self.model_type == "CLIP": |
| self.model = build_clip_text_encoder(model_path, pretrain=True) |
|
|
| if self.model_type == "OPENCLIP": |
| self.model, self.tokenizer = build_openclip_text_encoder(model_name, model_path) |
|
|
| self.model.eval() |
|
|
| def forward_text(self, text, prompt="a "): |
| if self.model_type == "CLIP": |
| return get_clip_embeddings(self.model, text, prompt) |
|
|
| if self.model_type == "OPENCLIP": |
| return get_openclip_embeddings(self.model, self.tokenizer, text, prompt) |
|
|
