Deploying PixArt-Sigma: A Guide to Image Generation on AWS Trainium

Solution Overview

Learn how to effectively deploy the PixArt-Sigma model on AWS Trainium to achieve high-quality image generation.

Step-by-Step Implementation

Step 1: Prerequisites and Setup

Prepare your development environment on appropriate AWS instances.

Step 2: Downloading and Compiling PixArt-Sigma for AWS Trainium

Follow our guide to ensure the successful compilation of the model for optimal performance.

Step 3: Deploying the Model on AWS Trainium

Execute the inference process to generate stunning images from your prompts.

Conclusion

Explore the initial steps of deploying diffusion transformers on AWS, with a focus on PixArt-Sigma’s capabilities.

Deploying PixArt-Sigma on AWS Trainium: A Step-by-Step Guide

In the rapidly-changing world of generative models, PixArt-Sigma stands out as a powerful diffusion transformer model capable of generating images at stunning 4K resolution. This state-of-the-art model significantly outperforms its predecessors, including PixArt-Alpha, through enhancements in both dataset and architectural design. By leveraging AWS’s purpose-built AI chips, Trainium and Inferentia, users can deploy these large generative models more cost-effectively and efficiently. In this first installment of our series, we’ll demonstrate how to deploy PixArt-Sigma on AWS Trainium instances.

Solution Overview

To successfully deploy the PixArt-Sigma model and generate high-quality images, you’ll need to follow a series of well-defined steps:

1. Pre-requisites and Setup
2. Download and Compile the PixArt-Sigma Model for AWS Trainium
3. Deploy the Model on AWS Trainium to Generate Images

Step 1: Prerequisites and Setup

Before diving into the deployment, you’ll need to set up a dedicated environment. Use one of the following instance types: trn1.32xlarge, trn2.48xlarge, or inf2. Here’s how to get started:

1. Launch an Instance:

   • Choose a suitable instance from AWS with Neuron DLAMI. For detailed guidance, refer to the Get Started with Neuron documentation.

2. Set Up a Jupyter Notebook Server:

   • Follow a user guide to set up your server effectively.

3. Clone the AWS Neuron Samples Repository:

   git clone https://github.com/aws-neuron/aws-neuron-samples.git

4. Navigate to the Notebook:

   cd aws-neuron-samples/torch-neuronx/inference

The provided example script is catered to the trn2 instance but can be easily modified to suit trn1 or inf2 instances.

Step 2: Download and Compile the PixArt-Sigma Model for AWS Trainium

This step involves downloading the model and ensuring it’s optimized for AWS Trainium.

1. Download the Model:
   The cache-hf-model.py script within the repository will guide you in downloading PixArt-Sigma from Hugging Face. Alternatively, use the huggingface-cli.

2. Structure of the Repository:
   The repository contains several important files, including:

   • compile_latency_optimized.sh: Script for latency-optimized compilation.
   • hf_pretrained_pixart_sigma_1k_latency_optimized.ipynb: Notebook for running the latency-optimized model.
   • neuron_pixart_sigma: Directory containing essential scripts and helper functions.

3. Sharding Layers:
   The example uses Neuron-specific wrapper classes for model components. These classes help in compiling models and sharding the attention layer across multiple devices to enhance performance.

   For example:

   class InferenceTextEncoderWrapper(nn.Module):
      # Implementation details

   This optimization allows you to replace standard linear layers with NeuronX Distributed components.

4. Compile Individual Sub-models:
   Each component of PixArt-Sigma (text encoder, transformer, and VAE decoder) must be compiled using the right techniques.

   Example for tracing the decoder:

   compiled_decoder = torch_neuronx.trace(
      decoder,
      sample_inputs,
      compiler_workdir=f"{compiler_workdir}/decoder",
      compiler_args=compiler_flags,
      inline_weights_to_neff=False
   )

Step 3: Deploy the Model on AWS Trainium to Generate Images

Now that the model is compiled, it’s time to run inference with PixArt-Sigma.

1. Create a Diffusers Pipeline:
   Utilize the Hugging Face diffusers library to set up your model-specific pipeline:

   pipe: PixArtSigmaPipeline = PixArtSigmaPipeline.from_pretrained(
      "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
      torch_dtype=torch.bfloat16,
      local_files_only=True,
      cache_dir="pixart_sigma_hf_cache_dir_1024"
   )

2. Load Compiled Models:
   Load your compiled models into the pipeline for seamless image generation:

   vae_decoder_wrapper.model = torch_neuronx.DataParallel(
      torch.jit.load(decoder_model_path), [0, 1, 2, 3], False
   )

3. Compose a Prompt:
   Craft a specific prompt to guide the image generation process effectively:

   prompt = "a photo of an astronaut riding a horse on Mars"
   negative_prompt = "mountains"

4. Generate an Image:
   Finally, generate the image by passing the prompt to the PixArt model pipeline:

   images = pipe(
      prompt=prompt,
      negative_prompt=negative_prompt,
      num_images_per_prompt=1,
      height=1024,
      width=1024,
      num_inference_steps=25
   ).images

5. Save the Images:
   Store the generated images for future reference:

   for idx, img in enumerate(images): 
      img.save(f"image_{idx}.png")

Cleanup

To avoid unnecessary charges, remember to stop your EC2 instance via the AWS Management Console or the AWS CLI.

Conclusion

This post outlined the steps required to deploy PixArt-Sigma, an innovative diffusion transformer, on AWS Trainium instances. Stay tuned for future installments in this series, where we will explore running various diffusion transformers and optimizing them for different tasks using Neuron.

About the Authors

The authors of this post are AWS Solutions Architects with a wealth of experience in machine learning, cloud infrastructure, and AI applications. They collaborate with customers to harness technology effectively for innovative solutions within the competitive landscape of AI.

With PixArt-Sigma, the future of high-resolution image generation looks brighter than ever. Happy deploying!