import gradio as gr
import torch
import os
from diffusers import StableDiffusionXLPipeline, AutoencoderKL, LCMScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from share_btn import community_icon_html, loading_icon_html, share_js
from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
import lora
import copy
import json
import gc
import random
from urllib.parse import quote

lora_list = hf_hub_download(repo_id="multimodalart/LoraTheExplorer", filename="sdxl_loras.json", repo_type="space")
with open(lora_list, "r") as file:
    data = json.load(file)
    sdxl_loras_raw = [
        {
            "image": item["image"],
            "title": item["title"],
            "repo": item["repo"],
            "trigger_word": item["trigger_word"],
            "weights": item["weights"],
            "is_compatible": item["is_compatible"],
            "is_pivotal": item.get("is_pivotal", False),
            "text_embedding_weights": item.get("text_embedding_weights", None),
            "likes": item.get("likes", 0),
            "downloads": item.get("downloads", 0),
            "is_nc": item.get("is_nc", False),
            "new": item.get("new", False),
        }
        for item in data
    ]

device = "cuda" 

state_dicts = {}

for item in sdxl_loras_raw:
    saved_name = hf_hub_download(item["repo"], item["weights"])
    
    if not saved_name.endswith('.safetensors'):
        state_dict = torch.load(saved_name)
    else:
        state_dict = load_file(saved_name)
    
    state_dicts[item["repo"]] = {
        "saved_name": saved_name,
        "state_dict": state_dict
    }

sdxl_loras_raw_new = [item for item in sdxl_loras_raw if item.get("new") == True]

sdxl_loras_raw = [item for item in sdxl_loras_raw if item.get("new") != True]
    
lcm_lora_id = "lcm-sd/lcm-sdxl-base-1.0-lora"

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
)
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to(device)
pipe.load_lora_weights(lcm_lora_id, weight_name="lcm_sdxl_lora.safetensors", adapter_name="lcm_lora", use_auth_token=os.getenv('HF_TOKEN'))
last_lora = ""
last_merged = False
last_fused = False
js = '''
var button = document.getElementById('button');
// Add a click event listener to the button
button.addEventListener('click', function() {
    element.classList.add('selected');
});
'''

