(self, prompt: str)
| 28 | self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path) |
| 29 | |
| 30 | def generate_image(self, prompt: str) -> Image.Image: |
| 31 | |
| 32 | batch_messages = [] |
| 33 | |
| 34 | |
| 35 | messages = [ |
| 36 | {"role": "system", "content": "You are a helpful assistant."}, |
| 37 | {"role": "user", "content": f"Please generate image based on the following caption: {prompt}"} |
| 38 | ] |
| 39 | input_text = self.tokenizer.apply_chat_template( |
| 40 | messages, |
| 41 | tokenize=False, |
| 42 | add_generation_prompt=True) |
| 43 | input_text += f"<im_start><S{self.config.scale}>" |
| 44 | |
| 45 | batch_messages.append(input_text) |
| 46 | |
| 47 | # tokenize as a batch |
| 48 | inputs = self.tokenizer(batch_messages, return_tensors="pt", padding=True, truncation=True, padding_side="left") |
| 49 | |
| 50 | gen_ids, output_image = self.model.generate_images( |
| 51 | inputs.input_ids.to(self.device), |
| 52 | inputs.attention_mask.to(self.device), |
| 53 | max_new_tokens=self.config.seq_len, |
| 54 | do_sample=True, |
| 55 | top_p=self.config.top_p, |
| 56 | top_k=self.config.top_k) |
| 57 | |
| 58 | print(output_image) |
| 59 | return output_image[0] |
| 60 | |
| 61 | |
| 62 | def main(): |
no test coverage detected