from huggingface_hub import from_pretrained_keras
from keras_cv import models
import gradio as gr

from tensorflow import keras

keras.mixed_precision.set_global_policy("mixed_float16")

# prepare model

sd_dreambooth_model = models.StableDiffusion(
    img_width=512, img_height=512
)
db_diffusion_model = from_pretrained_keras("Jayabalambika/my_repo")
sd_dreambooth_model._diffusion_model = db_diffusion_model

# generate images
def infer(prompt):
    generated_images = sd_dreambooth_model.text_to_image(
        prompt, batch_size=2
    )
    return generated_images 
    
output = gr.Gallery(label="Outputs").style(grid=(1,2))

# customize interface
title = "Dreambooth Demo on mobile phone Images"
description = "This is a dreambooth model fine-tuned on mobile phone images. To try it, input the concept with {mobile phones}."
examples=[["a photo of mobile phone in outer space"]]
gr.Interface(infer, inputs=["text"], outputs=[output], title=title, description=description, examples=examples).queue().launch()