Class PosterGenerator

inference.py:96–157

class PosterGenerator:
    def __init__(self, 
                 pipeline_path="black-forest-labs/FLUX.1-dev",
                 custom_transformer_path=None,
                 qwen_model_path=None,
                 device="cuda:0"):
        
        self.device = torch.device(device) if isinstance(device, str) else device
        
        # Load Qwen model for prompt rewriting
        if qwen_model_path and os.path.exists(qwen_model_path):
            print(f"Loading Qwen model from: {qwen_model_path}")
            self.qwen_agent = QwenRecapAgent(qwen_model_path, device=self.device)
        else:
            self.qwen_agent = None
        
        # Load Flux pipeline
        print(f"Loading Flux pipeline from: {pipeline_path}")
        self.pipeline = FluxPipeline.from_pretrained(pipeline_path, torch_dtype=torch.bfloat16)
        
        # Load custom transformer if provided
        if custom_transformer_path:
            print(f"Loading custom transformer from: {custom_transformer_path}")
            self.pipeline.transformer = FluxTransformer2DModel.from_pretrained(
                custom_transformer_path, 
                torch_dtype=torch.bfloat16
            )
        
        self.pipeline.to(self.device)
    
    def generate(self, 
                 prompt, 
                 enable_recap=True,
                 width=832, 
                 height=1216, 
                 num_inference_steps=28, 
                 guidance_scale=3.5, 
                 seed=None):
        
        # Prompt rewriting
        if enable_recap and self.qwen_agent:
            final_prompt = self.qwen_agent.recap_prompt(prompt)
        else:
            final_prompt = prompt
        
        # Poster generation
        if seed is None:
            seed = random.randint(1, 2**32 - 1)
        
        generator = torch.Generator(device=self.device).manual_seed(seed)
        
        with torch.inference_mode():
            image = self.pipeline(
                prompt=final_prompt,
                generator=generator,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                width=width,