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 Img2Img
• GPU Colab Inpainting
• GPU Colab - Tile / Texture generation
• GPU Colab - Loading Pytorch ckpt Weights
• 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 [email protected]: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 StableDiffusion
from PIL import Image

generator = StableDiffusion(
    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,
)

# for image to image :
img = generator.generate(
    "A Halloween bedroom",
    num_steps=50,
    unconditional_guidance_scale=7.5,
    temperature=1,
    batch_size=1,
    input_image="/path/to/img.png"
)


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.

Using img2img.py from the git repo

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

python img2img.py --prompt="a high quality sketch of people standing with sun and grass , watercolor , pencil color" --input="img.jpeg"

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

Check out the img2img.py file for more options, including the number of steps.

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

Inpainting

Image2Image

1. a high quality sketch of people standing with sun and grass , watercolor , pencil color

References

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