| 53 | tokenizer = pipe.tokenizer |
| 54 | |
| 55 | def get_text_embedding(prompts): |
| 56 | with torch.no_grad(): |
| 57 | uc = [] |
| 58 | for text in prompts: |
| 59 | tokens = tokenizer(text, |
| 60 | truncation=True, |
| 61 | max_length=tokenizer.model_max_length, |
| 62 | return_length=True, |
| 63 | return_overflowing_tokens=False, |
| 64 | padding="do_not_pad", |
| 65 | ).input_ids |
| 66 | if 'photo of a' in text[:15]: |
| 67 | print(text) |
| 68 | uc.append(pipe.text_encoder(torch.cuda.LongTensor(tokens).reshape(1,-1))[0][:, 4:].reshape(-1, 768)) |
| 69 | else: |
| 70 | uc.append(pipe.text_encoder(torch.cuda.LongTensor(tokens).reshape(1,-1))[0][:, 1:].reshape(-1, 768)) |
| 71 | return torch.cat(uc, 0).float() |
| 72 | |
| 73 | embeds = {} |
| 74 | count = 1 |