Stable Diffusion in TensorFlow / Keras

A Keras / Tensorflow implementation of Stable Diffusion.

The weights were ported from the original implementation.

Colab Notebooks

The easiest way to try it out is to use one of the Colab notebooks:

• GPU Colab
• GPU Colab + Mixed Precision
  • ~10s generation time per image (512x512) on default Colab GPU without drop in quality (source)
• TPU Colab.
  • Slower than GPU for single-image generation, faster for large batch of 8+ images (source).
• GPU Colab with Gradio

Installation

Install as a python package

Install using pip with the git repo:

pip install git+https://github.com/divamgupta/stable-diffusion-tensorflow

Installing using the repo

Download the repo, either by downloading the zip file or by cloning the repo with git:

git clone git@github.com:divamgupta/stable-diffusion-tensorflow.git

Using pip without a virtual environment

Install dependencies using the requirements.txt file or the requirements_m1.txt file,:

pip install -r requirements.txt

Using a virtual environment with virtualenv

1. Create your virtual environment for python3:

   python3 -m venv venv
   

2. Activate your virtualenv:

   source venv/bin/activate
   

3. Install dependencies using the requirements.txt file or the requirements_m1.txt file,:

   pip install -r requirements.txt
   

Usage

Using the Python interface

If you installed the package, you can use it as follows:

from stable_diffusion_tf.stable_diffusion import Text2Image
from PIL import Image

generator = Text2Image(
    img_height=512,
    img_width=512,
    jit_compile=False,
)
img = generator.generate(
    "An astronaut riding a horse",
    num_steps=50,
    unconditional_guidance_scale=7.5,
    temperature=1,
    batch_size=1,
)
Image.fromarray(img[0]).save("output.png")

Using text2image.py from the git repo

Assuming you have installed the required packages, you can generate images from a text prompt using:

python text2image.py --prompt="An astronaut riding a horse"

The generated image will be named output.png on the root of the repo. If you want to use a different name, use the --output flag.

python text2image.py --prompt="An astronaut riding a horse" --output="my_image.png"

Check out the text2image.py file for more options, including image size, number of steps, etc.

Example outputs

The following outputs have been generated using this implementation:

1. A epic and beautiful rococo werewolf drinking coffee, in a burning coffee shop. ultra-detailed. anime, pixiv, uhd 8k cryengine, octane render

2. Spider-Gwen Gwen-Stacy Skyscraper Pink White Pink-White Spiderman Photo-realistic 4K

3. A vision of paradise, Unreal Engine

References

1. https://github.com/CompVis/stable-diffusion
2. https://github.com/geohot/tinygrad/blob/master/examples/stable_diffusion.py