import gradio as gr
from huggingface_hub import login, HfFileSystem, HfApi

from diffusers import DiffusionPipeline, StableDiffusionXLPipeline
import torch
import copy
import os
import spaces
import random

is_shared_ui = True if "fffiloni/sd-xl-lora-fusion" in os.environ['SPACE_ID'] else False
hf_token = os.environ.get("HF_TOKEN")
login(token = hf_token)

fs = HfFileSystem(token=hf_token)
api = HfApi()

original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)

def get_files(file_paths):
    last_files = {}  # Dictionary to store the last file for each path

    for file_path in file_paths:
        # Split the file path into directory and file components
        directory, file_name = file_path.rsplit('/', 1)
    
        # Update the last file for the current path
        last_files[directory] = file_name
    
    # Extract the last files from the dictionary
    result = list(last_files.values())

    return result

def load_sfts(repo_1_id, repo_2_id):
      
    # List all ".safetensors" files in repos
    sfts_available_files_1 = fs.glob(f"{repo_1_id}/*.safetensors")
    sfts_available_files_1 = get_files(sfts_available_files_1)
    sfts_available_files_2 = fs.glob(f"{repo_2_id}/*.safetensors")
    sfts_available_files_2 = get_files(sfts_available_files_2)
    
    return gr.update(choices=sfts_available_files_1, value=sfts_available_files_1[0], visible=True), gr.update(choices=sfts_available_files_2, value=sfts_available_files_2[0], visible=True)
    
@spaces.GPU
def infer(lora_1_id, lora_1_sfts, lora_2_id, lora_2_sfts, prompt, negative_prompt, lora_1_scale, lora_2_scale, seed):

    unet = copy.deepcopy(original_pipe.unet)
    text_encoder = copy.deepcopy(original_pipe.text_encoder)
    text_encoder_2 = copy.deepcopy(original_pipe.text_encoder_2)

    pipe = StableDiffusionXLPipeline(
        vae = original_pipe.vae,
        text_encoder = text_encoder,
        text_encoder_2 = text_encoder_2,
        scheduler = original_pipe.scheduler,
        tokenizer = original_pipe.tokenizer,
        tokenizer_2 = original_pipe.tokenizer_2,
        unet = unet
    )

    pipe.to("cuda")

    pipe.load_lora_weights(
        lora_1_id,
        weight_name = lora_1_sfts,        
        low_cpu_mem_usage = True,
        use_auth_token = True
    )

    pipe.fuse_lora(lora_1_scale)

    pipe.load_lora_weights(
        lora_2_id,
        weight_name = lora_2_sfts,        
        low_cpu_mem_usage = True,
        use_auth_token = True
    )

    pipe.fuse_lora(lora_2_scale)

    if negative_prompt == "" :
        negative_prompt = None
    
    if seed < 0 :
        seed = random.randint(0, 423538377342)
    
    generator = torch.Generator(device="cuda").manual_seed(seed)

    image = pipe(
        prompt = prompt,
        negative_prompt = negative_prompt,
        num_inference_steps = 25,
        width = 1024,
        height = 1024,
        generator = generator
    ).images[0]

    pipe.unfuse_lora()
    
    return image, seed

