import gradio as gr
from textual_inversio_with_blueloss import TextualInversion

display_choices = ["matrix", "dragon born", "birb style", "pool rooms", "minecraft concept art"]
repo_id_embeds=["sd-concepts-library/matrix::in <hatman-matrix> world",
                "sd-concepts-library/dragonborn::with <dragonborn> concept", 
                "sd-concepts-library/birb-style::in <birb-style> concept", 
                "sd-concepts-library/poolrooms::with <poolrooms>", 
                "sd-concepts-library/minecraft-concept-art::with <minecraft-concept-art> concept"
                ]

textualInversion = TextualInversion(pretrained_model_name_or_path = "CompVis/stable-diffusion-v1-4", repo_id_embeds=repo_id_embeds)

def generate_image(prompt, selected_concept, grayscale_image):
    return textualInversion.generate_image(prompt, display_choices.index(selected_concept), grayscale_image=grayscale_image)

demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter your prompt"),
        gr.Dropdown(choices=display_choices, label="Select concept", value=display_choices[0]),
        gr.Checkbox(label="Grayscale Image", value=False)
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Textual Inversion Image Generator",
    description="Generate images using textual inversion concepts",
    examples=[["a flying dog", display_choices[0], False]],
    allow_flagging=False
)

# Launch the app
demo.launch()