danwaldie's picture
Update app.py
eafc480
raw
history blame
658 Bytes
from huggingface_hub import from_pretrained_keras
from keras_cv import models
import gradio as gr
sd_dreambooth_model = models.StableDiffusion(
img_width=512, img_height=512
)
db_diffusion_model = from_pretrained_keras("danwaldie/dreambooth_teddy")
sd_dreambooth_model._diffusion_model = db_diffusion_model
# generate images
def infer(prompt):
generated_images = sd_dreambooth_model.text_to_image(
prompt
)
return generated_images
output = gr.Gallery(label="Outputs").style(grid=(2,2))
# pass function, input type for prompt, the output for multiple images
gr.Interface(infer, inputs=["text"], outputs=[output]).launch()