| 89 | return text_embeddings, uncond_embeddings, cond_embeddings |
| 90 | |
| 91 | def process_input_embeddings(input_embeddings): |
| 92 | assert isinstance(input_embeddings, (tuple, list)) |
| 93 | if len(input_embeddings) == 3: |
| 94 | # input_embeddings: text_embeddings, uncond_embeddings, cond_embeddings |
| 95 | # Assume `uncond_embeddings` is full (has batch size the same as cond_embeddings) |
| 96 | _, uncond_embeddings, cond_embeddings = input_embeddings |
| 97 | assert uncond_embeddings.shape[0] == cond_embeddings.shape[0], f"{uncond_embeddings.shape[0]} != {cond_embeddings.shape[0]}" |
| 98 | return input_embeddings |
| 99 | elif len(input_embeddings) == 2: |
| 100 | # input_embeddings: uncond_embeddings, cond_embeddings |
| 101 | # uncond_embeddings may have only one item |
| 102 | uncond_embeddings, cond_embeddings = input_embeddings |
| 103 | if uncond_embeddings.shape[0] == 1: |
| 104 | uncond_embeddings = uncond_embeddings.expand(cond_embeddings.shape) |
| 105 | # We follow the convention: negative (unconditional) prompt comes first |
| 106 | text_embeddings = torch.cat((uncond_embeddings, cond_embeddings), dim=0) |
| 107 | return text_embeddings, uncond_embeddings, cond_embeddings |
| 108 | else: |
| 109 | raise ValueError(f"input_embeddings length: {len(input_embeddings)}") |