| 79 | return args |
| 80 | |
| 81 | def _encode_text_prompt( |
| 82 | tokenizer, |
| 83 | text_encoder, |
| 84 | prompt, |
| 85 | device, |
| 86 | batch_size |
| 87 | ): |
| 88 | assert isinstance(prompt, str) |
| 89 | |
| 90 | text_inputs = tokenizer( |
| 91 | prompt, |
| 92 | padding="max_length", |
| 93 | max_length=tokenizer.model_max_length, |
| 94 | truncation=True, |
| 95 | return_tensors="pt", |
| 96 | ) |
| 97 | text_input_ids = text_inputs.input_ids |
| 98 | untruncated_ids = tokenizer( |
| 99 | prompt, padding="longest", return_tensors="pt" |
| 100 | ).input_ids |
| 101 | |
| 102 | if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( |
| 103 | text_input_ids, untruncated_ids |
| 104 | ): |
| 105 | removed_text = tokenizer.batch_decode( |
| 106 | untruncated_ids[:, tokenizer.model_max_length - 1 : -1] |
| 107 | ) |
| 108 | logger.warning( |
| 109 | "The following part of your input was truncated because CLIP can only handle sequences up to" |
| 110 | f" {tokenizer.model_max_length} tokens: {removed_text}" |
| 111 | ) |
| 112 | |
| 113 | if ( |
| 114 | hasattr(text_encoder.config, "use_attention_mask") |
| 115 | and text_encoder.config.use_attention_mask |
| 116 | ): |
| 117 | attention_mask = text_inputs.attention_mask.to(device) |
| 118 | else: |
| 119 | attention_mask = None |
| 120 | |
| 121 | prompt_embeds = text_encoder( |
| 122 | text_input_ids.to(device), |
| 123 | attention_mask=attention_mask, |
| 124 | ) |
| 125 | prompt_embeds = prompt_embeds[0] |
| 126 | |
| 127 | prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device) |
| 128 | |
| 129 | bs_embed, seq_len, _ = prompt_embeds.shape |
| 130 | prompt_embeds = prompt_embeds.repeat(1, batch_size, 1) |
| 131 | prompt_embeds = prompt_embeds.view( |
| 132 | bs_embed * batch_size, seq_len, -1 |
| 133 | ) |
| 134 | |
| 135 | return prompt_embeds |
| 136 | |
| 137 | def main(args): |
| 138 | args.pretrained_model_name_or_path = "ashawkey/imagedream-ipmv-diffusers" |