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

dreambooth_model = models.StableDiffusion(img_width=256, img_height=256)

diffusion_model = from_pretrained_keras("moizsajid/dreambooth-markhor")
dreambooth_model._diffusion_model = diffusion_model

# generate images
def infer(prompt: str, negative_prompt: str, num_imgs_to_gen: int, num_steps: int, guidance_scale: float):
    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(
    infer, [
        gr.Textbox(label="Positive Prompt", value="a markhor in space"),
        gr.Textbox(label="Negative Prompt", value="bad anatomy, blurry"),
        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=10),
    ], [
        gr.Gallery(show_label=False),
    ],
    title="Dreambooth Markhor Demo",
    description = "This model is fine-tuned on images of Markhor from the internet (iStock). To use the demo, please add {markhor} to the input string.",
    examples = [["a picture of markhor upside down", "", 4, 100, 10]],
    ).launch()