def update_selection(selected_state: gr.SelectData, sdxl_loras, is_new=False):
    lora_repo = sdxl_loras[selected_state.index]["repo"]
    instance_prompt = sdxl_loras[selected_state.index]["trigger_word"]
    new_placeholder = "Type a prompt. This LoRA applies for all prompts, no need for a trigger word" if instance_prompt == "" else "Type a prompt to use your selected LoRA"
    weight_name = sdxl_loras[selected_state.index]["weights"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
    is_compatible = sdxl_loras[selected_state.index]["is_compatible"]
    is_pivotal = sdxl_loras[selected_state.index]["is_pivotal"]
    
    use_with_diffusers = f'''
    ## Using [`{lora_repo}`](https://huggingface.co/{lora_repo})
                        
    ## Use it with diffusers:
    '''
    if is_compatible:
        use_with_diffusers += f'''
        from diffusers import StableDiffusionXLPipeline
        import torch
    
        model_path = "stabilityai/stable-diffusion-xl-base-1.0"
        pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
        pipe.to("cuda")
        pipe.load_lora_weights("{lora_repo}", weight_name="{weight_name}")
    
        prompt = "{instance_prompt}..."
        lora_scale= 0.9
        image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5, cross_attention_kwargs={{"scale": lora_scale}}).images[0]
        image.save("image.png")
        '''
    elif not is_pivotal:
        use_with_diffusers += "This LoRA is not compatible with diffusers natively yet. But you can still use it on diffusers with `bmaltais/kohya_ss` LoRA class, check out this [Google Colab](https://colab.research.google.com/drive/14aEJsKdEQ9_kyfsiV6JDok799kxPul0j )"
    else:
        use_with_diffusers += f"This LoRA is not compatible with diffusers natively yet. But you can still use it on diffusers with sdxl-cog `TokenEmbeddingsHandler` class, check out the [model repo](https://huggingface.co/{lora_repo}#inference-with-🧨-diffusers)"
    use_with_uis = f'''
    ## Use it with Comfy UI, Invoke AI, SD.Next, AUTO1111: 

    ### Download the `*.safetensors` weights of [here](https://huggingface.co/{lora_repo}/resolve/main/{weight_name})
    
    - [ComfyUI guide](https://comfyanonymous.github.io/ComfyUI_examples/lora/)
    - [Invoke AI guide](https://invoke-ai.github.io/InvokeAI/features/CONCEPTS/?h=lora#using-loras)
    - [SD.Next guide](https://github.com/vladmandic/automatic)
    - [AUTOMATIC1111 guide](https://stable-diffusion-art.com/lora/)
    '''
    if(is_new):
        if(selected_state.index == 0):
            selected_state.index = -9999
        else:
            selected_state.index *= -1
    
    return (
        updated_text,
        instance_prompt,
        gr.update(placeholder=new_placeholder),
        selected_state,
        use_with_diffusers,
        use_with_uis,
        gr.Gallery(selected_index=None)
    )


def check_selected(selected_state):
    if not selected_state:
        raise gr.Error("You must select a LoRA")

def run_lora(prompt, negative, lora_scale, selected_state, sdxl_loras, sdxl_loras_new, progress=gr.Progress(track_tqdm=True)):
    global last_lora, last_merged, last_fused, pipe
    print("Index when running ", selected_state.index)
    if(selected_state.index < 0):
        if(selected_state.index == -9999):
            selected_state.index = 0
        else:
             selected_state.index *= -1
        sdxl_loras = sdxl_loras_new
    print("Selected State: ", selected_state.index)
    print(sdxl_loras[selected_state.index]["repo"])
    if negative == "":
        negative = None

    if not selected_state:
        raise gr.Error("You must select a LoRA")
    repo_name = sdxl_loras[selected_state.index]["repo"]
    weight_name = sdxl_loras[selected_state.index]["weights"]
    
    full_path_lora = state_dicts[repo_name]["saved_name"]
    loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
    cross_attention_kwargs = None
    if last_lora != repo_name:
        #if(last_fused):
            #pipe.unfuse_lora()
        pipe.load_lora_weights(loaded_state_dict, adapter_name=state_dicts[repo_name]["saved_name"])
        pipe.set_adapters([state_dicts[repo_name]["saved_name"], "lcm_lora"], adapter_weights=[0.8, 1.0])
        #last_fused = True
        is_pivotal = sdxl_loras[selected_state.index]["is_pivotal"]
        if(is_pivotal):
            #Add the textual inversion embeddings from pivotal tuning models
            text_embedding_name = sdxl_loras[selected_state.index]["text_embedding_weights"]
            text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
            tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
            embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
            embhandler = TokenEmbeddingsHandler(text_encoders, tokenizers)
            embhandler.load_embeddings(embedding_path)
            
    image = pipe(
        prompt=prompt,
        negative_prompt=negative,
        num_inference_steps=4,
        guidance_scale=0.5,
    ).images[0]
    
    last_lora = repo_name
    gc.collect()
    return image, gr.update(visible=True)

def shuffle_gallery(sdxl_loras):
    random.shuffle(sdxl_loras)
    return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras

def swap_gallery(order, sdxl_loras):
    if(order == "random"):
        return shuffle_gallery(sdxl_loras)
    else:
        sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get(order, 0), reverse=True)
        return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery

def deselect():
  return gr.Gallery(selected_index=None)

