(self,
pipeline_path="black-forest-labs/FLUX.1-dev",
custom_transformer_path=None,
qwen_model_path=None,
device="cuda:0")
| 95 | |
| 96 | class PosterGenerator: |
| 97 | def __init__(self, |
| 98 | pipeline_path="black-forest-labs/FLUX.1-dev", |
| 99 | custom_transformer_path=None, |
| 100 | qwen_model_path=None, |
| 101 | device="cuda:0"): |
| 102 | |
| 103 | self.device = torch.device(device) if isinstance(device, str) else device |
| 104 | |
| 105 | # Load Qwen model for prompt rewriting |
| 106 | if qwen_model_path and os.path.exists(qwen_model_path): |
| 107 | print(f"Loading Qwen model from: {qwen_model_path}") |
| 108 | self.qwen_agent = QwenRecapAgent(qwen_model_path, device=self.device) |
| 109 | else: |
| 110 | self.qwen_agent = None |
| 111 | |
| 112 | # Load Flux pipeline |
| 113 | print(f"Loading Flux pipeline from: {pipeline_path}") |
| 114 | self.pipeline = FluxPipeline.from_pretrained(pipeline_path, torch_dtype=torch.bfloat16) |
| 115 | |
| 116 | # Load custom transformer if provided |
| 117 | if custom_transformer_path: |
| 118 | print(f"Loading custom transformer from: {custom_transformer_path}") |
| 119 | self.pipeline.transformer = FluxTransformer2DModel.from_pretrained( |
| 120 | custom_transformer_path, |
| 121 | torch_dtype=torch.bfloat16 |
| 122 | ) |
| 123 | |
| 124 | self.pipeline.to(self.device) |
| 125 | |
| 126 | def generate(self, |
| 127 | prompt, |
nothing calls this directly
no test coverage detected