| from transformers import PreTrainedModel, PretrainedConfig | |
| from .module import ConditionalViT | |
| from sentence_transformers import SentenceTransformer | |
| class CondViTConfig(PretrainedConfig): | |
| model_type = "condvit" | |
| def __init__( | |
| self, | |
| input_resolution: int = 224, | |
| patch_size: int = 16, | |
| width: int = 768, | |
| layers: int = 12, | |
| heads: int = 12, | |
| output_dim: int = 512, | |
| n_categories: int = 10, | |
| lm_backbone: str = "sentence-transformers/sentence-t5-xl", | |
| lm_revision: str = "e0976ba9afd18be963c22c680367a3928c44fd22", | |
| device: str = "cpu", | |
| **kwargs | |
| ): | |
| self.input_resolution = input_resolution | |
| self.patch_size = patch_size | |
| self.width = width | |
| self.layers = layers | |
| self.heads = heads | |
| self.output_dim = output_dim | |
| self.n_categories = n_categories | |
| self.lm_backbone = lm_backbone | |
| self.lm_revision = lm_revision | |
| self.device = device | |
| super().__init__(**kwargs) | |
| class CondViTForEmbedding(PreTrainedModel): | |
| config_class = CondViTConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.condvit = ConditionalViT( | |
| input_resolution=config.input_resolution, | |
| patch_size=config.patch_size, | |
| width=config.width, | |
| layers=config.layers, | |
| heads=config.heads, | |
| output_dim=config.output_dim, | |
| ) | |
| if config.device: | |
| self.condvit.to(config.device) | |
| self.lm = SentenceTransformer( | |
| config.lm_backbone, revision=config.lm_revision, device=config.device | |
| ) | |
| def forward(self, pixel_values, texts=None): | |
| if texts is not None: | |
| text_embeddings = self.lm.encode( | |
| texts, | |
| convert_to_tensor=True, | |
| convert_to_numpy=False, | |
| ) | |
| text_embeddings = text_embeddings.to(pixel_values.device) | |
| else: | |
| text_embeddings = None | |
| return self.condvit(imgs=pixel_values, c=text_embeddings) | |