(self,
prompt,
enable_recap=True,
width=832,
height=1216,
num_inference_steps=28,
guidance_scale=3.5,
seed=None)
| 137 | self.pipeline.to(self.device) |
| 138 | |
| 139 | def generate(self, |
| 140 | prompt, |
| 141 | enable_recap=True, |
| 142 | width=832, |
| 143 | height=1216, |
| 144 | num_inference_steps=28, |
| 145 | guidance_scale=3.5, |
| 146 | seed=None): |
| 147 | |
| 148 | # Prompt rewriting |
| 149 | if enable_recap and self.qwen_agent: |
| 150 | final_prompt = self.qwen_agent.recap_prompt(prompt) |
| 151 | else: |
| 152 | final_prompt = prompt |
| 153 | |
| 154 | # Poster generation |
| 155 | if seed is None: |
| 156 | seed = random.randint(1, 2**32 - 1) |
| 157 | |
| 158 | generator = torch.Generator(device=self.device).manual_seed(seed) |
| 159 | |
| 160 | with torch.inference_mode(): |
| 161 | image = self.pipeline( |
| 162 | prompt=final_prompt, |
| 163 | generator=generator, |
| 164 | num_inference_steps=num_inference_steps, |
| 165 | guidance_scale=guidance_scale, |
| 166 | width=width, |
| 167 | height=height |
| 168 | ).images[0] |
| 169 | |
| 170 | return image, final_prompt, seed |
| 171 | |
| 172 | |
| 173 | def main(): |
no test coverage detected