import os
import gradio as gr
import json
from gradio_client import Client

with open('loras.json', 'r') as f:
    loras = json.load(f)
    

def infer (selected_index, prompt, style_prompt, inf_steps, guidance_scale, width, height, seed, lora_weight, progress=gr.Progress(track_tqdm=True)):    
    if selected_index is None:
        raise gr.Error("You must select a LoRA before proceeding.")
    
    # custom_model="lichorosario/dott_remastered_style_lora_sdxl"
    # weight_name="dott_style.safetensors"
    selected_lora = loras[selected_index]
    custom_model = selected_lora["repo"]
    trigger_word = selected_lora["trigger_word"]

    client = Client("fffiloni/sd-xl-custom-model")
    result = client.predict(
    		custom_model=custom_model,
    		api_name="/load_model"
    )
    weight_name = result[2]['value']

    client = Client("fffiloni/sd-xl-custom-model")
    prompt = trigger_word+". "+prompt+". "+style_prompt
    result = client.predict(
		custom_model=custom_model,
		weight_name=weight_name,
		prompt=prompt,
		inf_steps=inf_steps,
		guidance_scale=guidance_scale,
		width=width,
		height=height,
		seed=seed,
		lora_weight=lora_weight,
		api_name="/infer"
    )

    new_result = result + (prompt, )

    return new_result


def update_selection(evt: gr.SelectData):
    selected_lora = loras[evt.index]
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index
    )


css="""
#col-container{
    margin: 0 auto;
    max-width: 720px;
    text-align: left;
}
div#warning-duplicate {
    background-color: #ebf5ff;
    padding: 0 16px 16px;
    margin: 20px 0;
}
div#warning-duplicate > .gr-prose > h2, div#warning-duplicate > .gr-prose > p {
    color: #0f4592!important;
}
div#warning-duplicate strong {
    color: #0f4592;
}
p.actions {
    display: flex;
    align-items: center;
    margin: 20px 0;
}
div#warning-duplicate .actions a {
    display: inline-block;
    margin-right: 10px;
}
button#load_model_btn{
    height: 46px;
}
#status_info{
    font-size: 0.9em;
}
.custom-color {
    color: #030303 !important;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# lichorosario LoRA Portfolio")
    gr.Markdown(
        "### This is my portfolio.\n"
        "**Note**: Generation quality may vary. For best results, adjust the parameters.\n"
        "Special thanks to [@artificialguybr](https://huggingface.co/artificialguybr) and [@fffiloni](https://huggingface.co/fffiloni)."
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt_in = gr.Textbox(
            label="Your Prompt",
            info = "Dont' forget to include your trigger word if necessary"
            )   
            style_prompt_in = gr.Textbox(
            label="Your Style Prompt"
            )   
            selected_info = gr.Markdown("")
            used_prompt = gr.Textbox(
                label="Used prompt"
            )

        with gr.Column(scale=1):
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="LoRA Gallery",
                allow_preview=False,
                columns=2
            )

    with gr.Column(elem_id="col-container"):
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row():
                inf_steps = gr.Slider(
                    label="Inference steps",
                    minimum=12,
                    maximum=50,
                    step=1,
                    value=25
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=50.0,
                    step=0.1,
                    value=7.5
                )
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=2048,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=2048,
                    step=32,
                    value=1024,
                )
    
            with gr.Row():
                seed = gr.Slider(
                    label="Seed",
                    info = "-1 denotes a random seed",
                    minimum=-1,
                    maximum=423538377342,
                    step=1,
                    value=-1
                )
                last_used_seed = gr.Number(
                    label = "Last used seed",
                    info = "the seed used in the last generation",
                )
            lora_weight = gr.Slider(
                label="LoRa weigth",
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=1.0
            )
    submit_btn = gr.Button("Submit")
    image_out = gr.Image(label="Image output")
    selected_index = gr.State(None)

    submit_btn.click(
        fn = infer,
        inputs = [selected_index, prompt_in, style_prompt_in, inf_steps, guidance_scale, width, height, seed, lora_weight],
        outputs = [image_out, last_used_seed, used_prompt]
    )
    gallery.select(update_selection, outputs=[prompt_in, selected_info, selected_index])




demo.launch()