ยป Keras API reference / KerasCV / Models / StableDiffusion image-generation model

StableDiffusion image-generation model

[source]

StableDiffusion class

keras_cv.models.StableDiffusion(img_height=512, img_width=512, jit_compile=False)

Keras implementation of Stable Diffusion.

Stable Diffusion is a powerful image generation model that can be used, among other things, to generate pictures according to a short text description (called a "prompt").

Arguments

  • img_height: Height of the images to generate, in pixel. Note that only multiples of 128 are supported; the value provided will be rounded to the nearest valid value. Default: 512.
  • img_width: Width of the images to generate, in pixel. Note that only multiples of 128 are supported; the value provided will be rounded to the nearest valid value. Default: 512.
  • jit_compile: Whether to compile the underlying models to XLA. This can lead to a significant speedup on some systems. Default: False.

Example

from keras_cv.models import StableDiffusion
from PIL import Image

model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
img = model.text_to_image(
    prompt="A beautiful horse running through a field",
    batch_size=1,  # How many images to generate at once
    num_steps=25,  # Number of iterations (controls image quality)
    seed=123,  # Set this to always get the same image from the same prompt
)
Image.fromarray(img[0]).save("horse.png")
print("saved at horse.png")

References