from stable_diffusion_tf.stable_diffusion import StableDiffusion as StableDiffusionPy
import keras_cv
import gradio as gr
from tensorflow import keras

keras.mixed_precision.set_global_policy("mixed_float16")
# load keras model
resolution=512
sd_dreambooth_model_1=StableDiffusionPy(resolution, resolution, download_weights=False, jit_compile=True)

diffusion_model_pytorch_weights = keras.utils.get_file(
    origin="https://huggingface.co/riffusion/riffusion-model-v1/resolve/main/riffusion-model-v1.ckpt",
    file_hash="99a6eb51c18e16a6121180f3daa69344e571618b195533f67ae94be4eb135a57",
)

sd_dreambooth_model_1.load_weights_from_pytorch_ckpt(diffusion_model_pytorch_weights)

sd_dreambooth_model_1.diffusion_model.load_weights("/dreambooth_riffusion_model_currulao_v1")


def generate_images(prompt: str, num_steps: int, unconditional_guidance_scale: int, temperature: int):
    generated_img = sd_dreambooth_model_1.generate(
        prompt, 
        num_steps=num_steps,
        unconditional_guidance_scale=unconditional_guidance_scale,
        temperature=temperature,
        batch_size=1,
    )

    return generated_img


# pass function, input type for prompt, the output for multiple images
gr.Interface(
    title="Keras Dreambooth Riffusion-Currulao",
    description="""This SD model has been fine-tuned from Riffusion to generate Currulao spectrograms.
    To generate the concept, use the phrase 'a $currulao song' in your prompt.
    """,
    fn=generate_images,
    inputs=[
        gr.Textbox(label="Prompt", value="a $currulao song, lo-fi"),
        gr.Slider(label="Inference steps", value=50),
        gr.Slider(label="Guidance scale", value=7.5, maximum=15, minimum=0, step=0.5),
        gr.Slider(label='Temperature', value=1, maximum=1.5, minimum=0, step=0.1),
    ], 
    outputs=[
        gr.Gallery(show_label=False).style(grid=(1,2)),
    ],
    examples=[["a $currulao song", "low quality, deformed, dark", 2, 50, 7.5]],
    ).queue().launch(debug=True)