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

tf.keras.mixed_precision.set_global_policy("mixed_float16")

# load keras model
resolution = 512
dreambooth_model = models.StableDiffusion(
        img_width=resolution, img_height=resolution, jit_compile=True, 
    )
loaded_diffusion_model = from_pretrained_keras("keras-dreambooth/dreambooth_diffusion_clay_cups")
dreambooth_model._diffusion_model = loaded_diffusion_model


# generate images
def inference(prompt, negative_prompt, num_imgs_to_gen, num_steps, guidance_scale):
    generated_images = dreambooth_model.text_to_image(
        prompt,
        negative_prompt=negative_prompt,
        batch_size=num_imgs_to_gen,
        num_steps=num_steps,
        unconditional_guidance_scale=guidance_scale,
    )
    return generated_images 
    
    

# pass function, input type for prompt, the output for multiple images
gr.Interface(
    inference, [
        gr.Textbox(label="Positive Prompt", value="a bng clay textured unique tea cup designs"),
        gr.Textbox(label="Negative Prompt", value="(bad anatomy), (blurry), grain, low quality"),
        gr.Slider(label='Number of gen image', minimum=1, maximum=4, value=2, step=1),
        gr.Slider(label="Inference Steps",value=100),
        gr.Number(label='Guidance scale', value=7.5),
    ], [
        gr.Gallery(show_label=False).style(grid=(1,2)),
    ],
    title="A universe in Bengali clay",
    description = "This model has been fine-tuned to learn the intricacies of traditional Bengali clay pottery. I trained it using self-made images of clay cups and saucers. My goal was to make my imagination wild with this particular clay texture and image how the world would look in it. To use this demo, you should have {bng} {clay} in the input",
    examples = [["A set of bng clay textured unique tea cup designs, futuristic,  warm colors, fine details, golden ratio composition, trending on artstation, volumetric lighting, octane render", "(bad anatomy), (blurry), grain, low quality", 3, 100, 15],
               ["A world of bng clay textured building, modern architecture, futuristic, fine details, golden ratio composition, trending on artstation, volumetric lighting, octane render", "(bad anatomy), (blurry), grain, low quality", 3, 100, 15],
               ["A collection of bng clay engraved necklace, locket, inspired by Victorian era style, fine details, golden ratio composition, trending on artstation, volumetric lighting, octane render", "(bad anatomy), (blurry), grain, low quality", 3, 100, 15] 
               ],
    cache_examples=True
    ).queue().launch(debug=True)