/
ai-forever
/
diffusers-new
Обзор
Документация
Войти
/
ai-forever
/
diffusers-new
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
main
src/diffusers/pipelines/kandinsky/text_encoder.py
27 строк
1022 B
YiYi Xu
Add Kandinsky 2.1 (#3308)
26 май 2023, 00:28
Не верифицирован
26 май 2023, 00:28
03b7a84
Код
Авторство
О чём код?
import torch from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel class MCLIPConfig(XLMRobertaConfig): model_type = "M-CLIP" def __init__(self, transformerDimSize=1024, imageDimSize=768, **kwargs): self.transformerDimensions = transformerDimSize self.numDims = imageDimSize super().__init__(**kwargs) class MultilingualCLIP(PreTrainedModel): config_class = MCLIPConfig def __init__(self, config, *args, **kwargs): super().__init__(config, *args, **kwargs) self.transformer = XLMRobertaModel(config) self.LinearTransformation = torch.nn.Linear( in_features=config.transformerDimensions, out_features=config.numDims ) def forward(self, input_ids, attention_mask): embs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)[0] embs2 = (embs * attention_mask.unsqueeze(2)).sum(dim=1) / attention_mask.sum(dim=1)[:, None] return self.LinearTransformation(embs2), embs