Function main

inference.py:160–193

()

def main():
    parser = argparse.ArgumentParser(description='Generate poster')
    parser.add_argument('--prompt', type=str, required=True, help='Input poster description prompt')
    parser.add_argument('--enable_recap', action='store_true', default=True, help='Enable prompt rewriting (default: True)')
    parser.add_argument('--num_inference_steps', type=int, default=28, help='Number of inference steps')
    parser.add_argument('--guidance_scale', type=float, default=3.5, help='Guidance scale')
    parser.add_argument('--seed', type=int, default=42, help='Seed for random number generator')
    parser.add_argument('--pipeline_path', type=str, default="black-forest-labs/FLUX.1-dev", help='Flux pipeline path')
    parser.add_argument('--custom_transformer_path', type=str, default="PosterCraft/PosterCraft-v1_RL", help='Custom transformer path')
    parser.add_argument('--qwen_model_path', type=str, default="Qwen/Qwen3-8B", help='Qwen model path')
    args = parser.parse_args()

    generator = PosterGenerator(
        pipeline_path=args.pipeline_path,
        custom_transformer_path=args.custom_transformer_path,
        qwen_model_path=args.qwen_model_path,
        device="cuda:0"  
    )
    
    image, final_prompt, seed = generator.generate(
        prompt=args.prompt,
        enable_recap=args.enable_recap,
        width=832,
        height=1216,
        num_inference_steps=args.num_inference_steps,
        guidance_scale=args.guidance_scale,
        seed=args.seed  
    )
    
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = f"generated_poster_{timestamp}_{seed}.png"
    image.save(output_path)
    print(f"Image saved to: {output_path}")
    print(f"Final prompt used: {final_prompt}")


if __name__ == "__main__":