css="""
#col-container{
    margin: 0 auto;
    max-width: 750px;
    text-align: left;
}
div#warning-duplicate {
    background-color: #ebf5ff;
    padding: 0 10px 5px;
    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;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):

        if is_shared_ui:
            top_description = gr.HTML(f'''
                <div class="gr-prose">
                    <h2><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg>
                    Note: you might want to use private custom LoRa models</h2>
                    <p class="main-message">
                        To do so, <strong>duplicate the Space</strong> and run it on your own profile using <strong>your own access token</strong> and eventually a GPU (T4-small or A10G-small) for faster inference without waiting in the queue.<br />
                    </p>
                    <p class="actions">
                        <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true">
                            <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" />
                        </a>
                        to start using private models and skip the queue
                    </p>
                </div>
            ''', elem_id="warning-duplicate")

        title = gr.HTML(
        '''
        <h1 style="text-align: center;">SD-XL LoRA Fusion</h1>
        <p style="text-align: center;">
        Fuse 2 custom StableDiffusion-XL LoRa models <br />
        If you are running this demo in a duplicated private space, all your private LoRa models tagged will be automatically listed in LoRa IDs dropdowns
        </p>
        '''
        )
        
        # PART 1 • MODELS
        if not is_shared_ui:
            your_username = api.whoami()["name"]
            my_models = api.list_models(author=your_username, filter=["diffusers", "stable-diffusion-xl", 'lora'])
            model_names = [item.modelId for item in my_models]
            
            #print(model_names)           
            
            with gr.Row():
                
                with gr.Column():
            
                    lora_1_id = gr.Dropdown(
                        label = "LoRa 1 ID",
                        choices = model_names,
                        allow_custom_value = True
                        #placeholder = "username/model_id"
                    )
                    
                    lora_1_sfts = gr.Dropdown(
                        label = "Safetensors file",
                        visible=False
                    )
                
                with gr.Column():
                    
                    lora_2_id = gr.Dropdown(
                        label = "LoRa 2 ID",
                        choices = model_names,
                        allow_custom_value = True
                        #placeholder = "username/model_id"
                    )
    
                    lora_2_sfts = gr.Dropdown(
                        label = "Safetensors file",
                        visible=False
                    )
        else:
            with gr.Row():
                
                with gr.Column():
                    lora_1_id = gr.Textbox(
                        label = "LoRa 1 ID",
                        placeholder = "username/model_id"
                    )

                    lora_1_sfts = gr.Dropdown(
                        label = "Safetensors file",
                        visible=False
                    )

                with gr.Column():
                    lora_2_id = gr.Textbox(
                            label = "LoRa 2 ID",
                            placeholder = "username/model_id"
                        )
    
                    lora_2_sfts = gr.Dropdown(
                        label = "Safetensors file",
                        visible=False
                    )
        
        load_models_btn = gr.Button("1. Load models and .safetensors")

        # PART 2 • INFERENCE
        with gr.Column():
            with gr.Row():
            
                prompt = gr.Textbox(
                    label = "Your prompt",
                    show_label = True,
                    info = "Use your trigger words into a coherent prompt",
                    placeholder = "e.g: a triggerWordOne portrait in triggerWord2 style"
                )

            run_btn = gr.Button("2. Run", elem_id="run_button")
        
            output_image = gr.Image(
                label = "Output"
            )

        # Advanced Settings
        with gr.Accordion("Advanced Settings", open=False):
            
            with gr.Row():
                
                lora_1_scale = gr.Slider(
                    label = "LoRa 1 scale",
                    minimum = 0,
                    maximum = 1,
                    step = 0.1,
                    value = 0.7
                )
                
                lora_2_scale = gr.Slider(
                    label = "LoRa 2 scale",
                    minimum = 0,
                    maximum = 1,
                    step = 0.1,
                    value = 0.7
                )
            
            negative_prompt = gr.Textbox(
                label = "Negative prompt"
            )

            seed = gr.Slider(
                label = "Seed",
                info = "-1 denotes a random seed",
                minimum = -1,
                maximum = 423538377342,
                value = -1
            )

            last_used_seed = gr.Number(
                label = "Last used seed",
                info = "the seed used in the last generation",
            )
    
    # ACTIONS
    load_models_btn.click(
        fn = load_sfts, 
        inputs = [
            lora_1_id,
            lora_2_id
        ],
        outputs = [
            lora_1_sfts,
            lora_2_sfts
        ]
    )
    run_btn.click(
        fn = infer,
        inputs = [
            lora_1_id,
            lora_1_sfts,
            lora_2_id,
            lora_2_sfts,
            prompt,
            negative_prompt,
            lora_1_scale,
            lora_2_scale,
            seed
        ],
        outputs = [
            output_image, 
            last_used_seed
        ]
    )

demo.queue(concurrency_count=2).launch()