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_hogwarts_legacy")
dreambooth_model._diffusion_model = loaded_diffusion_model


# generate images
def generate_images(prompt: str, negative_prompt: str, num_imgs_to_gen: int, inference_steps: int, guidance_scale: float):
    output_images = dreambooth_model.text_to_image(
        prompt,
        negative_prompt=negative_prompt,
        batch_size=num_imgs_to_gen,
        num_steps=inference_steps,
        unconditional_guidance_scale=guidance_scale,
    )
    return output_images
    
# Define the UI
with gr.Blocks() as demo:
    gr.HTML("<h2 style=\"font-size: 2em; font-weight: bold\" align=\"center\">Keras Dreambooth - Hogwarts Legacy Demo</h2>")
    gr.HTML("<p style=\"font-size: 14; font-weight: normal\" align=\"left\">This model has been fine-tuned to learn the concept of Hogwarts Legacy student characters.<br>To use this demo, you should have append your prompt with string <q>hogwarts [legacy] student</q></p>")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Positive Prompt", value="a photo of a female hogwarts [legacy] student posing outside hogwarts castle")
            negative_prompt = gr.Textbox(label="Negative Prompt", value="out of frame, blurry, cropped, noisy")
            samples = gr.Slider(label="Number of Images", minimum=1, maximum=6, value=4, step=1)
            inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=50, step=1)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10, value=7.5, step=0.1)
            run = gr.Button(value="Run")
        with gr.Column():
            gallery = gr.Gallery(label="Outputs").style(grid=(1,2))

    run.click(fn=generate_images, inputs=[prompt, negative_prompt, samples, inference_steps, guidance_scale], outputs=gallery)
    
    gr.Examples([["hyperrealistic photo of a smiling female hogwarts [legacy] student posing outside hogwarts castle, 4K, highly detailed, intricate outfit design, art by miho hirano",
                  "ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off",
                  4, 100, 7.5],
                 ["digital art of a male hogwarts [legacy] student holding a wand, high quality, 8k",
                  "ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off",
                  4, 50, 7.5]],
                [prompt, negative_prompt, samples, inference_steps, guidance_scale], gallery, generate_images, cache_examples=True)
    gr.Markdown('Demo created by [Terrence Goh](https://huggingface.co/tgohblio/)')

demo.queue(concurrency_count=3)
demo.launch()