Note: This is a fork as a playground.

A Keras / Tensorflow implementation of Stable Diffusion.

The weights were ported from the original implementation.

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

Install using pip with the git repo:

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

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

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

pip install -r requirements.txt

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

Follow the steps for installing tensorflow with miniconda: install tensorflow with pip or use an existing conda environment.

Then with your conda environment activated, install the dependencies using the requirements.txt file or the requirements_m1.txt file:

pip install -r requirements.txt

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")

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(s) will be stored on the output directory in the following format output/Your-Prompt/[# sample]-[# batch]-s[seed]-[timestamp].png.

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

If you have a lower-end GPU, using --jit --mp might help.

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

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