PixArt-α with less than 8GB VRAM

Import what we need

We are going to import a few libraries:

import torch
from diffusers import PixArtAlphaPipeline
from transformers import T5EncoderModel
import gc

We will make use of the garbage collector to remove the models from memory after they have been used.

Initialize parameters

We cannot generate multiple images at once (in parallel / batch_size), as this would increase memory usage quite a lot. It is also not a very good idea to generate just one image, since loading the models into memory also takes some time. The ideal solution to keep memory usage to a minimum and to maximize generation speed, is to use a generation queue to generate one image after another (serial).

queue = []

# Generate image with this prompt. Keep the rest of the parameters as default
queue.extend([{ 'prompt': 'Oppenheimer sits on the beach on a chair, watching a nuclear explosion with a huge mushroom cloud, 1200mm' }])

# Generate image using specific values for all parameters
queue.extend([{
  'prompt': 'pirate ship trapped in a cosmic malestrom nebula',
  'width': 1024,
  'height': 1024,
  'seed': 1152753,
  'cfg': 5,
  'steps': 30,
}])

# Generate 4 images with this prompt. Do not use seed here or they will all come out the same
queue.extend([{ 'prompt': 'supercar', 'cfg': 4 } for _ in range(3)])

Embeddings and Transformer

The next step is to convert the prompt into tokens, then into an embedding and finally pass it through the transformer that applies the attention mechanisms. In PixArt-α, as in Stable Diffusion, the text encoder already takes care of producing a transformed embedding, so the last two parts are already joined.

Also, to avoid doing everything from scratch, let's abstract the code a bit more using the pipeline provided by Hugging Face and thus we save tokenizing ourselves.

We load the text encoder with 16-bit precision (fp16):

text_encoder = T5EncoderModel.from_pretrained(
  'PixArt-alpha/PixArt-XL-2-1024-MS',
  subfolder='text_encoder',
  torch_dtype=torch.float16,
  device_map='auto',
)

The device_map='auto' parameter loads the models (in this case the text encoder) wherever possible. It first utilizes the graphics card memory, then starts using RAM, and finally, the hard drive (although for the sake of your mental health, I hope it doesn't have to resort to RAM, let alone the disk).

We assign this text encoder to the pipeline and also tell it that we don't want to use any transformer, thus avoiding loading it into memory for now (transformer=None refers to the transformer responsible for cleaning the noise).

pipe = PixArtAlphaPipeline.from_pretrained(
  'PixArt-alpha/PixArt-XL-2-1024-MS',
  torch_dtype=torch.float16,
  text_encoder=text_encoder,
  transformer=None,
  device_map='auto',
)

Ir order to avoid reloading this model, we will process all prompts and save the result for later use:

with torch.no_grad():
  for generation in queue:
    generation['embeddings'] = pipe.encode_prompt(generation['prompt'])

As we no longer need the text encoder, we remove references to the model so that the garbage collector can clean up and we also clear the CUDA cache:

del text_encoder
del pipe
gc.collect()
torch.cuda.empty_cache()

print(f'Maximum memory used: {torch.cuda.max_memory_allocated(device="cuda")}')

Image generation

Diffusion models work on a tensor filled with noise, gradually cleaning it through a series of steps managed by the scheduler. When this loop ends, the tensor is passed through the VAE to obtain the final image.

In the case of the pipeline that we are using, all this process is already abstracted into a single function. So, let's instantiate another pipeline again, but this time we won't load the text encoder. In other words, this pipeline will contain the remaining components: the transformer model that is in charge of cleaning the noise (akin to the U-Net in Stable Diffusion), the scheduler and the VAE, which responsible for converting the noise-cleaned tensor into an image.

pipe = PixArtAlphaPipeline.from_pretrained(
  'PixArt-alpha/PixArt-XL-2-1024-MS',
  torch_dtype=torch.float16,
  text_encoder=None,
).to('cuda')

In this case it is necessary to use to('cuda') instead of device_map='auto'.

Now we use a loop to process all the embedding tensors, one after another. Remember that this function already takes care of generating a tensor filled with noise and cleaning it over several steps.

for i, generation in enumerate(queue, start=1):
  generator = torch.Generator(device='cuda')

  if 'seed' in generation:
    generator.manual_seed(generation['seed'])
  else:
    generator.seed()

  image = pipe(
    negative_prompt=None,
    width=generation['width'] if 'width' in generation else 1024,
    height=generation['height'] if 'height' in generation else 1024,
    guidance_scale=generation['cfg'] if 'cfg' in generation else 7,
    num_inference_steps=generation['steps'] if 'steps' in generation else 20,
    generator=generator,
    prompt_embeds=generation['embeddings'][0],
    prompt_attention_mask=generation['embeddings'][1],
    negative_prompt_embeds=generation['embeddings'][2],
    negative_prompt_attention_mask=generation['embeddings'][3],
    num_images_per_prompt=1,
  ).images[0]

  image.save(f'image_{i}.png')

In this loop we have instantiated a generator to which we assign the seed we have defined (generator.manual_seed(generation['seed'])) or a random seed otherwise (generator.seed()).

After that, all we have to do is pass all the necessary arguments to the pipeline:

• negative_prompt: PixArt-α does not accept a negative prompt (as far as I know).
• width and height: The size of the image we have specified or its default value (1024).
• guidance_scale: The value of CFG that we have specified or its default value (7).
• num_inference_steps: The steps value we have specified or its default value (20).
• generator: The generator containing the seed.
• prompt_embeds, prompt_attention_mask, negative_prompt_embeds and negative_prompt_attention_mask: These values are the ones returned by the text encoder and we had saved to use them here (they are inside the list in this order).
• num_images_per_prompt: The number of images that are generated at once (batch size). If we optimize for memory usage it does not make sense to change this value.

The pipeline returns a dictionary (images) containing the images already decoded. Since num_images_per_prompt is always going to be 1, we can directly access the single image by images[0]. We save the images to disk.

The results are these:

Using the VAE manually

As I mentioned above, we cannot use the pipeline with the vae=None parameter as it results in an error. If it were possible, we could run the VAE separately as well (or if we don't want to use the abstraction provided by the pipeline).

In this hypothetical scenario, if we use the output_type='latent' parameter in the pipeline, it will return a tensor in the latent space within the images property (don't use [0] here).

for generation in queue:
  # generator = ...

  generation['latents'] = pipe(
    # Rest of parameters
    output_type='latent',
  ).images

We cannot clean the tensor and pass it through the VAE in the same loop, since both models would be loaded into memory and we would exceed 8GB of memory. It must be done in two loops.

After finishing the first loop we can remove the transformer since we won't need it anymore (although it seems to be optional).

del pipe.transformer
gc.collect()
torch.cuda.empty_cache()

And now, we start the second one that decodes the tensor from the latent space to generate a tensor in the image space:

with torch.no_grad():
  for i, generation in enumerate(queue, start=1):
    image = pipe.vae.decode(
      generation['latents'] / pipe.vae.config.scaling_factor,
      return_dict=False,
    )[0]

    image = pipe.image_processor.postprocess(image, output_type='pil')[0]
    image.save(f'image_{i}.png')

Remember that as we saw in the previous article, the VAE scale factor must be taken into account.

The postprocess method is responsible for converting the tensor into a Pillow image (output_type='pil'). This is the same process we did in the other article using the ToPILImage method from the torchvision library.

Finally we save the image in the same way.