from huggingface_hub import from_pretrained_keras
import keras_cv
import gradio as gr
from tensorflow import keras

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


def generate_images(prompt: str, negative_prompt: str, num_imgs_to_gen: int, num_steps: int, guidance_scale: float):
    generated_img = 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_img


# pass function, input type for prompt, the output for multiple images
gr.Interface(
    title="Keras Dreambooth - Marvin the Paranoid Android",
    description="""This SD model has been fine-tuned to learn the concept of Marvin the Paranoid Android from The Hitchhiker's Guide to the Galaxy.

    To generate your own Marvin, use the phrase 'paranoid marvin a robot' in your prompt.
    """,
    fn=generate_images,
    inputs=[
        gr.Textbox(label="Positive Prompt", value="a photo of paranoid marvin a robot"),
        gr.Textbox(label="Negative Prompt", value="low quality, deformed"),
        gr.Slider(label='Number of gen image', minimum=1, maximum=4, value=2, step=1),
        gr.Slider(label="Inference Steps", value=50),
        gr.Slider(label='Guidance scale', value=7.5, maximum=15, minimum=0, step=0.5),
    ], 
    outputs=[
        gr.Gallery(show_label=False).style(grid=(1,2)),
    ],
    examples=[["a drawing of a white lowpoly paranoid marvin a robot, high quality, 4k, trending on artstation", "low quality, deformed, dark", 2, 50, 7.5]],
    ).queue().launch(debug=True)