| 103 | |
| 104 | |
| 105 | class PosterGenerator: |
| 106 | def __init__(self, |
| 107 | pipeline_path="black-forest-labs/FLUX.1-dev", |
| 108 | custom_transformer_path=None, |
| 109 | qwen_model_path=None, |
| 110 | device="cuda:0"): |
| 111 | |
| 112 | self.device = torch.device(device) if isinstance(device, str) else device |
| 113 | |
| 114 | # Load Qwen model for prompt rewriting |
| 115 | if qwen_model_path and os.path.exists(qwen_model_path): |
| 116 | print(f"Loading Qwen model from: {qwen_model_path}") |
| 117 | self.qwen_agent = QwenRecapAgent(qwen_model_path, device=self.device) |
| 118 | else: |
| 119 | self.qwen_agent = None |
| 120 | |
| 121 | # Load Flux pipeline |
| 122 | print(f"Loading Flux pipeline from: {pipeline_path}") |
| 123 | self.pipeline = FluxPipeline.from_pretrained(pipeline_path, torch_dtype=torch.bfloat16) |
| 124 | |
| 125 | # Load custom transformer if provided |
| 126 | if custom_transformer_path: |
| 127 | print(f"Loading custom transformer from: {custom_transformer_path}") |
| 128 | self.pipeline.transformer = FluxTransformer2DModel.from_pretrained( |
| 129 | custom_transformer_path, |
| 130 | torch_dtype=torch.bfloat16 |
| 131 | ) |
| 132 | |
| 133 | if self.device.type != 'cpu': |
| 134 | print("Enabling model CPU offload for FLUX pipeline.") |
| 135 | self.pipeline.enable_model_cpu_offload() |
| 136 | else: |
| 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, |