SD 2.1 on keras-cv toy.py

gistfile1.txt

#! /usr/bin/env python

from datetime import datetime
from keras_cv.models import StableDiffusionV2
from PIL import Image

BATCH_SIZE = 1
IMAGE_SIZE = 384
PROMPT = 'beautiful bird, photo-realistic'

model = StableDiffusionV2(img_height=IMAGE_SIZE, img_width=IMAGE_SIZE, jit_compile=True)

image = model.text_to_image(
    prompt=PROMPT,
    batch_size=BATCH_SIZE,
    num_steps=50,
    unconditional_guidance_scale=7.5,
    seed=0xdeadbeef,
)

tstamp = datetime.now().strftime('%Y%m%d-%H%M')
for i in range(0, BATCH_SIZE):
    Image.fromarray(image[i]).save(f'output-toy-{tstamp}-{i}.png')

print('Done!')