Using Pruna with AWS AMI

Steps

  • You have an AWS account.

  • You have subscribed to the Pruna AMI.

  • You have launched an EC2 instance from the Pruna AMI.

  • Make sure to use the ubuntu user (you can switch to the ubuntu user by running sudo su - ubuntu).

Minimal Example

You can now run the following script to test the setup. Bear in mind that since you are using the Pruna AMI, you do not need to provide a token when calling the smash function. You can also run any of our other examples in the tutorials section.

import torch
from diffusers import StableDiffusionPipeline
from pruna import smash, SmashConfig

# Define the model you want to smash
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config['compilers'] = ['step_caching']

# Smash the model without a token
smashed_model = smash(
        model=pipe,
        smash_config=smash_config,
    )

# Run the model on a prompt
prompt = "a photo of an astronaut riding a horse on mars"
image = smashed_model(prompt).images[0]