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

conda create --name pydf python=3.8
conda activate pydf 
git clone [email protected]:wangcaihua/stable_diffusion.git
cd stable_diffusion
pip install -r requirements.txt
python app.py

Note: if you have problem in GPU, set jit_compile=True in app.py

Usage

Using web interface

Input what you like in prompt, such as "An astronaut riding a horse"

Input as image and a mask, we will change background for your, just like the following:

As you see, the masked part kept.

The result as following,

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