| 106 | |
| 107 | # --- Poster Generator Class --- |
| 108 | class PosterGenerator: |
| 109 | def __init__(self, pipeline_path, qwen_model_path, custom_weights_path, device): |
| 110 | self.device = device |
| 111 | self.qwen_agent = self._load_qwen_agent(qwen_model_path) |
| 112 | self.pipeline = self._load_flux_pipeline(pipeline_path, custom_weights_path) |
| 113 | |
| 114 | def _load_qwen_agent(self, qwen_model_path): |
| 115 | if not qwen_model_path: |
| 116 | return None |
| 117 | return QwenRecapAgent(model_path=qwen_model_path, device_map=str(self.device)) |
| 118 | |
| 119 | def _load_flux_pipeline(self, pipeline_path, custom_weights_path): |
| 120 | pipeline = FluxPipeline.from_pretrained(pipeline_path, torch_dtype=torch.bfloat16) |
| 121 | |
| 122 | if custom_weights_path and os.path.exists(custom_weights_path): |
| 123 | logging.info(f"Loading custom Transformer from directory: {custom_weights_path}") |
| 124 | transformer = FluxTransformer2DModel.from_pretrained( |
| 125 | custom_weights_path, torch_dtype=torch.bfloat16 |
| 126 | ) |
| 127 | pipeline.transformer = transformer |
| 128 | |
| 129 | pipeline.to(self.device) |
| 130 | return pipeline |
| 131 | |
| 132 | def generate(self, prompt, enable_recap, **kwargs): |
| 133 | final_prompt = prompt |
| 134 | if enable_recap: |
| 135 | if not self.qwen_agent: |
| 136 | raise gr.Error("Recap is enabled, but the recap model is not available. Check model path.") |
| 137 | final_prompt = self.qwen_agent.recap_prompt(prompt) |
| 138 | |
| 139 | generator = torch.Generator(device=self.device).manual_seed(kwargs['seed']) |
| 140 | |
| 141 | with torch.inference_mode(): |
| 142 | image = self.pipeline( |
| 143 | prompt=final_prompt, |
| 144 | generator=generator, |
| 145 | num_inference_steps=kwargs['num_inference_steps'], |
| 146 | guidance_scale=kwargs['guidance_scale'], |
| 147 | width=kwargs['width'], |
| 148 | height=kwargs['height'] |
| 149 | ).images[0] |
| 150 | |
| 151 | return image, final_prompt |
| 152 | |
| 153 | # --- Global Model Initialization --- |
| 154 | generator = PosterGenerator( |