with gr.Blocks(css="custom.css") as demo:
    gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
    gr_sdxl_loras_new = gr.State(value=sdxl_loras_raw_new)
    title = gr.HTML(
        """<h1><img src="https://i.imgur.com/vT48NAO.png" alt="LoRA">LCM LoRA the Explorer</h1>
        Combine loading an [LCM LoRA](#) with your favorite SDXL LoRa and run LoRAs in only 4 steps. Check out our blog to see how this works.
        """,
        elem_id="title",
    )
    selected_state = gr.State()
    with gr.Row(elem_id="main_app"):
        with gr.Group(elem_id="gallery_box"):
            order_gallery = gr.Radio(choices=["random", "likes"], value="random", label="Order by", elem_id="order_radio")
            selected_loras = gr.Gallery(label="Selected LoRAs", height=80, show_share_button=False, visible=False, elem_id="gallery_selected", )
            new_gallery = gr.Gallery(label="New LoRAs", elem_id="gallery_new", columns=3, value=[(item["image"], item["title"]) for item in sdxl_loras_raw_new], allow_preview=False, show_share_button=False)
            gallery = gr.Gallery(
                #value=[(item["image"], item["title"]) for item in sdxl_loras],
                label="SDXL LoRA Gallery",
                allow_preview=False,
                columns=3,
                elem_id="gallery",
                show_share_button=False,
                height=784
            )
        with gr.Column():
            prompt_title = gr.Markdown(
                value="### Click on a LoRA in the gallery to select it",
                visible=True,
                elem_id="selected_lora",
            )
            with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, placeholder="Type a prompt after selecting a LoRA", elem_id="prompt")
                button = gr.Button("Run", elem_id="run_button")
            with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
                community_icon = gr.HTML(community_icon_html)
                loading_icon = gr.HTML(loading_icon_html)
                share_button = gr.Button("Share to community", elem_id="share-btn")
            result = gr.Image(
                interactive=False, label="Generated Image", elem_id="result-image"
            )
            with gr.Accordion("Advanced options", open=False):
                negative = gr.Textbox(label="Negative Prompt")
                weight = gr.Slider(0, 10, value=0.8, step=0.1, label="LoRA weight")
    with gr.Column(elem_id="extra_info"):
        with gr.Accordion(
            "Use it with: 🧨  diffusers, ComfyUI, Invoke AI, SD.Next, AUTO1111",
            open=False,
            elem_id="accordion",
        ):
            with gr.Row():
                use_diffusers = gr.Markdown("""## Select a LoRA first 🤗""")
                use_uis = gr.Markdown()
        with gr.Accordion("Submit a LoRA! 📥", open=False):
            submit_title = gr.Markdown(
                "### Streamlined submission coming soon! Until then [suggest your LoRA in the community tab](https://huggingface.co/spaces/multimodalart/LoraTheExplorer/discussions) 🤗"
            )
            with gr.Group(elem_id="soon"):
                submit_source = gr.Radio(
                    ["Hugging Face", "CivitAI"],
                    label="LoRA source",
                    value="Hugging Face",
                )
                with gr.Row():
                    submit_source_hf = gr.Textbox(
                        label="Hugging Face Model Repo",
                        info="In the format `username/model_id`",
                    )
                    submit_safetensors_hf = gr.Textbox(
                        label="Safetensors filename",
                        info="The filename `*.safetensors` in the model repo",
                    )
                with gr.Row():
                    submit_trigger_word_hf = gr.Textbox(label="Trigger word")
                    submit_image = gr.Image(
                        label="Example image (optional if the repo already contains images)"
                    )
                submit_button = gr.Button("Submit!")
                submit_disclaimer = gr.Markdown(
                    "This is a curated gallery by me, [apolinário (multimodal.art)](https://twitter.com/multimodalart). I'll try to include as many cool LoRAs as they are submitted! You can [duplicate this Space](https://huggingface.co/spaces/multimodalart/LoraTheExplorer?duplicate=true) to use it privately, and add your own LoRAs by editing `sdxl_loras.json` in the Files tab of your private space."
                )
    order_gallery.change(
        fn=swap_gallery,
        inputs=[order_gallery, gr_sdxl_loras],
        outputs=[gallery, gr_sdxl_loras],
        queue=False
    )
    gallery.select(
        fn=update_selection,
        inputs=[gr_sdxl_loras],
        outputs=[prompt_title, prompt, prompt, selected_state, use_diffusers, use_uis, new_gallery],
        queue=False,
        show_progress=False
    )
    new_gallery.select(
        fn=update_selection,
        inputs=[gr_sdxl_loras_new, gr.State(True)],
        outputs=[prompt_title, prompt, prompt, selected_state, use_diffusers, use_uis, gallery],
        queue=False,
        show_progress=False
    )
    prompt.submit(
        fn=check_selected,
        inputs=[selected_state],
        queue=False,
        show_progress=False
    ).success(
        fn=run_lora,
        inputs=[prompt, negative, weight, selected_state, gr_sdxl_loras, gr_sdxl_loras_new],
        outputs=[result, share_group],
    )
    button.click(
        fn=check_selected,
        inputs=[selected_state],
        queue=False,
        show_progress=False
    ).success(
        fn=run_lora,
        inputs=[prompt, negative, weight, selected_state, gr_sdxl_loras, gr_sdxl_loras_new],
        outputs=[result, share_group],
    )
    share_button.click(None, [], [], js=share_js)
    demo.load(fn=shuffle_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras], queue=False, js=js)
demo.queue(max_size=20)
demo.launch(share=True)