Install and Run Stable Diffusion 2 on Ubuntu

Introduction:

Stable Diffusion 2 is an advanced AI tool that has garnered attention for its impressive capabilities in generating high-quality content. To harness its potential, it is crucial to have a properly set up environment on your Ubuntu system. This guide will walk you through the necessary prerequisites and the installation process for Stable Diffusion 2, ensuring a seamless experience.

Prerequisites:

Before starting the installation process, ensure that you have the following components installed and set up on your system:

  • Python 3.10+ with Conda: The latest version of Python is essential for compatibility with the required libraries and packages. Conda serves as a package manager to streamline the installation of necessary dependencies.
  • Transformers and XFormers: These libraries provide state-of-the-art models and tools for natural language processing and machine learning tasks.
  • PyTorch (Torch): A popular open-source machine learning library that accelerates the development of AI applications.
  • Diffusers: A package specifically designed for Stable Diffusion 2.
  • Hardware Requirements: Ensure that you have a server with at least 20GB of storage and a 10GB+ GPU to handle the processing demands of Stable Diffusion 2.

Install Prerequisites

conda create -n stable_diffusion python=3.10
conda activate stable_diffusion
conda install transformers xformers torch diffusers -c pytorch -c huggingface
#If you want to use pip
pip install transformers xformers torch diffusers

Run Stable Diffusion 2

import torch
from diffusers import DiffusionPipeline
from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
# Workaround for not accepting attention shape using VAE for Flash Attention
pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)

Generate Stable Diffusion Images

pipe("An image of a squirrel in Picasso style").images[0]

Frequent Warnings or Errors

Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment.
To avoid above warning do following...

pip install accelerate