Cost-effective AI image generation with PixArt-Σ inference on AWS Trainium and AWS Inferentia

This article discusses deploying PixArt-Sigma, a high-resolution AI image generation model, on AWS Trainium and Inferentia hardware for cost-effective inference.

• PixArt-Sigma is a diffusion transformer capable of generating 4K resolution images
• The deployment process involves compiling three key model components:
  • Text encoder (4B parameters)
  • Denoising transformer (700M parameters)
  • VAE decoder
• Uses advanced techniques like tensor parallelism and sharded attention layers
• Demonstrates how to compile and run the model on AWS Neuron-powered instances
• Provides a step-by-step guide for model setup, compilation, and image generation

The solution aims to enable cost-effective, high-quality AI image generation using AWS's purpose-built AI acceleration hardware.